From 83a5d52e7b0375caa5ead07dbc7de45d96db6d3d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 16:05:56 +0000 Subject: [PATCH] Wrap the DiT training forward in bf16 autocast 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. --- .../core/training/diffusion_dit_trainer.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index efaa4d9408..6a81b91804 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -25,6 +25,7 @@ from __future__ import annotations import gc import random import time +from contextlib import nullcontext from dataclasses import dataclass from pathlib import Path from typing import Any, Callable, Optional @@ -577,11 +578,21 @@ def run_dit_lora_training( else (t.to(device) if t is not None else None) for t in emb ) - model_pred = spec.forward( - transformer, noisy, timesteps, sigmas, emb_dev, cfg, device, weight_dtype + # bf16 autocast around the forward + loss, matching the diffusers dreambooth + # scripts' accelerator.autocast: it reconciles the fp32 LoRA params with the + # bnb 4-bit base matmuls in one compute dtype. Without it the 4-bit backward + # on FLUX dies with an illegal-address / CUBLAS failure. + autocast = ( + torch.autocast(device_type = "cuda", dtype = torch.bfloat16) + if device == "cuda" + else nullcontext() ) - target = noise - latents - loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean") + with autocast: + model_pred = spec.forward( + transformer, noisy, timesteps, sigmas, emb_dev, cfg, device, weight_dtype + ) + target = noise - latents + loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean") (loss / cfg.gradient_accumulation_steps).backward() step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps