From 5f725bacf22120498e6563eeeba767352eb51b3a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 09:26:05 +0000 Subject: [PATCH] Add base_precision speed modes to DiT training: bf16 2.3-2.6x, int8, fp8 New base_precision config for the DiT trainers: nf4 (unchanged default) | bf16 | int8 | fp8 | auto, advertised per family + per machine through /api/train/diffusion/info (precision_modes, recommended_precision, supports_compile) so the UI can gate the selector. - bf16: dense transformer + regional torch.compile (auto-armed). The measured speed mode: 2.3x nf4 on FLUX (1.81 -> 4.12 steps/s), 2.6x on Z-Image (2.5 -> 6.38 steps/s) on B200, at dense-weight VRAM (FLUX 24.7 GB / Z-Image 13.6 GB peak vs 10.4 / 4.7 for nf4). - int8: torchao weight-only int8 on the frozen base, quantized AFTER add_adapter (quantizing first trips peft 0.18's TorchaoLoraLinear, which is incompatible with the torchao 0.16 config API). Runs eager: inductor rejects the int8 subclass training graph (aliased subclass outputs), so compile is force-disabled for it. - fp8: torchao convert_to_float8_training on the frozen linears (filter skips lora_ modules, proj_out, non-divisible-by-16 dims, pad_inner_dim), applied after add_adapter, compile auto-armed. Works and round-trips, but measured SLOWER than compiled bf16 at LoRA-training shapes (FLUX 3.15 vs 4.12 steps/s; Z-Image similar), so it is an explicit opt-in and auto never picks it. - auto: free VRAM (measured before load) + dense-size table -> bf16 when it fits with headroom, int8 in the middle band, else nf4. Prequant bnb repos always resolve to nf4; dense modes on them are rejected at validation with a pointer to the family's dense base. Two crashes found and fixed along the way: - The cuDNN SDPA backend's training graph fails on the FLUX attention shapes (torch 2.10 + cu130, B200): mha_graph.execute errors, then the context degrades into illegal memory accesses. The perf-flag guard now pins flash/mem-efficient SDPA for the run (mathematically equivalent, snapshot/restored). nf4 escaped it by routing attention differently. - Regional compile now uses dynamic=True (the inference layer's proven default): dynamic=False specialisation fused a gemm_and_bias epilogue that failed with CUBLAS_STATUS_EXECUTION_FAILED on the FLUX training graph; dynamic=True is also faster (Z-Image 3.84 -> 6.38 steps/s). Verified: 98 backend tests green (new test_diffusion_base_precision.py: validation, auto policy table, fp8 filter, compile gating, /info fields); per-mode 40-step runs on FLUX + Z-Image with loss means inside the nf4 envelope and adapter round-trip generation through the normal LoRA path for bf16-, fp8-, and int8-trained adapters. --- .../core/training/diffusion_dit_trainer.py | 224 ++++++++++++++---- .../core/training/diffusion_train_common.py | 82 ++++++- studio/backend/models/training.py | 14 ++ .../tests/test_diffusion_base_precision.py | 209 ++++++++++++++++ .../tests/test_diffusion_train_perf.py | 7 +- 5 files changed, 486 insertions(+), 50 deletions(-) create mode 100644 studio/backend/tests/test_diffusion_base_precision.py diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 4da16d764a..a7852a52b4 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -47,6 +47,7 @@ from core.training.diffusion_train_common import ( _publish_to_lora_catalog, _restore_perf_flags, discover_image_caption_pairs, + repo_is_prequantized, ) # Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks @@ -74,6 +75,9 @@ class _FamilySpec: lora_targets: tuple[str, ...] # bf16 only (Z-Image overflows fp16 and its RoPE/embedder run in fp32). force_bf16: bool + # Approximate dense-bf16 transformer weight size, used by base_precision="auto" to + # decide which mode fits the free VRAM (with headroom for activations + optimizer). + dense_bf16_gb: float # Phased load, so the (multi-GB) transformer never has to coexist with the text # encoders + VAE: ``load_conditioners`` builds the pipeline WITHOUT its transformer # (encode_prompt + VAE only) and returns (pipe, vae); ``load_transformer`` loads the @@ -155,12 +159,9 @@ def _bnb_4bit_config(): ) -def _repo_is_prequantized(base_model: str) -> bool: - """Heuristic: a repo whose name marks a bitsandbytes 4-bit build already ships a - quantized transformer, so we load it as-is rather than re-quantizing on the fly. A - dense (bf16) base instead gets on-the-fly nf4 quantization for QLoRA.""" - name = str(base_model or "").lower() - return "bnb-4bit" in name or "-4bit" in name or "int4" in name or "nf4" in name +# Kept as a module name for existing callers/tests; the heuristic itself moved to +# diffusion_train_common so config validation can use it without importing this module. +_repo_is_prequantized = repo_is_prequantized def _load_quantized_transformer(transformer_cls, cfg): @@ -191,24 +192,132 @@ def _load_pipe_without_transformer(pipe_cls, cfg, device): return pipe, pipe.vae -def _load_dit_transformer(transformer_cls, cfg, device, qlora): - """Load the transformer alone. A prequant (bnb-4bit) repo carries its quantization - config and loads 4-bit as-is; a dense base is quantized to nf4 on the fly when - ``qlora`` (the default) so QLoRA still fits.""" +def _load_dit_transformer(transformer_cls, cfg, device, base_precision): + """Load the transformer alone in the resolved ``base_precision``: + + - nf4: a prequant (bnb-4bit) repo carries its quantization config and loads 4-bit + as-is; a dense base is quantized to nf4 on the fly. The memory floor. + - bf16 / fp8: the dense transformer (fp8 converts its frozen linears to float8 + training compute AFTER the LoRA attaches; storage stays bf16). + - int8: the dense transformer quantized in place to torchao weight-only int8 (the + PEFT-attachable scheme), roughly halving the bf16 weight footprint.""" import torch - if qlora and not _repo_is_prequantized(cfg.base_model): - return _load_quantized_transformer(transformer_cls, cfg) - transformer = transformer_cls.from_pretrained( + if base_precision == "nf4": + if not repo_is_prequantized(cfg.base_model): + return _load_quantized_transformer(transformer_cls, cfg) + transformer = transformer_cls.from_pretrained( + cfg.base_model, + subfolder = "transformer", + torch_dtype = torch.bfloat16, + token = cfg.hf_token, + ) + # A prequant load is already device-placed by bitsandbytes. + if not getattr(transformer, "is_loaded_in_4bit", False): + transformer = transformer.to(device) + return transformer + + # Dense load for bf16 / fp8 / int8. int8 quantizes AFTER the LoRA attaches (see + # _int8_quantize_base): quantizing first makes peft dispatch its TorchaoLoraLinear + # wrapper, whose peft-0.18 constructor is incompatible with the torchao-0.16 config API + # (missing get_apply_tensor_subclass). + return transformer_cls.from_pretrained( cfg.base_model, subfolder = "transformer", torch_dtype = torch.bfloat16, token = cfg.hf_token, + ).to(device) + + +def _int8_quantize_base(transformer) -> None: + """torchao weight-only int8 on the big frozen linears, applied after add_adapter so + the base_layer inside each LoRA wrapper quantizes while the adapters stay high + precision. ``make_filter_fn`` (shared with the inference quant layer) keeps only + Linears with >= 512 features -- which also naturally skips the rank-sized LoRA + matrices -- and drops the M=1 modulation projections int8 kernels reject.""" + from core.inference.diffusion_transformer_quant import exclude_tokens_for_scheme, make_filter_fn + from torchao.quantization import Int8WeightOnlyConfig, quantize_ + + quantize_( + transformer, + Int8WeightOnlyConfig(), + filter_fn = make_filter_fn(512, exclude_name_tokens = exclude_tokens_for_scheme("int8")), ) - # A bnb-quantized load is already device-placed; a dense one still sits on CPU. - if not getattr(transformer, "is_loaded_in_4bit", False): - transformer = transformer.to(device) - return transformer + + +def _fp8_module_filter(mod, fqn: str) -> bool: + """Which frozen linears get float8 training compute: skip anything LoRA-owned (the + adapters must stay high precision -- PEFT has no float8 base support), the output + projection, and shapes float8 kernels reject (dims not divisible by 16), matching the + diffusers FLUX2 reference filter.""" + import torch.nn as nn + + if not isinstance(mod, nn.Linear): + return False + if "lora_" in fqn: + return False + if fqn.endswith("proj_out") or ".proj_out." in fqn: + return False + return mod.in_features % 16 == 0 and mod.out_features % 16 == 0 + + +def _apply_fp8_training(transformer, on_event) -> bool: + """Convert the frozen base linears to torchao float8 training compute (dynamic scaling; + weights stay bf16 in memory). Applied AFTER add_adapter so the filter can exclude the + LoRA modules. Never fatal: on any failure the run continues in bf16 with a warning.""" + try: + from torchao.float8 import Float8LinearConfig, convert_to_float8_training + + convert_to_float8_training( + transformer, + module_filter_fn = _fp8_module_filter, + config = Float8LinearConfig(pad_inner_dim = True), + ) + return True + except Exception as exc: # noqa: BLE001 -- fp8 is an optimisation, never fatal + _emit(on_event, "warning", message = f"fp8 training unavailable, using bf16 compute: {exc}") + return False + + +def _pick_auto_precision(prequant, device, free_gb, dense_gb, capability, has_fp8) -> str: + """Pure policy for base_precision="auto": nf4 for a prequant base or no CUDA; else the + fastest dense mode whose weights + headroom (activations, optimizer, cache) fit the + free VRAM at decision time. bf16 + regional compile is the measured speed winner + (2.3-2.6x over nf4 on B200); fp8 stays an explicit opt-in because torchao float8's + dynamic-scaling overhead made it SLOWER than compiled bf16 at LoRA-training shapes on + the same hardware. ``capability``/``has_fp8`` remain parameters so the policy can be + revisited per GPU generation without changing callers.""" + _ = capability, has_fp8 + if prequant or device != "cuda" or not free_gb or not dense_gb: + return "nf4" + headroom = 1.5 + if free_gb > dense_gb * headroom: + return "bf16" + if free_gb > dense_gb * 0.55 * headroom: + return "int8" + return "nf4" + + +def _resolve_base_precision(cfg, spec, device) -> str: + """Resolve "auto" against the live GPU (free VRAM measured BEFORE anything loads); + explicit modes pass through (normalized() already validated them).""" + mode = (cfg.base_precision or "nf4").strip().lower() + if mode != "auto": + return mode + prequant = repo_is_prequantized(cfg.base_model) + free_gb = None + capability = None + has_fp8 = False + if device == "cuda": + try: + import torch + + free_gb = torch.cuda.mem_get_info()[0] / 1e9 + capability = torch.cuda.get_device_capability() + has_fp8 = hasattr(torch, "float8_e4m3fn") + except Exception: # noqa: BLE001 -- probe failure -> the safe mode + pass + return _pick_auto_precision(prequant, device, free_gb, spec.dense_bf16_gb, capability, has_fp8) # ── FLUX.1-dev ──────────────────────────────────────────────────────────────── @@ -217,9 +326,9 @@ def _flux_load_conditioners(cfg, device, weight_dtype): return _load_pipe_without_transformer(FluxPipeline, cfg, device) -def _flux_load_transformer(cfg, device, weight_dtype, qlora): +def _flux_load_transformer(cfg, device, weight_dtype, base_precision): from diffusers import FluxTransformer2DModel - return _load_dit_transformer(FluxTransformer2DModel, cfg, device, qlora) + return _load_dit_transformer(FluxTransformer2DModel, cfg, device, base_precision) def _flux_encode_prompts(pipe, captions, device): @@ -325,12 +434,12 @@ def _qwen_load_conditioners(cfg, device, weight_dtype): return _load_pipe_without_transformer(QwenImagePipeline, cfg, device) -def _qwen_load_transformer(cfg, device, weight_dtype, qlora): +def _qwen_load_transformer(cfg, device, weight_dtype, base_precision): # The prequant default (unsloth/Qwen-Image-2512-unsloth-bnb-4bit) ships the transformer - # 4-bit and loads trainable as-is; the dense 20B Qwen/Qwen-Image is quantized to nf4 on - # the fly so QLoRA still fits. + # 4-bit and loads trainable as-is under nf4; the dense modes need the 20B + # Qwen/Qwen-Image base. from diffusers import QwenImageTransformer2DModel - return _load_dit_transformer(QwenImageTransformer2DModel, cfg, device, qlora) + return _load_dit_transformer(QwenImageTransformer2DModel, cfg, device, base_precision) def _qwen_encode_prompts(pipe, captions, device): @@ -443,11 +552,11 @@ def _zimage_load_conditioners(cfg, device, weight_dtype): return _load_pipe_without_transformer(ZImagePipeline, cfg, device) -def _zimage_load_transformer(cfg, device, weight_dtype, qlora): - # Prequant default loads 4-bit as-is; the dense bf16 Tongyi-MAI base is quantized to nf4 - # on the fly. Z-Image is bf16 only (its RoPE/embedder run fp32; fp16 overflows). +def _zimage_load_transformer(cfg, device, weight_dtype, base_precision): + # Prequant default loads 4-bit as-is under nf4; the dense modes use the bf16 Tongyi-MAI + # base. Z-Image is bf16 only (its RoPE/embedder run fp32; fp16 overflows). from diffusers import ZImageTransformer2DModel - return _load_dit_transformer(ZImageTransformer2DModel, cfg, device, qlora) + return _load_dit_transformer(ZImageTransformer2DModel, cfg, device, base_precision) def _zimage_encode_prompts(pipe, captions, device): @@ -513,6 +622,7 @@ _SPECS: dict[str, _FamilySpec] = { family = "flux.1", lora_targets = _FLUX_TARGETS, force_bf16 = False, + dense_bf16_gb = 23.8, load_conditioners = _flux_load_conditioners, load_transformer = _flux_load_transformer, encode_prompts = _flux_encode_prompts, @@ -526,6 +636,7 @@ _SPECS: dict[str, _FamilySpec] = { family = "qwen-image", lora_targets = _QWEN_TARGETS, force_bf16 = True, + dense_bf16_gb = 41.0, load_conditioners = _qwen_load_conditioners, load_transformer = _qwen_load_transformer, encode_prompts = _qwen_encode_prompts, @@ -539,6 +650,7 @@ _SPECS: dict[str, _FamilySpec] = { family = "z-image", lora_targets = _ZIMAGE_TARGETS, force_bf16 = True, + dense_bf16_gb = 12.3, load_conditioners = _zimage_load_conditioners, load_transformer = _zimage_load_transformer, encode_prompts = _zimage_encode_prompts, @@ -684,23 +796,36 @@ def _sample_cached_latents(cache, idxs, variant_rng, device): return lat_a + lat_b * torch.randn_like(lat_a) -def _should_compile(cfg, base_is_bnb, device) -> bool: +def _should_compile(cfg, base_is_bnb, device, base_precision = "nf4") -> bool: mode = (cfg.compile_transformer or "auto").strip().lower() if device != "cuda" or mode == "off": return False + # torch.compile cannot trace the torchao int8 subclass in training (inductor rejects + # the aliased subclass graph outputs), so int8 always runs eager. + if base_precision == "int8": + return False if mode == "on": return True - # auto: regional compile is a clean win on a dense base but fragile over bitsandbytes - # 4-bit modules (graph breaks in the dequant path), so it stays off for QLoRA. - return not base_is_bnb + # auto: regional compile is the whole point of the dense modes (measured 2.6x on + # Z-Image bf16) but fragile over bitsandbytes 4-bit modules (graph breaks in the + # dequant path), so it stays off for QLoRA. fp8 is only competitive compiled. + return base_precision in ("bf16", "fp8") -def _maybe_compile_transformer(transformer, cfg, base_is_bnb, device, on_event) -> bool: +def _maybe_compile_transformer( + transformer, cfg, base_is_bnb, device, on_event, base_precision = "nf4" +) -> bool: """Regionally compile the transformer blocks (diffusers compile_repeated_blocks) after the LoRA is attached. Never fatal: a wrap failure falls back to eager with a warning event, and dynamo's suppress_errors keeps a frame that fails to COMPILE at the first step running eager instead of raising mid-run.""" - if not _should_compile(cfg, base_is_bnb, device): + if not _should_compile(cfg, base_is_bnb, device, base_precision): + if base_precision == "fp8": + _emit( + on_event, + "warning", + message = "fp8 training without torch.compile is slow; enable compile for the speedup.", + ) return False import torch @@ -719,10 +844,12 @@ def _maybe_compile_transformer(transformer, cfg, base_is_bnb, device, on_event) setattr(dynamo_cfg, attr, max(getattr(dynamo_cfg, attr) or 0, 64)) if hasattr(dynamo_cfg, "suppress_errors"): dynamo_cfg.suppress_errors = True - # Shapes are static under the latent cache (fixed resolution/batch/text bucket), so - # dynamic=False lets inductor specialise. fullgraph only on a dense base: bnb 4-bit - # layers graph-break by design. - fn(fullgraph = not base_is_bnb, dynamic = False) + # dynamic=True matches the inference speed layer's proven default: on torch 2.10 / + # B200 the dynamic=False specialisation fused a gemm_and_bias epilogue that failed + # with CUBLAS_STATUS_EXECUTION_FAILED then an illegal memory access on the FLUX + # training graph. fullgraph only on a dense base: bnb 4-bit layers graph-break by + # design. + fn(fullgraph = not base_is_bnb, dynamic = True) return True except Exception as exc: # noqa: BLE001 -- optimisation only, never fatal _emit(on_event, "warning", message = f"torch.compile disabled (eager fallback): {exc}") @@ -858,10 +985,12 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto # the same seed-deterministic sequence whether or not the cache is enabled. variant_rng = random.Random(cfg.seed + 1) - # Phase 3: only now load the transformer (QLoRA nf4 by default: the prequant Qwen / - # Z-Image repos are already 4-bit; FLUX quantizes on the fly). - transformer = spec.load_transformer(cfg, device, weight_dtype, qlora = True) - base_is_bnb = True # the DiT base is always a bitsandbytes 4-bit QLoRA load today + # Phase 3: only now load the transformer, in the resolved base precision (nf4 QLoRA by + # default; bf16 / int8 / fp8 are the dense speed modes; "auto" picks from free VRAM + # measured before the load). + base_precision = _resolve_base_precision(cfg, spec, device) + transformer = spec.load_transformer(cfg, device, weight_dtype, base_precision) + base_is_bnb = base_precision == "nf4" # Freeze the base; attach the trainable LoRA to the transformer. transformer.requires_grad_(False) @@ -887,7 +1016,16 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto cast_training_params(transformer, dtype = torch.float32) lora_params = [p for p in transformer.parameters() if p.requires_grad] - compiled = _maybe_compile_transformer(transformer, cfg, base_is_bnb, device, on_event) + # int8 / fp8 convert the frozen base linears AFTER the LoRA attaches, so the adapter + # modules are excluded and stay high precision. + if base_precision == "int8": + _int8_quantize_base(transformer) + if base_precision == "fp8" and not _apply_fp8_training(transformer, on_event): + base_precision = "bf16" + + compiled = _maybe_compile_transformer( + transformer, cfg, base_is_bnb, device, on_event, base_precision + ) # Compiled Qwen graphs need one fixed text length across steps: pin the pad bucket to # the dataset's longest caption (encode_prompt already caps it at 1024 tokens). qwen_pad_to = None @@ -908,7 +1046,7 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto num_training_steps = cfg.train_steps, ) - _emit(on_event, "model_load_completed", compiled = compiled) + _emit(on_event, "model_load_completed", compiled = compiled, base_precision = base_precision) transformer.train() n_images = len(image_paths) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index f1989f6305..c2bcab4bcf 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -124,6 +124,35 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None return "sdxl" +def repo_is_prequantized(base_model: str) -> bool: + """Heuristic: a repo whose name marks a bitsandbytes 4-bit build already ships a + quantized transformer, so it loads as-is for nf4 and cannot serve the dense + (bf16/int8/fp8) base precisions.""" + name = str(base_model or "").lower() + return "bnb-4bit" in name or "-4bit" in name or "int4" in name or "nf4" in name + + +def train_precision_modes() -> tuple[list[str], str]: + """(supported base_precision modes, recommended pick) for the current machine: nf4 + always works; bf16/int8/auto need CUDA; fp8 needs an fp8-capable GPU (sm89+). Used by + the /info endpoint so the UI can gate the precision selector. Never raises.""" + modes = ["nf4"] + recommended = "nf4" + try: + import torch + + if torch.cuda.is_available(): + modes += ["bf16", "int8"] + major, minor = torch.cuda.get_device_capability() + if (major, minor) >= (8, 9) and hasattr(torch, "float8_e4m3fn"): + modes.append("fp8") + modes.append("auto") + recommended = "auto" + except Exception: # noqa: BLE001 -- no torch / probe failure -> nf4 only + pass + return modes, recommended + + def get_trainer(family: str) -> Callable[..., str]: """Return the training entrypoint for ``family``. Imports the trainer module lazily so this shared module stays free of the heavy trainer imports (and any import cycle).""" @@ -178,12 +207,16 @@ def family_train_infos() -> list[dict[str, Any]]: the family registry so it stays in sync with what the trainers actually support.""" from core.inference.diffusion_families import detect_family + dit_modes, dit_recommended = train_precision_modes() infos: list[dict[str, Any]] = [] for name in trainable_family_names(): fam = detect_family("", override = name) if fam is None: continue repos = list(fam.train_base_repos) or [fam.base_repo] + # base_precision / compile apply to the DiT trainer only; SDXL keeps its + # mixed_precision lever, so the UI hides the selector for it. + is_dit = name in ("flux.1", "qwen-image", "z-image") infos.append( { "name": name, @@ -192,6 +225,9 @@ def family_train_infos() -> list[dict[str, Any]]: "base_repos": repos, "defaults": train_defaults(name), "vram_note": _FAMILY_VRAM_NOTES.get(name, ""), + "precision_modes": dit_modes if is_dit else [], + "recommended_precision": dit_recommended if is_dit else "nf4", + "supports_compile": is_dit, } ) return infos @@ -240,6 +276,12 @@ class DiffusionLoraConfig: # TF32 matmuls + high fp32 matmul precision + cudnn autotuning for the run. Near-lossless; # disable for strict bit-reproducibility A/Bs. enable_tf32: bool = True + # DiT base transformer precision: "nf4" (bitsandbytes QLoRA, the memory floor and the + # default), "bf16" (dense, fastest eager, compile-friendly), "int8" (torchao + # weight-only, half of bf16), "fp8" (torchao float8 training compute on the frozen + # linears, Ada/Hopper/Blackwell + compile), or "auto" (pick by free VRAM + GPU class). + # Non-nf4 modes need a dense base repo (not a prequant bnb-4bit one). SDXL ignores it. + base_precision: str = "nf4" # How often to emit a progress event (in optimizer steps). log_every: int = 1 # Optional explicit family override ("sdxl" / "flux.1" / ...); None = detect from @@ -276,6 +318,21 @@ class DiffusionLoraConfig: compile_transformer = str(self.compile_transformer or "auto").strip().lower() if compile_transformer not in ("off", "on", "auto"): raise ValueError("compile_transformer must be one of off / on / auto") + base_precision = str(self.base_precision or "nf4").strip().lower() + if base_precision not in ("nf4", "bf16", "int8", "fp8", "auto"): + raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / auto") + if base_precision in ("bf16", "int8", "fp8"): + if repo_is_prequantized(self.base_model): + raise ValueError( + f"base_precision={base_precision!r} needs a dense base repo, but " + f"'{self.base_model}' is already bitsandbytes-quantized. Pick the " + f"family's dense (bf16) base repo for this mode, or use nf4/auto." + ) + if self.mixed_precision != "bf16": + raise ValueError( + f"base_precision={base_precision!r} trains in bf16 compute; set " + f"mixed_precision to bf16." + ) # learning_rate can arrive as a string ("1e-4") from the Studio config path, which # preserves it as a string after validation; coerce so AdamW receives a float. try: @@ -298,6 +355,7 @@ class DiffusionLoraConfig: hf_token = token or None, cache_variants = int(self.cache_variants), compile_transformer = compile_transformer, + base_precision = base_precision, resolved_family = resolved_family, ) @@ -428,6 +486,18 @@ def _apply_perf_flags( torch.set_float32_matmul_precision("high") if cudnn_benchmark: torch.backends.cudnn.benchmark = True + # The cuDNN SDPA backend's TRAINING graph is broken for the FLUX attention shapes + # on torch 2.10 + cu130 (B200): mha_graph.execute fails, then poisons the context + # into illegal memory accesses. Flash / mem-efficient SDPA are mathematically + # equivalent, so pin those for the run (restored on exit). + cuda_backends = getattr(torch.backends, "cuda", None) + if cuda_backends is not None and hasattr(cuda_backends, "enable_cudnn_sdp"): + try: + snap["cudnn_sdp"] = bool(cuda_backends.cudnn_sdp_enabled()) + except Exception: # noqa: BLE001 -- flag unreadable: skip the tweak entirely + snap["cudnn_sdp"] = None + if snap["cudnn_sdp"]: + cuda_backends.enable_cudnn_sdp(False) except Exception: # noqa: BLE001 -- perf flags are never fatal pass return snap @@ -441,13 +511,15 @@ def _restore_perf_flags(snap: Optional[dict]) -> None: from core.inference.diffusion_speed import restore_backend_flags restore_backend_flags(snap.get("flags")) - if snap.get("matmul_precision"): - try: - import torch + try: + import torch + if snap.get("matmul_precision"): torch.set_float32_matmul_precision(snap["matmul_precision"]) - except Exception: # noqa: BLE001 -- best-effort restore - pass + if snap.get("cudnn_sdp"): + torch.backends.cuda.enable_cudnn_sdp(True) + except Exception: # noqa: BLE001 -- best-effort restore + pass def _assert_trusted_base_model(base_model: str) -> None: diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 8dafb7dfa4..58a9779609 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -727,6 +727,14 @@ class DiffusionTrainingStartRequest(BaseModel): enable_tf32: bool = Field( True, description = "TF32 matmuls + cudnn autotuning (near-lossless speedup)" ) + base_precision: Literal["nf4", "bf16", "int8", "fp8", "auto"] = Field( + "nf4", + description = ( + "DiT base transformer precision: nf4 QLoRA (memory floor, default), bf16 dense, " + "int8 torchao weight-only, fp8 float8 training compute (Ada/Hopper/Blackwell), " + "or auto (pick by free VRAM + GPU class). Dense modes need a non-prequant base." + ), + ) class DiffusionTrainingStopRequest(BaseModel): @@ -801,6 +809,12 @@ class DiffusionTrainableFamily(BaseModel): base_repos: List[str] = Field(default_factory = list) defaults: dict = Field(default_factory = dict) vram_note: str = "" + # base_precision modes this machine supports for the family (empty = the family has no + # precision selector, e.g. SDXL), plus the recommended pick and whether regional + # torch.compile applies. Defaults keep older backends' payloads valid. + precision_modes: List[str] = Field(default_factory = list) + recommended_precision: str = "nf4" + supports_compile: bool = False class DiffusionTrainingInfoResponse(BaseModel): diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py new file mode 100644 index 0000000000..6c59a84fee --- /dev/null +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""CPU-only unit tests for the DiT base_precision work. + +Covers the new precision plumbing the precision PR adds: the ``base_precision`` +config validation (dense-vs-prequant + mixed-precision gating), the prequant-repo +heuristic and its trainer alias, the pure ``auto`` precision policy table, the +explicit-mode passthrough of ``_resolve_base_precision``, the fp8 module filter, the +fp8 branch of the compile policy, the ``train_precision_modes`` machine probe, the +family-info precision fields, and the request-model ``base_precision`` field. No GPU / +model load: every helper here is pure or name-based, so the config validation runs on +name matching (``resolve_trainable_family`` is offline) and the torch probe is monkeypatched. +""" + +from __future__ import annotations + +import pytest +import torch.nn as nn + +import core.training.diffusion_train_common as common +from core.training import diffusion_dit_trainer as dit +from core.training.diffusion_train_common import ( + DiffusionLoraConfig, + _config_from_dict, + repo_is_prequantized, + train_precision_modes, +) +from models.training import DiffusionTrainingStartRequest + +# A dense (non-prequant) DiT base and a prequant bnb-4bit base. Both resolve a trainer +# family from their names alone, so normalized() runs without a network call. +_FLUX_DENSE = "black-forest-labs/FLUX.1-dev" +_Z_PREQUANT = "unsloth/Z-Image-Turbo-unsloth-bnb-4bit" + + +def _cfg(base_model = _FLUX_DENSE, **kw) -> DiffusionLoraConfig: + return DiffusionLoraConfig(base_model = base_model, data_dir = "d", output_dir = "o", **kw) + + +# ── base_precision validation ───────────────────────────────────────────────── +def test_base_precision_validation(): + # Default normalizes to the nf4 memory floor. + assert _cfg().normalized().base_precision == "nf4" + + # An unknown mode is rejected by name. + with pytest.raises(ValueError, match = "base_precision"): + _cfg(base_precision = "banana").normalized() + + # A dense mode is case/space-insensitive and stored lowered: " FP8 " on a dense base + # with bf16 compute normalizes cleanly to "fp8". + norm = _cfg(base_precision = " FP8 ", mixed_precision = "bf16").normalized() + assert norm.base_precision == "fp8" + + # A dense mode against a prequant (bnb-4bit) base is refused: the repo already ships a + # 4-bit transformer and cannot serve the dense precisions. + with pytest.raises(ValueError, match = "dense base repo"): + _cfg(base_model = _Z_PREQUANT, base_precision = "bf16").normalized() + + # A dense mode with non-bf16 compute is refused: these modes train in bf16 compute. + with pytest.raises(ValueError, match = "bf16 compute"): + _cfg(base_precision = "int8", mixed_precision = "fp16").normalized() + + # "auto" is ACCEPTED by normalized() even on a prequant base: the concrete mode is + # resolved at runtime against the live GPU, not at config validation. + assert _cfg(base_model = _Z_PREQUANT, base_precision = "auto").normalized().base_precision == "auto" + + +# ── repo_is_prequantized heuristic + trainer alias ──────────────────────────── +@pytest.mark.parametrize( + "repo, expected", + [ + ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", True), + ("some/model-4bit", True), + ("some/model-int4", True), + ("some/model-nf4", True), + ("black-forest-labs/FLUX.1-dev", False), + ("Tongyi-MAI/Z-Image-Turbo", False), + ], +) +def test_repo_is_prequantized_cases(repo, expected): + assert repo_is_prequantized(repo) is expected + + +def test_repo_is_prequantized_alias_is_same_object(): + # The trainer keeps a module-level alias for callers/tests; it must be the exact same + # function object as the common heuristic (moved there for config validation). + assert dit._repo_is_prequantized is repo_is_prequantized + + +# ── _pick_auto_precision policy table (pure) ────────────────────────────────── +def test_pick_auto_precision_policy_table(): + p = dit._pick_auto_precision + + # A prequant base always resolves to nf4 (it can only serve 4-bit). + assert p(True, "cuda", 140, 23.8, (10, 0), True) == "nf4" + # No CUDA -> nf4 (the dense modes need a GPU). + assert p(False, "cpu", 140, 23.8, (10, 0), True) == "nf4" + # Missing free-VRAM number -> the safe nf4 mode. + assert p(False, "cuda", None, 23.8, (10, 0), True) == "nf4" + + # Plenty of free VRAM -> bf16 regardless of fp8 capability: compiled bf16 measured + # FASTER than torchao float8 at LoRA-training shapes, so fp8 is opt-in only. + assert p(False, "cuda", 140, 23.8, (10, 0), True) == "bf16" + assert p(False, "cuda", 140, 23.8, (8, 0), True) == "bf16" + assert p(False, "cuda", 140, 23.8, (10, 0), False) == "bf16" + + # Middle band (25 > 23.8 * 0.55 * 1.5 = 19.6, but not > 23.8 * 1.5 = 35.7) -> int8. + assert p(False, "cuda", 25, 23.8, (10, 0), True) == "int8" + # Too little free VRAM for even int8 -> nf4. + assert p(False, "cuda", 10, 23.8, (10, 0), True) == "nf4" + + +# ── _resolve_base_precision passthrough ─────────────────────────────────────── +def test_resolve_base_precision_passes_explicit_through(): + # An explicit mode passes straight through without probing the GPU (normalized() already + # validated it); the spec is only consulted for "auto". + spec = dit._SPECS["flux.1"] + cfg = _cfg(base_precision = "bf16") + assert dit._resolve_base_precision(cfg, spec, "cuda") == "bf16" + # Device is irrelevant for an explicit mode: still bf16 on cpu. + assert dit._resolve_base_precision(cfg, spec, "cpu") == "bf16" + + +# ── _fp8_module_filter ──────────────────────────────────────────────────────── +def test_fp8_module_filter(): + lin = nn.Linear(64, 64) + # A plain feed-forward Linear with divisible dims gets float8 training compute. + assert dit._fp8_module_filter(lin, "transformer_blocks.0.ff.net.0") is True + # A LoRA-owned module is skipped (adapters stay high precision). + assert dit._fp8_module_filter(lin, "transformer_blocks.0.attn.to_q.lora_A.default") is False + # The output projection is skipped. + assert dit._fp8_module_filter(lin, "proj_out") is False + # An in_features not divisible by 16 is rejected (float8 kernels reject the shape). + assert dit._fp8_module_filter(nn.Linear(30, 64), "transformer_blocks.0.ff.net.0") is False + # A non-Linear module is never float8. + assert dit._fp8_module_filter(nn.LayerNorm(64), "transformer_blocks.0.norm") is False + + +# ── _should_compile fp8 branch ──────────────────────────────────────────────── +def test_should_compile_fp8_branch(): + # fp8 is only competitive compiled, so auto arms compile for it on a dense (non-bnb) + # cuda base. + cfg = _cfg(compile_transformer = "auto") + assert dit._should_compile(cfg, False, "cuda", "fp8") is True + # fp8 forces compile under auto even when the base is (hypothetically) reported as bnb. + assert dit._should_compile(cfg, True, "cuda", "fp8") is True + # An explicit "off" still wins over fp8 -- compile stays off. + assert dit._should_compile(_cfg(compile_transformer = "off"), False, "cuda", "fp8") is False + + +# ── train_precision_modes machine probe ─────────────────────────────────────── +def test_train_precision_modes_no_cuda(monkeypatch): + # Patch the torch module attribute the function imports so it observes a CPU-only box: + # no CUDA -> the nf4-only floor with nf4 recommended, and it never raises. + import torch + + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + assert train_precision_modes() == (["nf4"], "nf4") + + +# ── family_train_infos precision fields ─────────────────────────────────────── +def test_family_train_infos_carries_precision_fields(monkeypatch): + # Pin the machine probe so the DiT families carry a deterministic mode list, while SDXL + # (no precision selector) stays empty regardless of the probe. + monkeypatch.setattr(common, "train_precision_modes", lambda: (["nf4", "bf16"], "auto")) + infos = {i["name"]: i for i in common.family_train_infos()} + + flux = infos["flux.1"] + assert flux["precision_modes"] == ["nf4", "bf16"] + assert flux["recommended_precision"] == "auto" + assert flux["supports_compile"] is True + + sdxl = infos["sdxl"] + assert sdxl["precision_modes"] == [] + assert sdxl["recommended_precision"] == "nf4" + assert sdxl["supports_compile"] is False + + +# ── request model base_precision field ──────────────────────────────────────── +def test_request_model_base_precision(): + # The request defaults to the nf4 memory floor. + req = DiffusionTrainingStartRequest(base_model = "x", data_dir = "d", output_dir = "o") + assert req.base_precision == "nf4" + + # An allowed dense mode is accepted. + assert ( + DiffusionTrainingStartRequest( + base_model = "x", data_dir = "d", output_dir = "o", base_precision = "fp8" + ).base_precision + == "fp8" + ) + + # An out-of-Literal value is rejected by pydantic. + with pytest.raises(Exception): + DiffusionTrainingStartRequest( + base_model = "x", data_dir = "d", output_dir = "o", base_precision = "int4" + ) + + # The generic Studio dict path carries base_precision through onto DiffusionLoraConfig. + cfg = _config_from_dict( + { + "base_model": _FLUX_DENSE, + "data_dir": "d", + "output_dir": "o", + "base_precision": "bf16", + } + ) + assert cfg.base_precision == "bf16" diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py index b8d805ead2..27a184a2fa 100644 --- a/studio/backend/tests/test_diffusion_train_perf.py +++ b/studio/backend/tests/test_diffusion_train_perf.py @@ -194,8 +194,11 @@ def test_should_compile_policy(): assert _should_compile(_cfg(compile_transformer = "on"), False, "cuda") is True # auto stays off over a bitsandbytes base (graph breaks in the dequant path). assert _should_compile(_cfg(compile_transformer = "auto"), True, "cuda") is False - # auto turns on for a dense (non-bnb) base on cuda. - assert _should_compile(_cfg(compile_transformer = "auto"), False, "cuda") is True + # auto turns on for the dense bf16 base precision on cuda. + assert ( + _should_compile(_cfg(compile_transformer = "auto"), False, "cuda", base_precision = "bf16") + is True + ) # Any mode is a no-op on cpu. for mode in ("off", "on", "auto"): assert _should_compile(_cfg(compile_transformer = mode), False, "cpu") is False