From cb9247e5376f5d58eadd48e6572d2b7e8277bbd7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 19:13:03 +0000 Subject: [PATCH] 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)"; };