diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 08b17bede3..300ddd4183 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -238,8 +238,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 @@ -258,7 +258,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). @@ -319,6 +319,64 @@ 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 + # 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 + + +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 + 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.quantization import quantize_ + quantize_( + transformer, + _mxfp8_training_config(), + 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, @@ -360,7 +418,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'." @@ -377,6 +435,21 @@ def _resolve_base_precision(cfg, spec, device) -> str: "torchao is missing or the non-functional Windows-ROCm stub. Use " "base_precision='nf4', 'bf16', 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. @@ -1060,8 +1133,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( @@ -1077,11 +1151,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 @@ -1273,8 +1350,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" @@ -1303,12 +1380,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 a4a9540cc4..879b88860f 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 @@ -384,6 +386,16 @@ 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 @@ -465,7 +477,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) # Permutation-cycle index sampler (shared with the DiT trainer): each dataset image is # visited once per cycle before any repeat, so a short run does not leave part of a @@ -484,6 +496,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) @@ -567,17 +580,23 @@ 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 385d4d8be1..7bc64b6f0c 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -128,7 +128,7 @@ 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 @@ -175,13 +175,14 @@ def has_functional_torchao() -> bool: def train_precision_modes() -> tuple[list[str], str]: """(supported base_precision modes, recommended pick) for the current machine: nf4 - always works; bf16/auto need a bf16-capable CUDA GPU (Ampere+); int8/fp8 additionally + always works; bf16/auto need a bf16-capable CUDA GPU (Ampere+); int8/fp8/mxfp8 additionally need a FUNCTIONAL torchao (their explicit paths import torchao with no fallback, and the - Windows-ROCm stub only looks installed). fp8 also needs an fp8-capable GPU (sm89+). The - dense modes all train in bf16 compute, which the DiT trainer requires, so a non-bf16 CUDA - GPU (T4/V100/RTX 20xx) is offered only nf4 -- otherwise /info would advertise a start that - evicts resident models and then fails the trainer's bf16 guard. Used by the /info endpoint - so the UI can gate the precision selector. Never raises.""" + Windows-ROCm stub only looks installed). fp8 also needs an fp8-capable GPU (sm89+); mxfp8 + (block-scaled fp8 compute) needs the Blackwell tensor cores (sm100+) its cuBLAS kernels + target. The dense modes all train in bf16 compute, which the DiT trainer requires, so a + non-bf16 CUDA GPU (T4/V100/RTX 20xx) is offered only nf4 -- otherwise /info would advertise + a start that evicts resident models and then fails the trainer's bf16 guard. Used by the + /info endpoint so the UI can gate the precision selector. Never raises.""" modes = ["nf4"] recommended = "nf4" try: @@ -194,6 +195,8 @@ def train_precision_modes() -> tuple[list[str], str]: major, minor = torch.cuda.get_device_capability() if torchao_ok and (major, minor) >= (8, 9) and hasattr(torch, "float8_e4m3fn"): modes.append("fp8") + if torchao_ok and (major, minor) >= (10, 0): + modes.append("mxfp8") modes.append("auto") recommended = "auto" except Exception: # noqa: BLE001 -- no torch / probe failure -> nf4 only @@ -332,8 +335,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 _DIT_TRAIN_FAMILIES # On a non-bf16 CUDA GPU the start route's preflight rejects EVERY DiT family (even nf4, # since the DiT trainer requires bf16 unconditionally on CUDA), so advertise no precision @@ -356,7 +360,9 @@ def family_train_infos() -> list[dict[str, Any]]: "vram_note": dit_block or _FAMILY_VRAM_NOTES.get(name, ""), "precision_modes": fam_modes, "recommended_precision": "nf4" if (not is_dit or dit_block) else dit_recommended, - "supports_compile": bool(is_dit and not dit_block), + # compile is offered everywhere (SDXL regional U-Net + DiT), except a DiT family + # the GPU can't train in bf16 (dit_block), where training is refused outright. + "supports_compile": bool(not dit_block), # Krea trains on Raw but previews adapters on Turbo; None elsewhere. "deploy_base": fam.deploy_base_repo, } @@ -455,13 +461,13 @@ 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") - # base_precision is a DiT-only lever (nf4/bf16/int8/fp8/auto for the transformer - # load); 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 mode-name validity check above still runs for every family. - if resolved_family != "sdxl" and 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") + # 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 + # mode-name validity check above still runs for every family. + if resolved_family != "sdxl" and 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 d6d5241072..f8120c42ee 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -740,11 +740,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_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index 66bebb28c6..8138e428a6 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -512,7 +512,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 ──────────────────────────────────────── diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index c40a0575a3..0e8e8a0370 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -9,6 +9,9 @@ exercised by the live GPU smokes, not here.""" from __future__ import annotations +import sys +import types + import pytest from core.training.diffusion_dit_trainer import ( @@ -17,15 +20,20 @@ from core.training.diffusion_dit_trainer import ( _QWEN_TARGETS, _SPECS, _ZIMAGE_TARGETS, + _apply_mxfp8_training, _assert_gated_access, + _mx_module_filter, _repo_is_prequantized, + _resolve_base_precision, _select_lora_targets, + _should_compile, run_dit_lora_training, ) from core.training.diffusion_train_common import ( DEFAULT_LORA_TARGETS, DiffusionLoraConfig, family_train_infos, + train_precision_modes, ) @@ -117,3 +125,176 @@ 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, + bias = False, +): + import torch.nn as nn + return nn.Linear(in_features, out_features, bias = bias) + + +def test_mx_module_filter_accepts_dense_block_linear(): + # 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. + 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 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 + warnings = [e for e in events if e["type"] == "warning"] + assert len(warnings) == 1 + 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+). 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): + # 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 4c7ed64213..f2971451e9 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -115,6 +115,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 e49f87ce53..729c975d82 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -393,6 +393,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_config_from_dict_epoch_mode_drops_max_steps_sentinel(): # The generic Studio epoch-mode payload sends max_steps: 0 as the "use epochs" sentinel. # The max_steps -> train_steps alias would copy that 0 and normalized() would reject diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 610e5ec935..94fc32507f 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 bf6e1b7abf..404ce164ec 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -94,7 +94,7 @@ const CUSTOM_BASE = "__custom__"; const UPLOAD_DATASET = "__upload__"; // The dense DiT base precisions: they load a dense (bf16) base and quantise/cast it, so the // backend rejects them for an already-quantised bnb-4bit repo. "nf4"/"auto" stay valid. -const DENSE_PRECISIONS = new Set(["bf16", "int8", "fp8"]); +const DENSE_PRECISIONS = new Set(["bf16", "int8", "fp8", "mxfp8"]); // Mirror the backend's repo_is_prequantized heuristic: a repo whose name marks a // bitsandbytes 4-bit build already ships a quantised transformer and cannot serve the dense // base precisions. Kept in sync with diffusion_train_common.repo_is_prequantized. @@ -197,17 +197,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(""); @@ -250,7 +255,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">( @@ -816,11 +821,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)"; }; @@ -883,63 +891,40 @@ 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. - {basePrequantized && ( - <> - {" "} - This base is already 4-bit quantised, so only nf4/auto apply; pick a dense - (bf16) base repo for the other modes. - - )} -

-
- {supportsCompile && ( -
- - { + precisionDirty.current = true; + setBasePrecision(e.target.value as typeof basePrecision); + }} + className={selectClass} + aria-label="Base precision" + > + {precisionModes.map((m) => ( + - - - -

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

-
- )} - + {precisionLabel(m)} + + ))} + +

+ 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. + {basePrequantized && ( + <> + {" "} + This base is already 4-bit quantised, so only nf4/auto apply; pick a dense + (bf16) base repo for the other modes. + + )} +

+ ) : (
@@ -958,6 +943,27 @@ export function DiffusionTrainPanel({

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

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

+
+ )} );