Gate DiT-training bf16 on native compute capability, not emulated is_bf16_supported

torch.cuda.is_bf16_supported() defaults to counting pre-Ampere bf16 EMULATION as
supported, so on a T4/V100/RTX 20xx the DiT-training bf16 gates all passed even
though the trainer requires native Ampere-or-newer bf16: /diffusion/info advertised
the DiT precision modes, /diffusion/start's preflight let the run through and freed
resident GPU models, then the trainer child hit the real unsupported bf16 path. The
inference device resolver already fixed this (issue #6658) by gating NVIDIA on
capability major >= 8; the training path never got it. Add a shared
native_bf16_supported() helper (NVIDIA cap major >= 8; ROCm keeps the trustworthy
is_bf16_supported()) and use it in the three DiT bf16 sites -- train_precision_modes,
bf16_unsupported_reason, and the trainer guard -- so a pre-Ampere card is offered
nf4 only and never advertises/evicts-then-fails. Tests now exercise the emulation
case (is_bf16_supported True but capability < 8).

studio/backend/core/training/diffusion_dit_trainer.py

@@ -52,6 +52,7 @@ from core.training.diffusion_train_common import (
     _restore_perf_flags,
     discover_image_caption_pairs,
     has_functional_torchao,
+    native_bf16_supported,
     PermutationBatchSampler,
     repo_is_prequantized,
     resolve_train_steps,
@@ -1237,8 +1238,10 @@ def run_dit_lora_training(
     # The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is
     # unsupported for real runs but keeps import/unit tests architecture-agnostic).
     # Fail fast on pre-Ampere CUDA (T4/V100/RTX 20xx): bf16 compute is required and the run
-    # would otherwise die deep in model load with an opaque dtype error.
-    if device == "cuda" and not torch.cuda.is_bf16_supported():
+    # would otherwise die deep in model load with an opaque dtype error. Gate on NATIVE bf16
+    # (capability major >= 8) -- is_bf16_supported() counts pre-Ampere emulation as supported,
+    # which is what this guard exists to reject; shared with the /info modes + start preflight.
+    if device == "cuda" and not native_bf16_supported():
         raise ValueError(
             "This trainer requires a bfloat16-capable GPU (Ampere or newer); "
             "this CUDA device does not support bf16."

studio/backend/core/training/diffusion_train_common.py

@@ -213,7 +213,7 @@ def train_precision_modes() -> tuple[list[str], str]:
     recommended = "nf4"
     try:
         import torch
-        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
+        if native_bf16_supported():
             modes.append("bf16")
             torchao_ok = has_functional_torchao()
             if torchao_ok:
@@ -292,6 +292,28 @@ _FAMILY_VRAM_NOTES = {
 _DIT_TRAIN_FAMILIES = frozenset({"flux.1", "qwen-image", "z-image", "krea-2"})
 
 
+def native_bf16_supported() -> bool:
+    """True only when the live CUDA GPU provides NATIVE bf16 compute, not pre-Ampere emulation.
+
+    ``torch.cuda.is_bf16_supported()`` defaults to counting EMULATED bf16, which every pre-Ampere
+    CUDA card (T4 / V100 / RTX 20xx) reports as supported even though the DiT trainer needs real
+    Ampere-or-newer bf16. Gate NVIDIA on compute capability major >= 8 instead -- the same #6658
+    fix the inference device resolver (``diffusion_device.py``) already uses; ROCm has no such
+    quirk, so ``is_bf16_supported()`` is trustworthy there. Never raises -- a probe failure or a
+    no-CUDA host returns False. Shared by the /info modes, the start preflight, and the trainer
+    guard so all three stay in sync."""
+    try:
+        import torch
+        if not torch.cuda.is_available():
+            return False
+        is_rocm = bool(getattr(getattr(torch, "version", None), "hip", None))
+        if is_rocm:
+            return bool(torch.cuda.is_bf16_supported())
+        return torch.cuda.get_device_capability()[0] >= 8
+    except Exception:  # noqa: BLE001 -- no torch / probe failure -> treat as unsupported
+        return False
+
+
 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
@@ -302,7 +324,7 @@ def bf16_unsupported_reason(resolved_family: str) -> Optional[str]:
         return None
     try:
         import torch
-        if torch.cuda.is_available() and not torch.cuda.is_bf16_supported():
+        if torch.cuda.is_available() and not native_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."

studio/backend/tests/test_diffusion_base_precision.py

@@ -142,13 +142,16 @@ def test_bf16_unsupported_reason(monkeypatch):
     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.
+    # A DiT family on a pre-Ampere CUDA GPU -> a clear reason. Pre-Ampere cards EMULATE bf16 and
+    # report is_bf16_supported() True, so the gate is native compute capability (major >= 8), not
+    # is_bf16_supported() -- otherwise the emulation case would slip through and evict-then-fail.
     monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda *a, **k: True)  # emulation reports True
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (7, 5))  # Turing
     assert "bfloat16" in (bf16_unsupported_reason("flux.1") or "")
 
-    # A bf16-capable GPU -> no reason.
-    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True)
+    # A NATIVE bf16-capable GPU (Ampere+) -> no reason.
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 0))
     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.
@@ -156,6 +159,24 @@ def test_bf16_unsupported_reason(monkeypatch):
     assert bf16_unsupported_reason("z-image") is None
 
 
+def test_native_bf16_supported_gates_on_capability(monkeypatch):
+    # Native bf16 is gated by compute capability (Ampere major >= 8), NOT is_bf16_supported(),
+    # which defaults to counting pre-Ampere emulation. A Turing card that emulates bf16 is
+    # correctly reported unsupported; Ampere+ is supported; a CPU-only host is always False.
+    import torch
+
+    from core.training.diffusion_train_common import native_bf16_supported
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda *a, **k: True)  # emulation reports True
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (7, 5))  # Turing
+    assert native_bf16_supported() is False
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 0))  # Ampere
+    assert native_bf16_supported() is True
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    assert native_bf16_supported() is False
+
+
 def test_training_precision_preflight_error(monkeypatch):
     # The start route calls this BEFORE evicting resident GPU workloads: it folds the bf16-GPU
     # requirement together with the explicit-int8 torchao requirement, so both fail fast instead
@@ -164,14 +185,17 @@ def test_training_precision_preflight_error(monkeypatch):
 
     from core.training.diffusion_train_common import training_precision_preflight_error
 
-    # Present a bf16-capable CUDA GPU so the int8 gate (not the bf16 gate) is what we exercise.
+    # Present a NATIVE bf16-capable CUDA GPU (Ampere+, cap major >= 8) so the int8 gate (not the
+    # bf16 gate) is what we exercise. bf16 is gated by capability, not is_bf16_supported() (which
+    # counts pre-Ampere emulation).
     monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda *a, **k: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 6))
 
-    # The bf16 gate takes precedence: a non-bf16 GPU rejects any DiT precision first.
-    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
+    # The bf16 gate takes precedence: a pre-Ampere GPU (emulated bf16) rejects any DiT precision.
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (7, 5))
     assert "bfloat16" in (training_precision_preflight_error("flux.1", "int8") or "")
-    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 6))
 
     # Explicit int8 on a DiT family with a NON-functional torchao -> a clear int8 reason
     # (its _int8_quantize_base has no fallback, so the child would otherwise raise post-eviction).

studio/backend/tests/test_diffusion_dit_trainer.py

@@ -302,3 +302,17 @@ def test_train_precision_modes_newer_blackwell_has_mxfp8(monkeypatch):
     _patch_capability(monkeypatch, (12, 0))
     modes, _ = train_precision_modes()
     assert "mxfp8" in modes
+
+
+def test_train_precision_modes_pre_ampere_is_nf4_only(monkeypatch):
+    # A pre-Ampere GPU EMULATES bf16 (is_bf16_supported() True) but has no native bf16 tensor
+    # cores; the DiT trainer requires native bf16, so /info must offer nf4 only. Otherwise it
+    # 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 *a, **k: True)  # emulation reports True
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (7, 5))  # Turing
+    modes, recommended = train_precision_modes()
+    assert modes == ["nf4"]
+    assert recommended == "nf4"