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.

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.

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.