From 6384fea272b45276ac983f2e0091a8ff70e17816 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 10:41:02 +0000 Subject: [PATCH] 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.