From 0f2f4334e201682f7ac2a18e1b0aaf74b4f416d4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 15:22:58 +0000 Subject: [PATCH] Reject dense DiT precisions on a CUDA-absent host before eviction; stabilize family-info tests The start-route preflight caught the bf16-GPU and int8-torchao requirements but not the dense precisions' CUDA requirement: on a GPU-less host bf16_unsupported_reason exempts CPU-only, so a bf16/fp8 (or int8-with-torchao) DiT request passed the preflight, evicted resident workloads, then raised only in the trainer child. Add the dense-mode CUDA gate mirroring _resolve_base_precision so the doomed run is rejected up front. Also pin bf16_unsupported_reason in the two positive-path family-info tests so they are deterministic across GPU types (a non-bf16 CUDA box would otherwise empty every DiT family's advertised modes). --- .../core/training/diffusion_train_common.py | 39 +++++++++++++------ .../tests/test_diffusion_base_precision.py | 19 +++++++++ 2 files changed, 46 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 480351ba73..95767a7af0 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -277,21 +277,36 @@ def bf16_unsupported_reason(resolved_family: str) -> Optional[str]: def training_precision_preflight_error(resolved_family: str, base_precision: str) -> Optional[str]: """Reason the requested DiT precision cannot run on this host, else None -- checked by the start route BEFORE evicting resident GPU workloads (the trainer's own checks fire only in the - child, after eviction). Two gates: the bf16-GPU requirement (bf16_unsupported_reason), and an - explicit int8 needing a FUNCTIONAL torchao (its _int8_quantize_base has no fallback, so - _resolve_base_precision would otherwise raise only after eviction). Never raises.""" + child, after eviction). Three gates, all mirroring _resolve_base_precision so a doomed run is + rejected before teardown: the bf16-GPU requirement (bf16_unsupported_reason); the dense + precisions (bf16/int8/fp8) requiring a CUDA GPU; and an explicit int8 needing a FUNCTIONAL + torchao (its _int8_quantize_base has no fallback). Never raises.""" reason = bf16_unsupported_reason(resolved_family) if reason: return reason - if ( - (resolved_family or "").strip().lower() in _DIT_TRAIN_FAMILIES - and (base_precision or "").strip().lower() == "int8" - and not has_functional_torchao() - ): - return ( - "base_precision='int8' needs a functional torchao install; this host's torchao is " - "missing or the non-functional Windows-ROCm stub. Use 'nf4', 'bf16', or 'auto'." - ) + fam = (resolved_family or "").strip().lower() + mode = (base_precision or "").strip().lower() + if fam in _DIT_TRAIN_FAMILIES and mode in ("bf16", "int8", "fp8"): + # The DiT trainer's dense precisions all require CUDA (_resolve_base_precision rejects + # bf16/int8/fp8 on device != "cuda"). bf16_unsupported_reason exempts a CPU-only host (the + # fp32 fallback for import/unit tests), so without this a dense request on a GPU-less host + # would pass the preflight, evict resident workloads, then raise only in the child. + try: + import torch + + has_cuda = torch.cuda.is_available() + except Exception: # noqa: BLE001 -- no torch / probe failure -> treat as no CUDA + has_cuda = False + if not has_cuda: + return ( + f"base_precision={mode!r} needs a CUDA GPU; this host has none. " + "Use base_precision='nf4' or 'auto'." + ) + if mode == "int8" and not has_functional_torchao(): + return ( + "base_precision='int8' needs a functional torchao install; this host's torchao is " + "missing or the non-functional Windows-ROCm stub. Use 'nf4', 'bf16', or 'auto'." + ) return None diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index a769568c30..66bebb28c6 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -100,6 +100,10 @@ def test_family_train_infos_drops_denied_fp8_for_qwen(monkeypatch): monkeypatch.setattr( common, "train_precision_modes", lambda: (["nf4", "bf16", "int8", "fp8", "auto"], "auto") ) + # family_train_infos reads the live GPU via bf16_unsupported_reason; pin it to "bf16 OK" so + # this positive-path assertion is deterministic across GPU types (a non-bf16 CUDA box would + # otherwise empty every DiT family's modes). The empty-on-non-bf16 path is covered separately. + monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: None) 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 @@ -187,6 +191,18 @@ def test_training_precision_preflight_error(monkeypatch): assert training_precision_preflight_error("sdxl", "int8") is None assert training_precision_preflight_error("", "int8") is None + # On a CUDA-ABSENT host, bf16_unsupported_reason exempts CPU-only, but the DiT trainer's dense + # precisions still require CUDA (mirroring _resolve_base_precision), so bf16/int8/fp8 for a DiT + # family are rejected UP FRONT rather than after eviction. nf4/auto (and SDXL) still pass. + monkeypatch.setattr(common, "has_functional_torchao", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + for dense in ("bf16", "int8", "fp8"): + reason = training_precision_preflight_error("flux.1", dense) + assert reason is not None and "CUDA" in reason + assert training_precision_preflight_error("flux.1", "nf4") is None + assert training_precision_preflight_error("flux.1", "auto") is None + assert training_precision_preflight_error("sdxl", "bf16") is None + def test_family_train_infos_empties_dit_modes_on_non_bf16(monkeypatch): # On a non-bf16 GPU the start route rejects EVERY DiT family (even nf4), so /info must not @@ -483,6 +499,9 @@ def test_family_train_infos_carries_precision_fields(monkeypatch): # Pin the machine probe so the DiT families carry a deterministic mode list, while SDXL # (no precision selector) stays empty regardless of the probe. monkeypatch.setattr(common, "train_precision_modes", lambda: (["nf4", "bf16"], "auto")) + # Also pin bf16_unsupported_reason (family_train_infos reads the live GPU through it): "bf16 OK" + # so this positive-path assertion is deterministic across GPU types, not just on CPU-only CI. + monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: None) infos = {i["name"]: i for i in common.family_train_infos()} flux = infos["flux.1"]