83a5d52e7b
The fp32 LoRA parameters and the bnb 4-bit base matmuls need a single compute dtype during the forward, exactly like the diffusers dreambooth scripts run under accelerator.autocast. Without it the 4-bit backward on FLUX.1-dev fails with an illegal-address CUBLAS error partway into the first step. Z-Image and Qwen-Image smokes are unaffected and the SDXL path (its own trainer) is untouched. |
||
|---|---|---|
| .. | ||
| __init__.py | ||
| diffusion_dit_trainer.py | ||
| diffusion_lora_trainer.py | ||
| diffusion_train_common.py | ||
| diffusion_training_service.py | ||
| resume.py | ||
| s3_dataset.py | ||
| trainer.py | ||
| training.py | ||
| worker.py | ||