diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index dbbcda32d7..480351ba73 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -274,6 +274,27 @@ def bf16_unsupported_reason(resolved_family: str) -> Optional[str]: return None +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.""" + 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'." + ) + 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 @@ -291,10 +312,17 @@ def family_train_infos() -> list[dict[str, Any]]: # 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 _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 [] + # On a non-bf16 CUDA GPU the start route's preflight rejects EVERY DiT family (even nf4, + # since the DiT trainer requires bf16 unconditionally on CUDA), so advertise no precision + # for it -- otherwise /info offers an nf4 DiT option that always 400s. Otherwise drop any + # scheme this family's DiT corrupts (fp8 on Qwen-Image: activation outliers exceed fp8's + # range; the inference path denies the same set), so the UI never offers a mode + # normalized() would then reject. + dit_block = bf16_unsupported_reason(name) if is_dit else None + if not is_dit or dit_block: + fam_modes: list[str] = [] + else: + fam_modes = [m for m in dit_modes if not _family_denied(name, m)] infos.append( { "name": name, @@ -302,10 +330,10 @@ def family_train_infos() -> list[dict[str, Any]]: "default_base": repos[0], "base_repos": repos, "defaults": train_defaults(name), - "vram_note": _FAMILY_VRAM_NOTES.get(name, ""), + "vram_note": dit_block or _FAMILY_VRAM_NOTES.get(name, ""), "precision_modes": fam_modes, - "recommended_precision": dit_recommended if is_dit else "nf4", - "supports_compile": is_dit, + "recommended_precision": "nf4" if (not is_dit or dit_block) else dit_recommended, + "supports_compile": bool(is_dit and not dit_block), } ) return infos diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 275dfe98dc..b08efa5753 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1238,16 +1238,18 @@ async def start_diffusion_training( 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 + # Preflight the requested DiT precision BEFORE freeing GPU residents: the DiT trainer's own + # checks (a bf16-capable GPU is required; an explicit int8 needs a functional torchao) fire + # only in the child, 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) or a + # stub-torchao host never tears down resident models for a run that cannot start. + from core.training.diffusion_train_common import training_precision_preflight_error - _bf16_reason = bf16_unsupported_reason(normalized_cfg.resolved_family) - if _bf16_reason: - raise HTTPException(status_code = 400, detail = _bf16_reason) + _precision_reason = training_precision_preflight_error( + normalized_cfg.resolved_family, normalized_cfg.base_precision + ) + if _precision_reason: + raise HTTPException(status_code = 400, detail = _precision_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 diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index c6cf0ec33c..a769568c30 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -152,6 +152,62 @@ def test_bf16_unsupported_reason(monkeypatch): assert bf16_unsupported_reason("z-image") is None +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 + # of only surfacing in the trainer child after the GPU has already been freed. + import torch + + 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. + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True) + + # The bf16 gate takes precedence: a non-bf16 GPU rejects any DiT precision first. + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False) + assert "bfloat16" in (training_precision_preflight_error("flux.1", "int8") or "") + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True) + + # 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). + monkeypatch.setattr(common, "has_functional_torchao", lambda: False) + reason = training_precision_preflight_error("qwen-image", "int8") + assert reason is not None and "int8" in reason and "torchao" in reason + + # The same int8 request is fine once torchao is functional. + monkeypatch.setattr(common, "has_functional_torchao", lambda: True) + assert training_precision_preflight_error("qwen-image", "int8") is None + + # With a broken torchao, only EXPLICIT int8 is gated -- nf4/bf16/auto pass, and the int8 + # gate never applies to a non-DiT (SDXL) or unknown family. + monkeypatch.setattr(common, "has_functional_torchao", lambda: False) + 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", "int8") is None + assert training_precision_preflight_error("", "int8") 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 + # advertise a DiT precision option that always 400s: the modes empty, the reason surfaces in + # vram_note, compile is off, and the recommendation degrades to nf4. SDXL (non-DiT) is exempt. + from core.training.diffusion_train_common import _DIT_TRAIN_FAMILIES, family_train_infos + + monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: "no bfloat16 on this GPU") + + infos = {info["name"]: info for info in family_train_infos()} + dit_seen = False + for name, info in infos.items(): + if name in _DIT_TRAIN_FAMILIES: + dit_seen = True + assert info["precision_modes"] == [] + assert info["vram_note"] == "no bfloat16 on this GPU" + assert info["recommended_precision"] == "nf4" + assert info["supports_compile"] is False + assert dit_seen # the registry must still expose at least one DiT family to have covered it + + 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 diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 07e524958e..156803778a 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -606,17 +606,17 @@ def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypat 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. + # A DiT precision the host cannot run (no bf16 GPU, or explicit int8 without a functional + # torchao) must 400 BEFORE resident GPU workloads are freed: otherwise the host tears down the + # user's chat/Images model and the run then dies in the trainer child. The route imports + # training_precision_preflight_error 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: ( + "core.training.diffusion_train_common.training_precision_preflight_error", + lambda fam, prec: ( "This trainer requires a bfloat16-capable GPU (Ampere or newer)." if fam != "sdxl" else None