[HaozheZhang6]

What does this PR do?

Follow-up to #13895 (merged), which fixed the fp16 LoRA "unscale" crash in examples/dreambooth/train_dreambooth_lora.py. The same bug is present in the other DreamBooth LoRA training scripts — this PR applies the same fix to all of them.

Root cause (same as #13895)

Under --mixed_precision="fp16", cast_training_params(models, dtype=torch.float32) keeps the trainable LoRA params in fp32. The in-loop validation pipeline is rebuilt around the live training transformer:

pipeline = SomePipeline.from_pretrained(..., transformer=unwrap_model(transformer), torch_dtype=weight_dtype)

log_validation then casts that pipeline to torch_dtype (= fp16), which downcasts the shared transformer's fp32 LoRA params back to fp16. The next optimizer step then raises:

ValueError: Attempting to unscale FP16 gradients.

The final (post-training) validation builds a fresh pipeline from the saved weights, so it is unaffected either way.

Affected scripts & fix

I verified each script has all three conditions (fp16 cast_training_params guard + in-loop validation sharing the live transformer + the fp16 downcast in log_validation). None of these were covered by #13895, and none have an open issue/PR.

Group A — pipeline.to(accelerator.device, dtype=torch_dtype) → drop dtype (keep the device move), exactly like #13895:

• train_dreambooth_lora_flux.py
• train_dreambooth_lora_flux_kontext.py
• train_dreambooth_lora_qwen_image.py
• train_dreambooth_lora_hidream.py
• advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py

Group B — the cast is pipeline.to(dtype=torch_dtype) (no device move) immediately followed by pipeline.enable_model_cpu_offload(), so the cast line is simply removed. Frozen weights already use weight_dtype (from from_pretrained(torch_dtype=weight_dtype)) and the offload call handles device placement:

• train_dreambooth_lora_z_image.py
• train_dreambooth_lora_flux2.py
• train_dreambooth_lora_flux2_img2img.py
• train_dreambooth_lora_flux2_klein.py
• train_dreambooth_lora_flux2_klein_img2img.py

(The torch_dtype parameter of log_validation is left in place for signature consistency, same as the merged base script.)

Testing

Same situation as #13895: this is a GPU-only codepath. On CPU accelerator.scaler is None (the fp16 GradScaler is CUDA-only), so the example CI cannot reproduce it, and the example suite has no --mixed_precision fp16 tests for this reason. The fix is mechanical and identical to the one already merged in #13895; the mechanism is verified there. ruff check and ruff format --check pass on all 10 changed files.

Before submitting

• Did you read the contributor guideline?
• Was this discussed/approved? Follow-up to merged Fix fp16 LoRA unscale crash after validation in train_dreambooth_lora.py #13895.
• Did you write any new necessary tests? Same GPU-only limitation as Fix fp16 LoRA unscale crash after validation in train_dreambooth_lora.py #13895 — can't be reproduced in the CPU example CI.

Who can review?

@sayakpaul (reviewed and merged #13895)

[HaozheZhang6]

Gentle nudge -- green for ~2 weeks. This fixes the fp16 LoRA unscale crash that triggers after validation. Could a maintainer take a look when there's time?