diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index fdc7fa96b2..51f56af84e 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/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. diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 82f89b5ee2..9d1dfc3475 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/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: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index c579e1ca24..275dfe98dc 100644 --- a/studio/backend/routes/training.py +++ b/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. diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index e9e28018b6..88f83ce307 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/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 diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index a03803f1e6..07e524958e 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/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)