Gate DiT training precision: deny fp8 for Qwen, gate explicit int8 on torchao, gate advertised dense modes + route on bf16

- normalized() + family_train_infos() mirror the inference fp8 deny for
  Qwen-Image (activation outliers exceed fp8's range and corrupt the trained
  result); int8 stays allowed and the UI no longer advertises fp8 for it.
- _resolve_base_precision() gates an explicit int8 on a FUNCTIONAL torchao, the
  same gate auto and /info already apply, so a missing/stub torchao fails fast
  instead of silently loading dense with compile disabled.
- train_precision_modes() gates the dense modes (bf16/int8/fp8/auto) on
  torch.cuda.is_bf16_supported(), so a non-bf16 CUDA GPU (T4/V100/RTX 20xx) is
  offered only nf4 instead of a start that evicts resident models and then fails.
- start_diffusion_training preflights bf16 support for the DiT families BEFORE
  _free_gpu_for_diffusion_training(), so any DiT start (nf4 included, since the
  trainer requires bf16 unconditionally on CUDA) fails fast without eviction.

studio/backend/core/training/diffusion_dit_trainer.py

@@ -343,6 +343,18 @@ 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'."
             )
+        # int8 has no runtime fallback (_int8_quantize_base imports torchao unconditionally),
+        # so an explicit int8 against a missing torchao or the Windows-ROCm stub would leave
+        # the transformer dense with compile disabled as if it were int8 -- the memory saving
+        # silently gone and a likely OOM. The auto pick and /info already gate on a FUNCTIONAL
+        # torchao; apply the same gate to the explicit request so it fails fast with a clear
+        # message. fp8 keeps its own graceful fallback (_apply_fp8_training), so this is int8-only.
+        if mode == "int8" and not has_functional_torchao():
+            raise ValueError(
+                "base_precision='int8' needs a functional torchao install; this host's "
+                "torchao is missing or the non-functional Windows-ROCm stub. Use "
+                "base_precision='nf4', 'bf16', 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/core/training/diffusion_train_common.py

@@ -174,15 +174,18 @@ 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 CUDA; int8/fp8 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+). Used by the /info endpoint
+    always works; bf16/auto need a bf16-capable CUDA GPU (Ampere+); int8/fp8 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."""
     modes = ["nf4"]
     recommended = "nf4"
     try:
         import torch
-        if torch.cuda.is_available():
+        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
             modes.append("bf16")
             torchao_ok = has_functional_torchao()
             if torchao_ok:
@@ -244,12 +247,40 @@ _FAMILY_VRAM_NOTES = {
     "z-image": "6B model, QLoRA (nf4) by default (~12 GB+). bf16 only.",
 }
 
+# The flow-matching DiT families (run by diffusion_dit_trainer). They expose the
+# base_precision / compile levers and require bf16 compute on CUDA; SDXL is absent because
+# it uses its own mixed_precision path. Kept as a set so the UI gate, the bf16 preflight,
+# and any future dispatch stay in sync.
+_DIT_TRAIN_FAMILIES = frozenset({"flux.1", "qwen-image", "z-image"})
+
+
+def bf16_unsupported_reason(resolved_family: str) -> Optional[str]:
+    """Return a user-facing error string if ``resolved_family`` needs bf16 compute that the
+    live GPU cannot provide, else None. The DiT trainer requires a bf16-capable GPU (Ampere
+    or newer) and otherwise raises deep in model load; the start route uses this to fail fast
+    BEFORE evicting resident GPU workloads. CPU-only hosts (which fall back to fp32 for
+    import/unit tests) and SDXL (its own mixed_precision path) are exempt. Never raises."""
+    if (resolved_family or "").strip().lower() not in _DIT_TRAIN_FAMILIES:
+        return None
+    try:
+        import torch
+
+        if torch.cuda.is_available() and not torch.cuda.is_bf16_supported():
+            return (
+                "This trainer requires a bfloat16-capable GPU (Ampere or newer); this CUDA "
+                "device does not support bf16. Train the DiT families on a newer GPU."
+            )
+    except Exception:  # noqa: BLE001 -- torch probe failure must not block a start
+        return None
+    return None
+
 
 def family_train_infos() -> list[dict[str, Any]]:
     """Describe every trainable family for the Train UI: name, label, the default + allowed
     base repos, the recommended starting hyperparameters, and a VRAM/access note. Built from
     the family registry so it stays in sync with what the trainers actually support."""
     from core.inference.diffusion_families import detect_family
+    from core.inference.diffusion_transformer_quant import _family_denied
 
     dit_modes, dit_recommended = train_precision_modes()
     infos: list[dict[str, Any]] = []
@@ -260,7 +291,11 @@ def family_train_infos() -> list[dict[str, Any]]:
         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.
-        is_dit = name in ("flux.1", "qwen-image", "z-image")
+        is_dit = name in _DIT_TRAIN_FAMILIES
+        # Drop any advertised scheme this family's DiT cannot use (fp8 corrupts Qwen-Image:
+        # activation outliers exceed fp8's range; the inference path denies the same set), so
+        # the UI never offers a mode that normalized() would then reject.
+        fam_modes = [m for m in dit_modes if not _family_denied(name, m)] if is_dit else []
         infos.append(
             {
                 "name": name,
@@ -269,7 +304,7 @@ def family_train_infos() -> list[dict[str, Any]]:
                 "base_repos": repos,
                 "defaults": train_defaults(name),
                 "vram_note": _FAMILY_VRAM_NOTES.get(name, ""),
-                "precision_modes": dit_modes if is_dit else [],
+                "precision_modes": fam_modes,
                 "recommended_precision": dit_recommended if is_dit else "nf4",
                 "supports_compile": is_dit,
             }
@@ -381,6 +416,19 @@ class DiffusionLoraConfig:
                     f"base_precision={base_precision!r} trains in bf16 compute; set "
                     f"mixed_precision to bf16."
                 )
+            # Some DiT families are corrupted by fp8's activation range: outliers exceed even
+            # per-row fp8's dynamic range, so the frozen linears' float8 training compute
+            # learns against a garbage forward pass. The inference path already denies these
+            # schemes; mirror that deny here so the run fails fast instead of silently
+            # producing a broken adapter. int8 (per-token) is unaffected and stays allowed.
+            from core.inference.diffusion_transformer_quant import _family_denied
+
+            if _family_denied(resolved_family, base_precision):
+                raise ValueError(
+                    f"base_precision={base_precision!r} is not supported for "
+                    f"{resolved_family}: its activations exceed fp8's range and corrupt the "
+                    f"trained result. Use 'nf4', 'int8', 'bf16', or 'auto'."
+                )
         # A zero/negative gamma would zero out (or invert) the min-SNR weight and
         # silently train on a degenerate loss; None is the documented disable.
         if self.snr_gamma is not None and float(self.snr_gamma) <= 0:

studio/backend/routes/training.py

@@ -1234,10 +1234,21 @@ async def start_diffusion_training(
     from core.training.diffusion_lora_trainer import _config_from_dict
 
     try:
-        _config_from_dict(config).normalized()
+        normalized_cfg = _config_from_dict(config).normalized()
     except ValueError as e:
         raise HTTPException(status_code = 400, detail = str(e))
 
+    # Preflight bf16 support for the DiT families BEFORE freeing GPU residents: the DiT
+    # trainer requires a bf16-capable GPU (Ampere or newer) and otherwise raises deep in
+    # model load -- which, from here, would happen only AFTER _free_gpu_for_diffusion_training()
+    # already evicted the user's chat/Images model. Fail fast (400) so a pre-Ampere GPU
+    # (T4 / V100 / RTX 20xx) never tears down resident models for a run that cannot start.
+    from core.training.diffusion_train_common import bf16_unsupported_reason
+
+    _bf16_reason = bf16_unsupported_reason(normalized_cfg.resolved_family)
+    if _bf16_reason:
+        raise HTTPException(status_code = 400, detail = _bf16_reason)
+
     # Run the trainers' trust gate here too (both assert the same predicate before
     # from_pretrained), so an untrusted/typoed base 400s BEFORE freeing GPU residents
     # instead of tearing down the user's chat/Images model and failing in the child.

studio/backend/tests/test_diffusion_base_precision.py

@@ -35,6 +35,9 @@ _Z_PREQUANT = "unsloth/Z-Image-Turbo-unsloth-bnb-4bit"
 # An SDXL base whose name LOOKS prequant (bnb-4bit): SDXL ignores base_precision, so the
 # dense-mode gates must not fire for it even with a dense mode + fp16 compute.
 _SDXL_PREQUANT_NAME = "some/sdxl-model-bnb-4bit"
+# A dense Qwen-Image base: its DiT is corrupted by fp8 (activation outliers), so fp8 is
+# denied for training the same way the inference path denies it.
+_QWEN_DENSE = "Qwen/Qwen-Image"
 
 
 def _cfg(base_model = _FLUX_DENSE, **kw) -> DiffusionLoraConfig:
@@ -69,6 +72,84 @@ def test_base_precision_validation():
     assert _cfg(base_model = _Z_PREQUANT, base_precision = "auto").normalized().base_precision == "auto"
 
 
+def test_base_precision_denies_fp8_for_corrupted_family():
+    # fp8 corrupts the Qwen-Image DiT (activation outliers exceed fp8's range), so a dense
+    # Qwen base with base_precision="fp8" is refused up front -- mirroring the inference deny.
+    with pytest.raises(ValueError, match = "fp8"):
+        _cfg(base_model = _QWEN_DENSE, base_precision = "fp8", mixed_precision = "bf16").normalized()
+
+    # The deny is fp8-specific: int8 (per-token, unaffected) and the other dense modes stay
+    # allowed for the same Qwen base.
+    for mode in ("nf4", "bf16", "int8", "auto"):
+        norm = _cfg(base_model = _QWEN_DENSE, base_precision = mode, mixed_precision = "bf16").normalized()
+        assert norm.resolved_family == "qwen-image"
+        assert norm.base_precision == mode
+
+    # A family the deny does not cover (FLUX) still accepts fp8.
+    flux = _cfg(base_model = _FLUX_DENSE, base_precision = "fp8", mixed_precision = "bf16").normalized()
+    assert flux.resolved_family == "flux.1"
+    assert flux.base_precision == "fp8"
+
+
+def test_family_train_infos_drops_denied_fp8_for_qwen(monkeypatch):
+    # /info advertises the machine's DiT modes per family, but a family whose DiT the mode
+    # corrupts must not offer it: with fp8 in the machine list, Qwen-Image drops fp8 while
+    # FLUX keeps it, so the UI never surfaces a mode normalized() would reject.
+    monkeypatch.setattr(
+        common, "train_precision_modes", lambda: (["nf4", "bf16", "int8", "fp8", "auto"], "auto")
+    )
+    infos = {i["name"]: i for i in common.family_train_infos()}
+    assert "fp8" not in infos["qwen-image"]["precision_modes"]
+    assert "int8" in infos["qwen-image"]["precision_modes"]  # int8 is fine on Qwen
+    assert "fp8" in infos["flux.1"]["precision_modes"]
+
+
+def test_resolve_base_precision_explicit_int8_gates_on_torchao(monkeypatch):
+    # Explicit int8 has no runtime fallback, so a missing/stub torchao must fail fast here
+    # rather than load dense with compile disabled. Gate the explicit request the same way
+    # auto + /info already gate it.
+    spec = dit._SPECS["flux.1"]
+    cfg = _cfg(base_precision = "int8")
+
+    monkeypatch.setattr(dit, "has_functional_torchao", lambda: False)  # torchao absent / stub
+    with pytest.raises(ValueError, match = "torchao"):
+        dit._resolve_base_precision(cfg, spec, "cuda")
+
+    # With a functional torchao the explicit int8 passes straight through.
+    monkeypatch.setattr(dit, "has_functional_torchao", lambda: True)
+    assert dit._resolve_base_precision(cfg, spec, "cuda") == "int8"
+
+    # The gate is int8-specific: explicit bf16/fp8 pass through regardless of torchao (fp8 has
+    # its own graceful fallback; bf16 needs no torchao).
+    monkeypatch.setattr(dit, "has_functional_torchao", lambda: False)
+    assert dit._resolve_base_precision(_cfg(base_precision = "bf16"), spec, "cuda") == "bf16"
+    assert dit._resolve_base_precision(_cfg(base_precision = "fp8"), spec, "cuda") == "fp8"
+
+
+def test_bf16_unsupported_reason(monkeypatch):
+    # The route uses this to fail fast on a non-bf16 GPU BEFORE evicting resident workloads.
+    import torch
+
+    from core.training.diffusion_train_common import bf16_unsupported_reason
+
+    # SDXL (own mixed_precision path) and unknown families are always exempt.
+    assert bf16_unsupported_reason("sdxl") is None
+    assert bf16_unsupported_reason("") is None
+
+    # A DiT family on a CUDA GPU without bf16 -> a clear reason.
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
+    assert "bfloat16" in (bf16_unsupported_reason("flux.1") or "")
+
+    # A bf16-capable GPU -> no reason.
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True)
+    assert bf16_unsupported_reason("qwen-image") is None
+
+    # A CPU-only host (fp32 fallback for import/unit tests) -> no reason even for a DiT family.
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    assert bf16_unsupported_reason("z-image") is None
+
+
 def test_base_precision_gates_skip_sdxl():
     # SDXL ignores base_precision, so the dense-mode gates (prequant base / non-bf16 compute)
     # must not fire for it: a prequant-looking SDXL name with base_precision="bf16" does not
@@ -307,6 +388,7 @@ def test_train_precision_modes_gates_int8_fp8_on_torchao(monkeypatch):
     import torch
 
     monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True)
     monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (10, 0))
 
     # No functional torchao (absent or stub): bf16 + auto only, int8/fp8 dropped.
@@ -322,6 +404,22 @@ def test_train_precision_modes_gates_int8_fp8_on_torchao(monkeypatch):
     assert "int8" in modes2 and "fp8" in modes2
 
 
+def test_train_precision_modes_gates_dense_on_bf16_support(monkeypatch):
+    # The dense modes (bf16/int8/fp8/auto) all train in bf16 compute, which the DiT trainer
+    # requires. On a CUDA GPU that cannot do bf16 (T4/V100/RTX 20xx), /info must offer ONLY
+    # nf4 -- otherwise the UI advertises a start that evicts resident models and then fails the
+    # trainer's bf16 guard.
+    import torch
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (7, 5))  # Turing, no bf16
+    monkeypatch.setattr(common, "has_functional_torchao", lambda: True)
+    modes, recommended = train_precision_modes()
+    assert modes == ["nf4"]
+    assert recommended == "nf4"
+
+
 # ── family_train_infos precision fields ───────────────────────────────────────
 def test_family_train_infos_carries_precision_fields(monkeypatch):
     # Pin the machine probe so the DiT families carry a deterministic mode list, while SDXL

studio/backend/tests/test_diffusion_training.py

@@ -605,6 +605,38 @@ def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypat
     assert client._fake.started_with is None
 
 
+def test_route_start_refuses_non_bf16_gpu_without_freeing_gpu(client, monkeypatch):
+    # A DiT family on a GPU that cannot do bf16 must 400 BEFORE resident GPU workloads are
+    # freed: otherwise the pre-Ampere GPU tears down the user's chat/Images model and the run
+    # then dies deep in model load at the trainer's bf16 guard. The route imports
+    # bf16_unsupported_reason locally, so patch it on its home module.
+    import routes.training as tr
+
+    freed = []
+    monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1))
+    monkeypatch.setattr(
+        "core.training.diffusion_train_common.bf16_unsupported_reason",
+        lambda fam: (
+            "This trainer requires a bfloat16-capable GPU (Ampere or newer)."
+            if fam != "sdxl"
+            else None
+        ),
+    )
+    r = client.post(
+        "/api/train/diffusion/start",
+        json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"},
+    )
+    assert r.status_code == 400
+    assert "bfloat16" in r.json()["detail"]
+    assert freed == []
+    assert client._fake.started_with is None
+
+    # SDXL (its own mixed_precision path) is exempt: the same probe returns None, so an SDXL
+    # start proceeds normally past the preflight.
+    r2 = client.post("/api/train/diffusion/start", json = _BODY)
+    assert r2.status_code == 200, r2.text
+
+
 # ── metric history + perf/family fields (PR A platform) ──────────────────────
 def test_apply_event_records_metric_history_and_perf():
     svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)