tsiendragon/qwen-image-finetune/qwen-image-edit-aipc-finetune-05-adaptation

Step 5 — Framework Adaptation (XPU + NF4 Patches)

Series position: Step 5 of 8.

Scope: this step adapts a clone of tsiendragon/qwen-image-finetune for XPU. The

framework is a custom training loop (not HF Trainer), so it hardcodes torch.cuda.* and

CUDA-only paths. If your clone is already adapted (has device_utils.py and the patches

below), some or all steps may already be done — run the §5.8 import test + checker to confirm

before re-applying anything.

Prerequisites:

• config.yaml with trainer: set (Step 4) — determines whether the Plus section (§5.7) applies
• The qwen-image-finetune project directory (Step 1)

Next step: when adaptation_check.py exits 0, proceed to skill qwen-image-edit-aipc-finetune-06-quantize.

Which trainer files to adapt

§5.1–§5.6 are framework-wide and always apply, including to the base trainer (`qwen_image_edit_trainer.py`) — even when your model uses the Plus trainer, because QwenImageEditPlusTrainer inherits the base trainer's training loop (that's where the NF4-A2 autocast lives) and its device-wrapped helpers. So: always do §5.1–§5.6.

Then, only if `detect_trainer.py` selected `QwenImageEditPlus` (Step 4), also do §5.7 — adapting the Plus trainer's own load_model. (The Plus trainer overrides only load_model and the multi-image methods; everything else it inherits from the base trainer you already adapted.)

5.1 Device-agnostic helpers (device_utils.py)

Replace each torch.cuda.* call in the framework with a device-agnostic wrapper that checks XPU first. This confines the XPU adaptation to the framework's own source files and leaves third-party libraries (Triton, bitsandbytes, accelerate) untouched.

In qwen-image-finetune the wrappers live at src/qflux/utils/device_utils.py. Create this file (before any other change if adapting a fresh clone) with exactly this content:

# src/qflux/utils/device_utils.py
# enable XPU: device-agnostic helpers — call these instead of torch.cuda.* directly.
import contextlib
import torch
def device_empty_cache() -> None:
    if torch.xpu.is_available():           # enable XPU
        torch.xpu.empty_cache()
    elif torch.cuda.is_available():
        torch.cuda.empty_cache()
def device_synchronize() -> None:
    if torch.xpu.is_available():           # enable XPU
        torch.xpu.synchronize()
    elif torch.cuda.is_available():
        torch.cuda.synchronize()
def device_manual_seed_all(seed: int) -> None:
    if torch.xpu.is_available():           # enable XPU
        torch.xpu.manual_seed_all(seed)
    else:
        torch.cuda.manual_seed_all(seed)
def device_context(device):
    """Device context manager replacing torch.cuda.device().
    torch.xpu.device("cpu") raises; this helper returns nullcontext() for
    CPU and the correct accelerator context otherwise.
    """
    device_str = str(device) if device is not None else "cpu"
    if device_str == "cpu":
        return contextlib.nullcontext()
    if torch.xpu.is_available():           # enable XPU
        return torch.xpu.device(device)
    return torch.cuda.device(device)

Replace each direct torch.cuda.* call in the framework:

Import the helpers from the submodule directly (as in the table); there is no need to re-export them from qflux/utils/__init__.py.

5.2 Find and fix residual patterns

The helpers cover functional calls. A separate scan is needed for string-based device comparisons and CUDA-only library imports, which helpers cannot address:

grep -rn "torch\.cuda\.\|flash_attention_2\|transformer_engine\|device\.type.*cuda\|Qwen/Qwen-Image-Edit" src/qflux/

The scan covers both trainer files (qwen_image_edit_trainer.py and qwen_image_edit_plus_trainer.py) — fix every hit in whichever one(s) your model uses.

Example device.type fix in a config validator:

if d.type == "cuda":
    if not torch.cuda.is_available():
        raise ValueError(f"CUDA not available but got device={d}.")
# enable XPU: validate XPU device availability
if d.type == "xpu" and not torch.xpu.is_available():
    raise ValueError(f"XPU not available but got device={d}.")

Why the `max_train_steps` fix needs both edits. train_epoch()'s guard saves a final

checkpoint and returns — but it only returns from train_epoch(). fit()'s

for epoch in range(...) loop keeps going, so every remaining epoch re-enters train_epoch(),

trips the guard again, and writes another full resumable checkpoint before returning. With

num_epochs set well above the epoch count max_train_steps actually reaches (the Step 4

recommender's default pairing), that is one redundant checkpoint directory per remaining epoch,

each carrying a complete copy of the optimizer state — enough to fill a disk on a long run. The

symptom is a pile of checkpoint-last-<epoch>-<step>-last directories that all share the same

step number and differ only by epoch. Nothing prunes them: train.checkpoints_total_limit is

declared in the schema but never wired into ProjectConfiguration, so it has no effect.

The copies are byte-identical, so a run that skipped edit (2) still trained correctly — the cost

is disk and save time, not the model. adaptation_check.py reports the missing break as a WARN

for that reason, not an ERROR.

Breaking the epoch loop as well makes the guard write exactly one final checkpoint. fit()'s own

closing save_checkpoint(...) reuses the same epoch and step, hence the same directory name, so it

overwrites rather than adding another.

5.3 accelerate single-card XPU config

qflux.main launches via accelerate launch --config_file <path>. Create an XPU-specific config accelerate_config_xpu.yaml in the project root:

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: NO
mixed_precision: 'no'
num_machines: 1
num_processes: 1
main_training_function: main
dynamo_backend: 'no'
use_cpu: false
deepspeed_config: {}

Key differences from a typical CUDA multi-GPU config:

YAML gotcha: dynamo_backend: no unquoted parses as YAML boolean False; accelerate then

calls .upper() on it and raises AttributeError: 'bool' object has no attribute 'upper'.

Always quote 'no' (and same for mixed_precision: 'no').

5.4 Pydantic schema patches

Two edits to src/qflux/data/config.py so the NF4 config the recommender emits validates against the upstream schema.

a) Add the `quantize_type` field to `ModelConfig`. Upstream ModelConfig has quantize: bool but no `quantize_type` field, and it is declared with ConfigDict(extra="forbid") — so any config carrying model.quantize_type: nf4 fails immediately with pydantic ... Extra inputs are not permitted. Add the field:

class ModelConfig(BaseModel):
    model_config = ConfigDict(extra="forbid")
    ...
    quantize: bool = False
    quantize_type: str | None = None   # enable XPU: "nf4" selects NF4 QLoRA loading

The NF4 loader (§5.5 / Step 6) reads quantize_type == "nf4" to pass use_nf4=True; without this field the whole NF4 recipe is unreachable.

b) Relax `num_workers` to `>= 0`. Upstream declares num_workers with a > 0 constraint. On Windows under spawn start-method, classes registered via trust_remote_code are not picklable across processes — workers crash on dataset iteration, and num_workers = 0 (the fix) is rejected by the schema. Relax the constraint to >= 0 in the dataset config validator. The recommender defaults to 1, but 0 should be acceptable as an escape hatch. (This is one of two Windows spawn-related fixes — the other is the mp.set_start_method guard in §5.2; they are independent.)

Cache key default (`prompt_empty_drop_keys`). Upstream CacheConfig defaults this to

["prompt_embed", ...] — note the missing s. The cached key is prompt_embeds, so

caption-dropout raises KeyError: 'prompt_embed'. The recommender sidesteps this by **emitting

cache.prompt_empty_drop_keys: ["prompt_embeds", "prompt_embeds_mask"]** in the config, so no

source edit is required as long as you use the generated config. (If you hand-write a config,

set this key yourself.)

5.5 BitsAndBytes XPU backend + NF4 QLoRA required patches

5.5.1 BitsAndBytes XPU backend (informational)

Knowing which bnb operations dispatch to which backend on XPU helps diagnose failures and explains why Step 1 requires Visual Studio + oneAPI. For bitsandbytes >= 0.48.2: