From acc604ffafeb0f919aec238664abfcf8b1f96cab Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 04:59:35 +0000 Subject: [PATCH] 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). --- .../core/training/diffusion_dit_trainer.py | 7 +++- .../core/training/diffusion_train_common.py | 26 +++++++++++- .../tests/test_diffusion_base_precision.py | 42 +++++++++++++++---- .../tests/test_diffusion_dit_trainer.py | 14 +++++++ 4 files changed, 76 insertions(+), 13 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 300ddd4183..d78ddf9d6b 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/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." diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 34e0ec4746..d62cbd1707 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/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." diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index ed6128c87b..61ddcad959 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/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). diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index c02342b021..8ba1940c89 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/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"