From cb9247e5376f5d58eadd48e6572d2b7e8277bbd7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 19:13:03 +0000 Subject: [PATCH 01/18] Add mxfp8 training base precision and SDXL U-Net regional compile - base_precision="mxfp8": torchao MX block-scaled float8 compute on the frozen base linears (Blackwell sm100+, cuBLAS kernels). Applied after add_adapter like fp8, never fatal, weights stay bf16 in memory. Measured 1.16x over compiled bf16 on Z-Image at 1024px batch 4 (16k tokens/step); a wash at small token counts, so it stays an explicit opt-in and auto never picks it. - SDXL: regionally compile the U-Net's BasicTransformerBlocks through the same never-fatal wrapper the DiT trainer uses. 1.35x steady state at 1024px batch 4 with same-seed loss parity (~1e-5 per step) and unchanged peak VRAM; ~30 s one-time warmup. Steady-state samples/sec now excludes step 1, matching the DiT trainer. - /info: mxfp8 advertised only on sm100+; supports_compile now true for sdxl. - NVFP4 training: not available in torchao 0.16 (no autograd path, no training recipe), so NVFP4 stays an inference-only quant for now. 193 diffusion backend tests green; frontend build clean. --- .../core/training/diffusion_dit_trainer.py | 69 +++++++++-- .../core/training/diffusion_lora_trainer.py | 34 +++-- .../core/training/diffusion_train_common.py | 25 ++-- studio/backend/models/training.py | 3 +- .../tests/test_diffusion_dit_trainer.py | 117 +++++++++++++++++- .../tests/test_diffusion_lora_trainer.py | 42 +++++++ .../backend/tests/test_diffusion_training.py | 14 +++ studio/frontend/src/features/images/api.ts | 7 +- .../images/train/diffusion-train-panel.tsx | 24 ++-- 9 files changed, 294 insertions(+), 41 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 916e5f95d8..984b7b5445 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -218,8 +218,8 @@ def _load_dit_transformer(transformer_cls, cfg, device, 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). + - bf16 / fp8 / mxfp8: the dense transformer (fp8/mxfp8 convert 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 @@ -238,7 +238,7 @@ def _load_dit_transformer(transformer_cls, cfg, device, base_precision): transformer = transformer.to(device) return transformer - # Dense load for bf16 / fp8 / int8. int8 quantizes AFTER the LoRA attaches (see + # Dense load for bf16 / fp8 / mxfp8 / 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). @@ -299,6 +299,45 @@ def _apply_fp8_training(transformer, on_event) -> bool: return False +def _mx_module_filter(mod, fqn: str) -> bool: + """Which frozen linears get mxfp8 training compute: skip anything LoRA-owned (the + adapters must stay high precision), the output projection (same guard as fp8), and + shapes the 32-wide MX block scaling cannot tile (dims not divisible by 32).""" + 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 % 32 == 0 and mod.out_features % 32 == 0 + + +def _apply_mxfp8_training(transformer, on_event) -> bool: + """Swap the frozen base linears to torchao MX float8 training compute (mxfp8, the + Blackwell-native block-scaled format; the swap is in place and the weights stay bf16 + in memory, so like fp8 this is a speed mode, not a memory mode). Applied AFTER + add_adapter so the filter can exclude the LoRA modules. Only competitive under + torch.compile and only ahead of compiled bf16 at large token counts (high resolution + or batch), which is why it stays an explicit opt-in rather than an "auto" pick. + Never fatal: on any failure the run continues in bf16 with a warning.""" + try: + from torchao.prototype.mx_formats import MXLinearConfig + from torchao.quantization import quantize_ + quantize_( + transformer, + MXLinearConfig.from_recipe_name("mxfp8_cublas"), + filter_fn = _mx_module_filter, + ) + return True + except Exception as exc: # noqa: BLE001 -- mxfp8 is an optimisation, never fatal + _emit( + on_event, "warning", message = f"mxfp8 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 @@ -330,7 +369,7 @@ def _resolve_base_precision(cfg, spec, device) -> str: transformer onto the CPU.""" mode = (cfg.base_precision or "nf4").strip().lower() if mode != "auto": - if mode in ("bf16", "int8", "fp8") and device != "cuda": + if mode in ("bf16", "int8", "fp8", "mxfp8") and device != "cuda": raise ValueError( f"base_precision={mode!r} needs a CUDA GPU; this host has none. " f"Use base_precision='nf4' or 'auto'." @@ -977,8 +1016,9 @@ def _should_compile( return True # 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") + # dequant path), so it stays off for QLoRA. fp8/mxfp8 are only competitive compiled + # (eager, their per-matmul dynamic casts run 4-5x slower than bf16). + return base_precision in ("bf16", "fp8", "mxfp8") def _maybe_compile_transformer( @@ -994,11 +1034,14 @@ def _maybe_compile_transformer( 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, base_precision): - if base_precision == "fp8": + if base_precision in ("fp8", "mxfp8"): _emit( on_event, "warning", - message = "fp8 training without torch.compile is slow; enable compile for the speedup.", + message = ( + f"{base_precision} training without torch.compile is slow; " + f"enable compile for the speedup." + ), ) return False import torch @@ -1177,8 +1220,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto variant_rng = random.Random(cfg.seed + 1) # 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). + # default; bf16 / int8 / fp8 / mxfp8 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" @@ -1207,12 +1250,14 @@ 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] - # int8 / fp8 convert the frozen base linears AFTER the LoRA attaches, so the adapter - # modules are excluded and stay high precision. + # int8 / fp8 / mxfp8 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" + if base_precision == "mxfp8" and not _apply_mxfp8_training(transformer, on_event): + base_precision = "bf16" compiled = _maybe_compile_transformer( transformer, cfg, base_is_bnb, device, on_event, base_precision diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 09d1c5618b..0004b3680f 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -26,7 +26,9 @@ latents are likewise precomputed into a small CPU cache (``cache_latents``) and freed. The cache stores the posterior's affine pair (mean/std, scale folded in), so every step still draws a fresh VAE sample -- distribution-identical to encoding in the loop, without keeping the VAE resident or paying a per-step encode. TF32 matmuls + cudnn -autotuning are enabled for the run under ``cfg.enable_tf32``. +autotuning are enabled for the run under ``cfg.enable_tf32``, and the U-Net's repeated +transformer blocks are regionally torch.compiled (``cfg.compile_transformer``, never +fatal -- any failure falls back to eager with a warning event). """ from __future__ import annotations @@ -353,6 +355,15 @@ def run_diffusion_lora_training( if weight_dtype != torch.float32: cast_training_params(unet, dtype = torch.float32) + # Regionally torch.compile the U-Net's repeated BasicTransformerBlocks through the + # DiT trainer's never-fatal wrapper (a wrap/compile failure falls back to eager + # with a warning event). The U-Net is a dense bf16 base here, the combination that + # wrapper compiles under "auto". + from core.training.diffusion_dit_trainer import _maybe_compile_transformer + compiled = _maybe_compile_transformer( + unet, cfg, False, device, on_event, base_precision = "bf16" + ) + lora_params = [p for p in unet.parameters() if p.requires_grad] optimizer = _make_lora_optimizer(lora_params, cfg.learning_rate) # The scheduler advances once per optimizer update: lr_sched.step() runs a single @@ -429,7 +440,7 @@ def run_diffusion_lora_training( # seed-deterministic sequence whether or not the cache is enabled. variant_rng = random.Random(cfg.seed + 1) - _emit(on_event, "model_load_completed") + _emit(on_event, "model_load_completed", compiled = compiled) def _next_batch() -> tuple[list[int], list[str], list[str]]: idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs))) @@ -442,6 +453,7 @@ def run_diffusion_lora_training( running_loss = 0.0 peak_gb = 0.0 t_start = time.time() + t_steady = None done = 0 for opt_step in range(cfg.train_steps): optimizer.zero_grad(set_to_none = True) @@ -524,17 +536,25 @@ def run_diffusion_lora_training( running_loss += step_loss done = opt_step + 1 + now = time.time() + if done == 1: + # Step 1 pays the one-time costs (cudnn autotune, torch.compile warmup), so + # the reported rate starts after it and reflects the steady state (the DiT + # trainer does the same). + t_steady = now if done % cfg.log_every == 0 or done == cfg.train_steps: # ``learning_rate`` (not ``lr``) is the field the Studio training pump reads, so # these progress events are directly consumable by the existing training # status/SSE machinery when the diffusion trainer is wired into the worker. if device == "cuda": peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2) - samples_per_second = round( - (done * cfg.train_batch_size * cfg.gradient_accumulation_steps) - / max(time.time() - t_start, 1e-6), - 3, - ) + per_step = cfg.train_batch_size * cfg.gradient_accumulation_steps + if t_steady is not None and done > 1: + samples_per_second = round( + (done - 1) * per_step / max(now - t_steady, 1e-6), 3 + ) + else: + samples_per_second = round(done * per_step / max(now - t_start, 1e-6), 3) _emit( on_event, "progress", diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 475f8fff30..dc205ebf5b 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -128,15 +128,17 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None 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.""" + (bf16/int8/fp8/mxfp8) 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.""" + always works; bf16/int8/auto need CUDA; fp8 needs an fp8-capable GPU (sm89+); mxfp8 + (block-scaled fp8 compute) needs the Blackwell tensor cores (sm100+) its cuBLAS + kernels target. Used by the /info endpoint so the UI can gate the precision + selector. Never raises.""" modes = ["nf4"] recommended = "nf4" try: @@ -146,6 +148,8 @@ def train_precision_modes() -> tuple[list[str], str]: major, minor = torch.cuda.get_device_capability() if (major, minor) >= (8, 9) and hasattr(torch, "float8_e4m3fn"): modes.append("fp8") + if (major, minor) >= (10, 0): + modes.append("mxfp8") modes.append("auto") recommended = "auto" except Exception: # noqa: BLE001 -- no torch / probe failure -> nf4 only @@ -219,8 +223,9 @@ def family_train_infos() -> list[dict[str, Any]]: 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. + # base_precision applies to the DiT trainer only; SDXL keeps its mixed_precision + # lever, so the UI hides the precision selector for it. compile applies everywhere: + # the SDXL trainer regionally compiles the U-Net's transformer blocks too. is_dit = name in ("flux.1", "qwen-image", "z-image", "krea-2") infos.append( { @@ -232,7 +237,7 @@ def family_train_infos() -> list[dict[str, Any]]: "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, + "supports_compile": True, } ) return infos @@ -329,9 +334,11 @@ class DiffusionLoraConfig: 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 base_precision not in ("nf4", "bf16", "int8", "fp8", "mxfp8", "auto"): + raise ValueError( + "base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto" + ) + if base_precision in ("bf16", "int8", "fp8", "mxfp8"): if repo_is_prequantized(self.base_model): raise ValueError( f"base_precision={base_precision!r} needs a dense base repo, but " diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index aa7b1f29fe..c85a35c880 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -736,11 +736,12 @@ 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( + base_precision: Literal["nf4", "bf16", "int8", "fp8", "mxfp8", "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), " + "mxfp8 block-scaled float8 compute (Blackwell, best at high resolution/batch), " "or auto (pick by free VRAM + GPU class). Dense modes need a non-prequant base." ), ) diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index f4e9c804ca..1ae037feb1 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -9,16 +9,25 @@ exercised by the live GPU smokes, not here.""" from __future__ import annotations +import sys + import pytest from core.training.diffusion_dit_trainer import ( _GATED_TRAIN_REPOS, _SPECS, + _apply_mxfp8_training, _assert_gated_access, + _mx_module_filter, _repo_is_prequantized, + _should_compile, run_dit_lora_training, ) -from core.training.diffusion_train_common import DiffusionLoraConfig, family_train_infos +from core.training.diffusion_train_common import ( + DiffusionLoraConfig, + family_train_infos, + train_precision_modes, +) def test_specs_cover_the_dit_families(): @@ -85,3 +94,109 @@ def test_family_train_infos_lists_dit_families(): assert "gated" in infos["flux.1"]["vram_note"].lower() # Z-Image defaults to the prequant nf4 repo for QLoRA. assert "4bit" in infos["z-image"]["default_base"].lower() + + +def test_family_train_infos_sdxl_supports_compile_without_precision_modes(monkeypatch): + # Regional compile now applies to every family (the SDXL trainer compiles its U-Net + # blocks too), but base_precision stays DiT-only, so SDXL advertises no precision modes + # while a DiT family (z-image) keeps its own. Pin the precision list so the assertion + # holds regardless of the test host's GPU capability. + import core.training.diffusion_train_common as dtc + + monkeypatch.setattr(dtc, "train_precision_modes", lambda: (["nf4", "bf16", "auto"], "auto")) + infos = {i["name"]: i for i in family_train_infos()} + assert infos["sdxl"]["supports_compile"] is True + assert infos["sdxl"]["precision_modes"] == [] + assert infos["z-image"]["supports_compile"] is True + assert infos["z-image"]["precision_modes"] == ["nf4", "bf16", "auto"] + + +# ── mxfp8 base precision (DiT dense speed mode) ─────────────────────────────── +def _linear(in_features, out_features): + import torch.nn as nn + + return nn.Linear(in_features, out_features) + + +def test_mx_module_filter_accepts_dense_block_linear(): + # A 3072x3072 attention/FFN linear at a normal block fqn is a valid mxfp8 target. + assert _mx_module_filter(_linear(3072, 3072), "blocks.0.ff.up") is True + + +def test_mx_module_filter_skips_lora_and_proj_out(): + # LoRA-owned modules (adapters stay high precision) and the output projection are + # excluded, mirroring the fp8 filter's guards. + lin = _linear(3072, 3072) + assert _mx_module_filter(lin, "blocks.0.attn.to_q.lora_A.default") is False + assert _mx_module_filter(lin, "proj_out") is False + assert _mx_module_filter(lin, "x.proj_out.y") is False + + +def test_mx_module_filter_rejects_non_block_aligned_dims(): + # MX block scaling tiles 32-wide, so a dim not divisible by 32 (3000) is rejected. + assert _mx_module_filter(_linear(3000, 3072), "blocks.0.ff.up") is False + + +def test_mx_module_filter_rejects_non_linear(): + import torch.nn as nn + + # A non-Linear module is never a target even if it exposes matching feature counts. + assert _mx_module_filter(nn.LayerNorm(3072), "blocks.0.norm") is False + + +def test_should_compile_auto_mxfp8_on_cuda(): + # auto compiles the dense speed modes on cuda; int8 stays eager (torchao subclass); + # an explicit "off" wins over the mode. + cfg = DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o") + assert _should_compile(cfg, False, "cuda", base_precision = "mxfp8") is True + assert _should_compile(cfg, False, "cuda", base_precision = "int8") is False + off = DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", compile_transformer = "off" + ) + assert _should_compile(off, False, "cuda", base_precision = "mxfp8") is False + + +def test_apply_mxfp8_training_failure_falls_back_with_warning(monkeypatch): + # An unavailable torchao MX path must never be fatal: force the import to raise, then + # assert the helper returns False and emits exactly one warning naming mxfp8. + monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", None) + events = [] + ok = _apply_mxfp8_training(object(), lambda e: events.append(e)) + assert ok is False + warnings = [e for e in events if e["type"] == "warning"] + assert len(warnings) == 1 + assert "mxfp8" in warnings[0]["message"] + + +def _patch_capability(monkeypatch, capability): + # Drive train_precision_modes' GPU probe: pretend CUDA is present at the given tensor + # core capability (fp8 needs sm89+, mxfp8 needs sm100+). + import torch + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability) + + +def test_train_precision_modes_blackwell_lists_mxfp8(monkeypatch): + # sm100 (Blackwell) exposes both fp8 and mxfp8, ordered before the "auto" pick. + _patch_capability(monkeypatch, (10, 0)) + modes, recommended = train_precision_modes() + assert "mxfp8" in modes and "fp8" in modes + assert modes.index("mxfp8") < modes.index("auto") + assert modes.index("fp8") < modes.index("auto") + assert recommended == "auto" + + +def test_train_precision_modes_ada_has_fp8_without_mxfp8(monkeypatch): + # sm89 (Ada) is fp8-capable but not block-scaled mxfp8-capable. + _patch_capability(monkeypatch, (8, 9)) + modes, _ = train_precision_modes() + assert "fp8" in modes + assert "mxfp8" not in modes + + +def test_train_precision_modes_newer_blackwell_has_mxfp8(monkeypatch): + # Any capability >= sm100 keeps mxfp8 (sm120 here). + _patch_capability(monkeypatch, (12, 0)) + modes, _ = train_precision_modes() + assert "mxfp8" in modes diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index 3465408a0e..e1297bf64e 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -103,6 +103,48 @@ def test_config_normalized_validation(kw): DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw).normalized() +def test_config_normalized_accepts_mxfp8_dense_base(): + # mxfp8 is a dense speed mode: a dense base + bf16 compute normalises through. + cfg = DiffusionLoraConfig( + base_model = "black-forest-labs/FLUX.1-dev", + data_dir = "d", + output_dir = "o", + base_precision = "mxfp8", + ).normalized() + assert cfg.base_precision == "mxfp8" + + +def test_config_normalized_mxfp8_rejects_prequant_base(): + # A prequant (bnb-4bit) base cannot serve the dense mxfp8 base precision. + with pytest.raises(ValueError, match = "mxfp8"): + DiffusionLoraConfig( + base_model = "unsloth/Qwen-Image-2512-unsloth-bnb-4bit", + data_dir = "d", + output_dir = "o", + base_precision = "mxfp8", + ).normalized() + + +def test_config_normalized_mxfp8_requires_bf16_compute(): + # Like the other dense modes, mxfp8 trains in bf16 compute; fp16 is refused. + with pytest.raises(ValueError, match = "mxfp8"): + DiffusionLoraConfig( + base_model = "black-forest-labs/FLUX.1-dev", + data_dir = "d", + output_dir = "o", + base_precision = "mxfp8", + mixed_precision = "fp16", + ).normalized() + + +def test_config_normalized_lists_mxfp8_in_invalid_mode_error(): + # The invalid-base_precision message enumerates the allowed modes, including mxfp8. + with pytest.raises(ValueError, match = "mxfp8"): + DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", base_precision = "bogus" + ).normalized() + + def _cfg(**kw): return DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 6320b7fc42..ff83dfc2b8 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -339,6 +339,20 @@ def test_request_model_num_epochs_bounds(): DiffusionTrainingStartRequest(**base, num_epochs = bad) +def test_request_model_base_precision_accepts_mxfp8(): + # The base_precision Literal now includes mxfp8 (the DiT dense speed mode); a bogus + # mode is still rejected. + from pydantic import ValidationError + + from models.training import DiffusionTrainingStartRequest + + base = {"base_model": "b", "data_dir": "d", "output_dir": "o"} + assert DiffusionTrainingStartRequest(**base).base_precision == "nf4" # default + assert DiffusionTrainingStartRequest(**base, base_precision = "mxfp8").base_precision == "mxfp8" + with pytest.raises(ValidationError): + DiffusionTrainingStartRequest(**base, base_precision = "bogus") + + def test_route_start_rejects_uncontained_paths(client): # An absolute path outside the Studio dataset roots is a 400, not silently accepted. r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "/etc"}) diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index eb8ecd785e..9291c2055d 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -287,9 +287,10 @@ export interface DiffusionTrainingStartRequest { lr_warmup_steps?: number; // DiT-family quantised base precision (nf4 QLoRA by default). Ignored for sdxl, which // uses mixed_precision instead. "auto" lets the backend pick per family. - base_precision?: "nf4" | "bf16" | "int8" | "fp8" | "auto"; - // Whether to torch.compile the transformer (DiT families that support it). "auto" lets - // the backend decide; "off"/"on" force it. + base_precision?: "nf4" | "bf16" | "int8" | "fp8" | "mxfp8" | "auto"; + // Whether to torch.compile the transformer (any family whose /info reports + // supports_compile; that includes the SDXL U-Net). "auto" lets the backend decide; + // "off"/"on" force it. compile_transformer?: "off" | "on" | "auto"; // Precompute + cache the VAE latents before the loop (skips re-encoding each epoch). cache_latents?: boolean; diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index a76e8e5209..6ef3599e34 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -182,17 +182,22 @@ export function DiffusionTrainPanel({ const isDiT = familyName !== "sdxl"; // The quantised base precisions this family can train in, with a stable fallback when the // backend does not report them (older backend, or a preset-only family). - const precisionModes = useMemo>(() => { + const precisionModes = useMemo< + Array<"nf4" | "bf16" | "int8" | "fp8" | "mxfp8" | "auto"> + >(() => { const reported = reportedFamily?.precision_modes?.filter( - (m): m is "nf4" | "bf16" | "int8" | "fp8" => - m === "nf4" || m === "bf16" || m === "int8" || m === "fp8", + (m): m is "nf4" | "bf16" | "int8" | "fp8" | "mxfp8" => + m === "nf4" || m === "bf16" || m === "int8" || m === "fp8" || m === "mxfp8", ); if (reported && reported.length > 0) return ["auto", ...reported]; + // Fallback without a backend report: the GPU-independent modes only (mxfp8 needs a + // Blackwell probe, so it is offered strictly when the backend advertises it). return ["auto", "nf4", "bf16", "int8", "fp8"]; }, [reportedFamily?.precision_modes]); - // Whether to show the torch.compile control. Default on for DiT families when the backend - // does not say otherwise; sdxl's U-Net path does not expose it here. - const supportsCompile = isDiT && (reportedFamily?.supports_compile ?? true); + // Whether to show the torch.compile control. The backend advertises this per family + // (the SDXL U-Net path compiles regionally too now); default on for DiT families when + // an older backend does not report it. + const supportsCompile = reportedFamily?.supports_compile ?? isDiT; const [baseChoice, setBaseChoice] = useState(family?.base_repos[0] ?? ""); const [customBase, setCustomBase] = useState(""); @@ -235,7 +240,7 @@ export function DiffusionTrainPanel({ // lets the backend pick the family's recommended mode. Re-seeded to the family's // recommendation on family change (unless the user picked one). const [basePrecision, setBasePrecision] = useState< - "nf4" | "bf16" | "int8" | "fp8" | "auto" + "nf4" | "bf16" | "int8" | "fp8" | "mxfp8" | "auto" >("auto"); // Whether to torch.compile the DiT transformer. "auto" defers to the backend. const [compileTransformer, setCompileTransformer] = useState<"off" | "on" | "auto">( @@ -737,11 +742,14 @@ export function DiffusionTrainPanel({ ); - const precisionLabel = (m: "nf4" | "bf16" | "int8" | "fp8" | "auto"): string => { + const precisionLabel = ( + m: "nf4" | "bf16" | "int8" | "fp8" | "mxfp8" | "auto", + ): string => { if (m === "auto") return "Auto (recommended)"; if (m === "nf4") return "nf4 (4-bit QLoRA, lowest VRAM)"; if (m === "bf16") return "bf16 (fastest, most VRAM)"; if (m === "int8") return "int8 (8-bit)"; + if (m === "mxfp8") return "mxfp8 (Blackwell, best at high res/batch)"; return "fp8 (experimental)"; }; From 4a3764e1d426652e80e37dc6ba8e071d5df4a62e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 3 Jul 2026 19:15:07 +0000 Subject: [PATCH 02/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_dit_trainer.py | 5 ++--- studio/backend/core/training/diffusion_lora_trainer.py | 5 ++--- studio/backend/core/training/diffusion_train_common.py | 4 +--- studio/backend/tests/test_diffusion_dit_trainer.py | 2 -- 4 files changed, 5 insertions(+), 11 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 984b7b5445..17b5aff61e 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -325,6 +325,7 @@ def _apply_mxfp8_training(transformer, on_event) -> bool: try: from torchao.prototype.mx_formats import MXLinearConfig from torchao.quantization import quantize_ + quantize_( transformer, MXLinearConfig.from_recipe_name("mxfp8_cublas"), @@ -332,9 +333,7 @@ def _apply_mxfp8_training(transformer, on_event) -> bool: ) return True except Exception as exc: # noqa: BLE001 -- mxfp8 is an optimisation, never fatal - _emit( - on_event, "warning", message = f"mxfp8 training unavailable, using bf16 compute: {exc}" - ) + _emit(on_event, "warning", message = f"mxfp8 training unavailable, using bf16 compute: {exc}") return False diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 0004b3680f..d22b7f737e 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -360,6 +360,7 @@ def run_diffusion_lora_training( # with a warning event). The U-Net is a dense bf16 base here, the combination that # wrapper compiles under "auto". from core.training.diffusion_dit_trainer import _maybe_compile_transformer + compiled = _maybe_compile_transformer( unet, cfg, False, device, on_event, base_precision = "bf16" ) @@ -550,9 +551,7 @@ def run_diffusion_lora_training( peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2) per_step = cfg.train_batch_size * cfg.gradient_accumulation_steps if t_steady is not None and done > 1: - samples_per_second = round( - (done - 1) * per_step / max(now - t_steady, 1e-6), 3 - ) + samples_per_second = round((done - 1) * per_step / max(now - t_steady, 1e-6), 3) else: samples_per_second = round(done * per_step / max(now - t_start, 1e-6), 3) _emit( diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index dc205ebf5b..c867fde7b1 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -335,9 +335,7 @@ class DiffusionLoraConfig: 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", "mxfp8", "auto"): - raise ValueError( - "base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto" - ) + raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto") if base_precision in ("bf16", "int8", "fp8", "mxfp8"): if repo_is_prequantized(self.base_model): raise ValueError( diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index 1ae037feb1..b2dbc969e4 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -114,7 +114,6 @@ def test_family_train_infos_sdxl_supports_compile_without_precision_modes(monkey # ── mxfp8 base precision (DiT dense speed mode) ─────────────────────────────── def _linear(in_features, out_features): import torch.nn as nn - return nn.Linear(in_features, out_features) @@ -172,7 +171,6 @@ def _patch_capability(monkeypatch, capability): # Drive train_precision_modes' GPU probe: pretend CUDA is present at the given tensor # core capability (fp8 needs sm89+, mxfp8 needs sm100+). import torch - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability) From 1521172b9d603df03751b25a85b2bd6a74bdf3c7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 19:18:54 +0000 Subject: [PATCH 03/18] Show the compile control for SDXL in the Train settings The compile select was nested inside the DiT-only branch of the precision area, so supports_compile=true for sdxl never rendered it. Lift it out of the ternary: the precision selects stay family-specific, the compile control follows supports_compile. Verified with a live Playwright pass (mxfp8 listed for DiT families on sm100, compile select present for SDXL, DiT precision selector still hidden for SDXL). --- .../images/train/diffusion-train-panel.tsx | 90 +++++++++---------- 1 file changed, 44 insertions(+), 46 deletions(-) diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 6ef3599e34..d3f089a42c 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -812,52 +812,29 @@ export function DiffusionTrainPanel({ {isDiT ? ( - <> -
- - -

- How the frozen base weights are quantised. nf4 (4-bit) uses the least VRAM; - bf16 is fastest but needs the most. Auto picks this family's recommended - mode. -

-
- {supportsCompile && ( -
- - -

- torch.compile the transformer. Adds a one-time warmup, then speeds up each - step. -

-
- )} - +
+ + +

+ How the frozen base weights are quantised. nf4 (4-bit) uses the least VRAM; + bf16 is fastest but needs the most. Auto picks this family's recommended + mode. +

+
) : (
@@ -876,6 +853,27 @@ export function DiffusionTrainPanel({

)} + {supportsCompile && ( +
+ + +

+ torch.compile the transformer blocks. Adds a one-time warmup, then speeds up + each step. +

+
+ )} ); From f30cdc939ef2f4673a8e7402147d912f9731260c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 01:10:01 +0000 Subject: [PATCH 04/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/sd_cpp_backend.py | 5 ++++- studio/backend/tests/test_sd_cpp_backend.py | 14 +++++++++++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 7632551661..d8cb21d409 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -893,7 +893,10 @@ class SdCppDiffusionBackend: lora_stage = Path(server_lora_dir) / f"gen_{os.urandom(6).hex()}" materialized = diffusion_lora.materialize_native_dir(lora_resolved, lora_stage) lora_payload = [ - {"path": f"{lora_stage.name}/{Path(m.path).name}", "multiplier": float(m.weight)} + { + "path": f"{lora_stage.name}/{Path(m.path).name}", + "multiplier": float(m.weight), + } for m in materialized ] try: diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 0b52ea6e14..be2fc22a4b 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -691,7 +691,11 @@ def _fake_materialize(resolved, dest): return out -def _patch_lora(monkeypatch, resolved, supported = True): +def _patch_lora( + monkeypatch, + resolved, + supported = True, +): from core.inference import diffusion_lora as dl monkeypatch.setattr(dl, "supports_lora", lambda **k: supported) @@ -706,7 +710,9 @@ def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch): eng = _FakeEngine() b = _loaded_backend(engine = eng) # mode = "oneshot" - _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]) + _patch_lora( + monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)] + ) b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)]) _, params, _, _ = eng.calls[0] assert params.lora_dir is not None and params.lora_apply_mode == "auto" @@ -725,7 +731,9 @@ def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tm servers: list = [] _run_server_load(monkeypatch, b, servers) servers[0].lora_dir = str(tmp_path) - _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]) + _patch_lora( + monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)] + ) b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)]) payload = servers[0].payloads[0] assert "lora" in payload and len(payload["lora"]) == 1 From c3b77c08d6560ba684fac8a996dae5dc9133bcd9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 01:34:35 +0000 Subject: [PATCH 05/18] Update sdxl supports_compile expectation after the U-Net compile change The SDXL trainer now regionally compiles its transformer blocks, so /info advertises supports_compile for every family; the precision selector stays DiT-only. --- studio/backend/tests/test_diffusion_base_precision.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index 1f8ca9318f..4c27efcc60 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -265,7 +265,9 @@ def test_family_train_infos_carries_precision_fields(monkeypatch): sdxl = infos["sdxl"] assert sdxl["precision_modes"] == [] assert sdxl["recommended_precision"] == "nf4" - assert sdxl["supports_compile"] is False + # The SDXL trainer regionally compiles its U-Net blocks too, so compile is advertised + # for every family; only the precision selector stays DiT-only. + assert sdxl["supports_compile"] is True # ── request model base_precision field ──────────────────────────────────────── From 8c7bf6d61bf9008fab49a4839d4b1fe1069848f2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 01:35:06 +0000 Subject: [PATCH 06/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_train_common.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index fd49c33423..4776645518 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -335,9 +335,7 @@ class DiffusionLoraConfig: 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", "mxfp8", "auto"): - raise ValueError( - "base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto" - ) + raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto") # base_precision is a DiT-only lever (the transformer load precision); SDXL uses its # own mixed_precision path and ignores base_precision entirely, so the dense-mode # gates (prequant base / non-bf16 compute) apply only to the DiT families. The From 5ea8eb958cebe2ea7064e1157fbcd6f5402217c8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 03:23:33 +0000 Subject: [PATCH 07/18] Support the torchao 0.17 mxfp8 recipe API torchao 0.17 removed MXLinearConfig from prototype.mx_formats in favour of MXFP8TrainingOpConfig.from_recipe shared with MoE training. _mxfp8_training_config tries the 0.16 API first and falls back to the 0.17 one; both feed quantize_. mxfp8 still degrades to bf16 with a warning when neither import resolves --- .../core/training/diffusion_dit_trainer.py | 19 +++++++++- .../tests/test_diffusion_dit_trainer.py | 38 ++++++++++++++++++- 2 files changed, 53 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index fd219c6794..e571e9bd0c 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -314,6 +314,22 @@ def _mx_module_filter(mod, fqn: str) -> bool: return mod.in_features % 32 == 0 and mod.out_features % 32 == 0 +def _mxfp8_training_config(): + """The torchao MX training config across the prototype API's revisions: torchao 0.16 + ships ``MXLinearConfig`` in ``prototype.mx_formats``; 0.17 removed it in favour of the + ``MXFP8TrainingOpConfig`` recipe API shared with MoE training. Both feed ``quantize_``. + Raises ImportError when neither API exists (mxfp8 then falls back to bf16).""" + try: + from torchao.prototype.mx_formats import MXLinearConfig + return MXLinearConfig.from_recipe_name("mxfp8_cublas") + except ImportError: + from torchao.prototype.moe_training.config import ( + MXFP8TrainingOpConfig, + MXFP8TrainingRecipe, + ) + return MXFP8TrainingOpConfig.from_recipe(MXFP8TrainingRecipe.MXFP8_RCEIL) + + def _apply_mxfp8_training(transformer, on_event) -> bool: """Swap the frozen base linears to torchao MX float8 training compute (mxfp8, the Blackwell-native block-scaled format; the swap is in place and the weights stay bf16 @@ -323,12 +339,11 @@ def _apply_mxfp8_training(transformer, on_event) -> bool: or batch), which is why it stays an explicit opt-in rather than an "auto" pick. Never fatal: on any failure the run continues in bf16 with a warning.""" try: - from torchao.prototype.mx_formats import MXLinearConfig from torchao.quantization import quantize_ quantize_( transformer, - MXLinearConfig.from_recipe_name("mxfp8_cublas"), + _mxfp8_training_config(), filter_fn = _mx_module_filter, ) return True diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index b2dbc969e4..ee7a6f3353 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -156,9 +156,11 @@ def test_should_compile_auto_mxfp8_on_cuda(): def test_apply_mxfp8_training_failure_falls_back_with_warning(monkeypatch): - # An unavailable torchao MX path must never be fatal: force the import to raise, then - # assert the helper returns False and emits exactly one warning naming mxfp8. + # An unavailable torchao MX path must never be fatal: force both API revisions' + # imports to raise, then assert the helper returns False and emits exactly one + # warning naming mxfp8. monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", None) + monkeypatch.setitem(sys.modules, "torchao.prototype.moe_training.config", None) events = [] ok = _apply_mxfp8_training(object(), lambda e: events.append(e)) assert ok is False @@ -167,6 +169,38 @@ def test_apply_mxfp8_training_failure_falls_back_with_warning(monkeypatch): assert "mxfp8" in warnings[0]["message"] +def test_mxfp8_training_config_falls_back_to_the_torchao_0_17_api(monkeypatch): + # torchao 0.17 removed prototype.mx_formats.MXLinearConfig in favour of the + # MXFP8TrainingOpConfig recipe API; the config helper must fall back to it so the + # advertised mxfp8 mode keeps engaging on those installs instead of silently + # training dense bf16. + from types import SimpleNamespace + + from core.training.diffusion_dit_trainer import _mxfp8_training_config + + calls = {} + + class _Recipe: + MXFP8_RCEIL = "mxfp8_rceil" + + class _OpConfig: + @staticmethod + def from_recipe(recipe): + calls["recipe"] = recipe + return "cfg-0.17" + + fake_config = SimpleNamespace( + MXFP8TrainingOpConfig = _OpConfig, MXFP8TrainingRecipe = _Recipe + ) + monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", None) + monkeypatch.setitem( + sys.modules, "torchao.prototype.moe_training", SimpleNamespace(config = fake_config) + ) + monkeypatch.setitem(sys.modules, "torchao.prototype.moe_training.config", fake_config) + assert _mxfp8_training_config() == "cfg-0.17" + assert calls["recipe"] == _Recipe.MXFP8_RCEIL + + def _patch_capability(monkeypatch, capability): # Drive train_precision_modes' GPU probe: pretend CUDA is present at the given tensor # core capability (fp8 needs sm89+, mxfp8 needs sm100+). From 003e28730cf052fe33d714ca683a6d7274c4d8a8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 03:31:53 +0000 Subject: [PATCH 08/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../backend/core/training/diffusion_dit_trainer.py | 1 - studio/backend/routes/models.py | 4 +--- studio/backend/tests/test_diffusion_dataset_api.py | 6 ++++-- studio/backend/tests/test_diffusion_dit_trainer.py | 9 +++++---- studio/backend/tests/test_diffusion_training.py | 6 ++++-- studio/backend/tests/test_local_model_format.py | 13 +++++++++---- 6 files changed, 23 insertions(+), 16 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 105fd0a822..de98350821 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -356,7 +356,6 @@ def _apply_mxfp8_training(transformer, on_event) -> bool: Never fatal: on any failure the run continues in bf16 with a warning.""" try: from torchao.quantization import quantize_ - quantize_( transformer, _mxfp8_training_config(), diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 54bc490052..d4501517c8 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -911,9 +911,7 @@ async def list_local_models( models = collect_local_models(models_root) # Tag each model with its task so the Images picker can filter to diffusion # (GGUF by architecture; local diffusers checkpoints by pipeline / family). - models = [ - m.model_copy(update = {"task": _local_model_task(m)}) for m in models - ] + models = [m.model_copy(update = {"task": _local_model_task(m)}) for m in models] return LocalModelListResponse( models_dir = str(models_root), diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index e1c2c0a186..b10dd04f4c 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -64,8 +64,10 @@ def test_list_images_caption_precedence(client, ds_root): # a.png -> sidecar (an explicit edit beats the metadata row), b.png -> metadata-only, # c.png -> none. (folder / "metadata.jsonl").write_text( - json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n" - + json.dumps({"file_name": "b.png", "text": "from metadata"}) + "\n", + json.dumps({"file_name": "a.png", "text": "from metadata"}) + + "\n" + + json.dumps({"file_name": "b.png", "text": "from metadata"}) + + "\n", encoding = "utf-8", ) (folder / "a.txt").write_text("edited sidecar", encoding = "utf-8") diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index 8d46200e23..b58868c701 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -67,7 +67,10 @@ def test_select_lora_targets_explicit_override_wins(): base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o" ).normalized() assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS - assert _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) == _FLUX_TARGETS + assert ( + _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) + == _FLUX_TARGETS + ) @pytest.mark.parametrize( @@ -215,9 +218,7 @@ def test_mxfp8_training_config_falls_back_to_the_torchao_0_17_api(monkeypatch): calls["recipe"] = recipe return "cfg-0.17" - fake_config = SimpleNamespace( - MXFP8TrainingOpConfig = _OpConfig, MXFP8TrainingRecipe = _Recipe - ) + fake_config = SimpleNamespace(MXFP8TrainingOpConfig = _OpConfig, MXFP8TrainingRecipe = _Recipe) monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", None) monkeypatch.setitem( sys.modules, "torchao.prototype.moe_training", SimpleNamespace(config = fake_config) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index d3748eea35..fd28213294 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -561,8 +561,10 @@ def test_diffusion_info_counts_metadata_captions(client, dataset_roots): (folder / "c.png").write_bytes(b"x") # a.png + b.png via metadata; a.png also has a sidecar (must count once); c.png none. (folder / "metadata.jsonl").write_text( - json.dumps({"file_name": "a.png", "text": "cap a"}) + "\n" - + json.dumps({"file_name": "b.png", "text": "cap b"}) + "\n", + json.dumps({"file_name": "a.png", "text": "cap a"}) + + "\n" + + json.dumps({"file_name": "b.png", "text": "cap b"}) + + "\n", encoding = "utf-8", ) (folder / "a.txt").write_text("edited a", encoding = "utf-8") diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index dd888eb50d..f9184774d6 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -121,7 +121,14 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): from models.models import LocalModelInfo # noqa: E402 -def _local(path, *, model_format = None, model_id = None, display_name = "m", id = "m"): +def _local( + path, + *, + model_format = None, + model_id = None, + display_name = "m", + id = "m", +): return LocalModelInfo( id = id, display_name = display_name, @@ -147,9 +154,7 @@ def test_local_task_tags_diffusers_by_family_id(tmp_path): d = tmp_path / "flux-checkpoint" _touch(d / "flux1-dev.safetensors") assert ( - models_route._local_model_task( - _local(d, model_id = "black-forest-labs/FLUX.1-dev") - ) + models_route._local_model_task(_local(d, model_id = "black-forest-labs/FLUX.1-dev")) == "text-to-image" ) From 3204874f7eb81ca7f3ebe7a622e606d8f070fcc2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 04:42:24 +0000 Subject: [PATCH 09/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_train_common.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index b91f719395..5cbd82d966 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -498,7 +498,9 @@ def discover_image_caption_pairs( # 2. metadata row keyed by file name (basename or the relative path; as_posix so a # Windows backslash path still matches the jsonl's forward-slash keys). if caption is None: - caption = meta_caption.get(img.name) or meta_caption.get(img.relative_to(root).as_posix()) + caption = meta_caption.get(img.name) or meta_caption.get( + img.relative_to(root).as_posix() + ) # 3. dreambooth instance prompt. if caption is None and instance_prompt: caption = instance_prompt From 5df56c080400c438f40fb4eae0d6486536fcbe98 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 08:21:07 +0000 Subject: [PATCH 10/18] Stub the torchao probe in the precision-mode capability tests train_precision_modes gates int8/fp8/mxfp8 on has_functional_torchao, and the Backend CI runner does not install torchao, so the three capability-gating tests collapsed to nf4/bf16/auto and failed. They exercise the CAPABILITY gate, not torchao presence: stub the probe functional alongside the CUDA capability patch. Validated with a torchao-blocked run (22 passed). --- studio/backend/tests/test_diffusion_dit_trainer.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index b58868c701..6d7736a106 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -230,10 +230,15 @@ def test_mxfp8_training_config_falls_back_to_the_torchao_0_17_api(monkeypatch): def _patch_capability(monkeypatch, capability): # Drive train_precision_modes' GPU probe: pretend CUDA is present at the given tensor - # core capability (fp8 needs sm89+, mxfp8 needs sm100+). + # core capability (fp8 needs sm89+, mxfp8 needs sm100+). The torchao probe is stubbed + # functional so these tests exercise the CAPABILITY gate on hosts without torchao + # (the CPU-only CI runner does not install it). import torch + + import core.training.diffusion_train_common as dtc monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability) + monkeypatch.setattr(dtc, "has_functional_torchao", lambda: True) def test_train_precision_modes_blackwell_lists_mxfp8(monkeypatch): From 310849726efd772974ef4d6932099a27c0b477a2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 08:21:39 +0000 Subject: [PATCH 11/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_dit_trainer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index 6d7736a106..d406bf99eb 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -236,6 +236,7 @@ def _patch_capability(monkeypatch, capability): import torch import core.training.diffusion_train_common as dtc + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability) monkeypatch.setattr(dtc, "has_functional_torchao", lambda: True) From 1bc77d0fd89013a248f3b704f43403f5045a9929 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 02:12:56 +0000 Subject: [PATCH 12/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/sd_cpp_backend.py | 4 +--- studio/backend/tests/test_diffusion_backend.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 78425833d6..75572b500a 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -698,9 +698,7 @@ class SdCppDiffusionBackend: loading = self._loading if loading is None or loading.error is not None: return () - return tuple( - r for r in (loading.repo_id, loading.base_repo, *loading.asset_repos) if r - ) + return tuple(r for r in (loading.repo_id, loading.base_repo, *loading.asset_repos) if r) # ── Generate ─────────────────────────────────────────────────────────── diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 8341bf46f6..5bf371f633 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1414,9 +1414,7 @@ def test_bad_mode_strings_fail_before_eviction(fake_runtime): {"text_encoder_quant": "fp3"}, ): with pytest.raises(ValueError): - backend.load_pipeline( - "unsloth/Z-Image-GGUF", gguf_filename = "m.gguf", **kwargs - ) + backend.load_pipeline("unsloth/Z-Image-GGUF", gguf_filename = "m.gguf", **kwargs) assert backend._state is not None From 1a928afbec01d2d09ca21e6f7a0a33a0aba27809 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 04:45:50 +0000 Subject: [PATCH 13/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_backend.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 345839b026..8c8332ded1 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2144,8 +2144,7 @@ def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch): ) assert root == "/cache/snap" assert ( - backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None) - is None + backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None) is None ) From ea1aa16d6f492f2e6ab6edf573c1e3cfd875b820 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 04:59:09 +0000 Subject: [PATCH 14/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_controlnet.py | 1 + 1 file changed, 1 insertion(+) diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index 3d4f7c98cf..50bf7159b0 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -263,6 +263,7 @@ def test_controlnet_pipe_rejects_family_without_classes(): with pytest.raises(ValueError, match = "not supported"): b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event()) + def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch): # An unload that lands while from_pipe is assembling must not let the wrapper # repopulate the cache around the torn-down base pipe. From 1c40fa855f6954939faef7276a464fc791b3a1c1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 07:59:21 +0000 Subject: [PATCH 15/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_train_common.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 30b6ec727a..bdf53d5b6b 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -597,7 +597,7 @@ def _plan_cache_variants( # fallback. Over this budget the default falls back to per-step VAE encoding. A fixed # constant (rather than a psutil RAM fraction) keeps the gate dependency-free and identical # across hosts; it is deliberately conservative, well under a typical training host's RAM. -_LATENT_CACHE_BUDGET_BYTES = 4 * 1024 ** 3 # 4 GiB +_LATENT_CACHE_BUDGET_BYTES = 4 * 1024**3 # 4 GiB # Returned by the cache builders when the estimated cache exceeds the budget: the caller # keeps the VAE resident and encodes each step's latents in-loop. A distinct sentinel from @@ -614,7 +614,9 @@ def _latent_cache_forced() -> bool: def _latent_cache_over_budget( - per_variant_bytes: int, total_variants: int, budget_bytes: Optional[int] = None + per_variant_bytes: int, + total_variants: int, + budget_bytes: Optional[int] = None, ) -> bool: """True when a cache of ``total_variants`` entries, each two fp32 tensors totalling ``per_variant_bytes``, is estimated to exceed ``budget_bytes``. ``per_variant_bytes`` is From 83bc33d04adce79673fcfb96f65f598ec145cba0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 11:43:54 +0000 Subject: [PATCH 16/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/diffusion.py | 1 - studio/backend/tests/test_diffusion_controlnet.py | 8 ++++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index a810bd22e1..7ce31d5ecf 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1564,7 +1564,6 @@ class DiffusionBackend: # local dir the user picked has no Hub scan and is exempt (fail-open there). if not getattr(resolved_cn, "is_local", False): from utils.security import evaluate_file_security - _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None) if _cn_fs.blocked: raise ValueError(_cn_fs.reason) diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index b7ae037197..bcb7393457 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -234,7 +234,6 @@ def _state(): def _allow_cn_security(monkeypatch): """Stub the Hub malware preflight to allow the load (hermetic, no network).""" import utils.security - monkeypatch.setattr( utils.security, "evaluate_file_security", @@ -277,7 +276,12 @@ def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch): class _TrapModel(_FakeCNModel): @classmethod - def from_pretrained(cls, path, torch_dtype = None, token = None): + def from_pretrained( + cls, + path, + torch_dtype = None, + token = None, + ): loaded["called"] = True return super().from_pretrained(path, torch_dtype = torch_dtype, token = token) From 6384fea272b45276ac983f2e0091a8ff70e17816 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 10:41:02 +0000 Subject: [PATCH 17/18] dit trainer: preserve biases under mxfp8, gate explicit mxfp8 to Blackwell - The torchao 0.17 MX training path swaps a matched frozen Linear's weight for a wrapper tensor whose linear override computes input @ weight_t and drops the bias, so mxfp8'ing a biased frozen linear silently loses its bias and corrupts the base output the LoRA regresses against (verified on Blackwell: the bias term is fully dropped). Skip biased linears in _mx_module_filter. - _resolve_base_precision re-checked explicit dense modes against the live device but only rejected CPU, so an explicit mxfp8 request on a non-Blackwell CUDA GPU passed and then crashed at the first MX GEMM after a full dense-transformer load. /info only advertises mxfp8 on sm100+; mirror that gate here and fail fast for a stale or direct client below Blackwell. --- .../core/training/diffusion_dit_trainer.py | 22 ++++++++++++ .../tests/test_diffusion_dit_trainer.py | 35 +++++++++++++++++-- 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 1fdceb4a07..b6c5cd6983 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -331,6 +331,12 @@ def _mx_module_filter(mod, fqn: str) -> bool: return False if fqn.endswith("proj_out") or ".proj_out." in fqn: return False + # Skip biased linears: the torchao 0.17 MX training path swaps the weight for a wrapper tensor + # whose linear override computes input @ weight_t and drops the bias entirely, so an mxfp8'd + # FROZEN base linear would silently lose its bias and change the output the LoRA regresses + # against (verified on Blackwell: the bias term is fully dropped). Keep biased linears in bf16. + if getattr(mod, "bias", None) is not None: + return False return mod.in_features % 32 == 0 and mod.out_features % 32 == 0 @@ -417,6 +423,22 @@ def _resolve_base_precision(cfg, spec, device) -> str: f"base_precision={mode!r} needs a CUDA GPU; this host has none. " f"Use base_precision='nf4' or 'auto'." ) + # mxfp8 needs Blackwell (sm100+): its MX GEMM has no kernel below sm100 and raises at the + # first training step, AFTER a full dense-transformer load. /info only advertises mxfp8 on + # sm100+ (train_precision_modes), so re-check it here to fail fast for a stale or direct + # client on an older CUDA GPU instead of crashing mid-run. + if mode == "mxfp8" and device == "cuda": + try: + import torch + + blackwell = torch.cuda.get_device_capability() >= (10, 0) + except Exception: # noqa: BLE001 -- probe failure -> treat as unsupported, fail fast + blackwell = False + if not blackwell: + raise ValueError( + "base_precision='mxfp8' needs a Blackwell (sm100+) GPU; this GPU is older. " + "Use base_precision='bf16', 'int8', 'nf4', or 'auto'." + ) return mode # auto may only resolve to the dense modes when the run uses bf16 compute, mirroring # the normalized() rule for explicit dense modes; otherwise stay on the nf4 floor. diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index d406bf99eb..aaad553107 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -10,6 +10,7 @@ exercised by the live GPU smokes, not here.""" from __future__ import annotations import sys +import types import pytest @@ -23,6 +24,7 @@ from core.training.diffusion_dit_trainer import ( _assert_gated_access, _mx_module_filter, _repo_is_prequantized, + _resolve_base_precision, _select_lora_targets, _should_compile, run_dit_lora_training, @@ -141,16 +143,43 @@ def test_family_train_infos_sdxl_supports_compile_without_precision_modes(monkey # ── mxfp8 base precision (DiT dense speed mode) ─────────────────────────────── -def _linear(in_features, out_features): +def _linear(in_features, out_features, bias = False): import torch.nn as nn - return nn.Linear(in_features, out_features) + return nn.Linear(in_features, out_features, bias = bias) def test_mx_module_filter_accepts_dense_block_linear(): - # A 3072x3072 attention/FFN linear at a normal block fqn is a valid mxfp8 target. + # A bias-free 3072x3072 attention/FFN linear at a normal block fqn is a valid mxfp8 target. assert _mx_module_filter(_linear(3072, 3072), "blocks.0.ff.up") is True +def test_mx_module_filter_skips_biased_linear(): + # The torchao 0.17 MX training path drops the bias term (its linear override computes + # input @ weight_t only), so an mxfp8'd biased FROZEN linear would silently lose its bias and + # corrupt the base output the LoRA regresses against. Biased linears must stay bf16. + assert _mx_module_filter(_linear(3072, 3072, bias = True), "blocks.0.ff.up") is False + + +def test_resolve_base_precision_explicit_mxfp8_requires_blackwell(monkeypatch): + # An explicit mxfp8 request on a non-Blackwell CUDA GPU must fail fast: its MX GEMM has no + # kernel below sm100 and would otherwise crash at the first training step, after a full dense + # transformer load. /info only advertises mxfp8 on sm100+, so this mirrors that gate. + import torch + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 9)) + cfg = types.SimpleNamespace(base_precision = "mxfp8", mixed_precision = "bf16", base_model = "x") + with pytest.raises(ValueError, match = "Blackwell"): + _resolve_base_precision(cfg, None, "cuda") + + +def test_resolve_base_precision_explicit_mxfp8_ok_on_blackwell(monkeypatch): + import torch + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 0)) + cfg = types.SimpleNamespace(base_precision = "mxfp8", mixed_precision = "bf16", base_model = "x") + assert _resolve_base_precision(cfg, None, "cuda") == "mxfp8" + + def test_mx_module_filter_skips_lora_and_proj_out(): # LoRA-owned modules (adapters stay high precision) and the output projection are # excluded, mirroring the fp8 filter's guards. From 477757eac790e89bd1b1ff99debe99d2a58e7412 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 10:42:52 +0000 Subject: [PATCH 18/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_dit_trainer.py | 1 - studio/backend/tests/test_diffusion_dit_trainer.py | 6 +++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index b6c5cd6983..325bf0d95e 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -430,7 +430,6 @@ def _resolve_base_precision(cfg, spec, device) -> str: if mode == "mxfp8" and device == "cuda": try: import torch - blackwell = torch.cuda.get_device_capability() >= (10, 0) except Exception: # noqa: BLE001 -- probe failure -> treat as unsupported, fail fast blackwell = False diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index aaad553107..0e8e8a0370 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -143,7 +143,11 @@ def test_family_train_infos_sdxl_supports_compile_without_precision_modes(monkey # ── mxfp8 base precision (DiT dense speed mode) ─────────────────────────────── -def _linear(in_features, out_features, bias = False): +def _linear( + in_features, + out_features, + bias = False, +): import torch.nn as nn return nn.Linear(in_features, out_features, bias = bias)