Gate DiT training on functional torchao for explicit int8; hide always-400 DiT modes on non-bf16 GPUs
The start route preflight only rejected non-bf16 GPUs; an explicit int8 request on a host with a missing or stub torchao passed the preflight, evicted resident GPU workloads, then died in the trainer child (its int8 base quantizer has no fallback). Fold both gates into training_precision_preflight_error so int8-without-torchao fails fast before eviction. Also empty the advertised DiT precision_modes (and surface the reason in vram_note, drop compile) whenever the bf16 preflight would reject the family, so /info never offers an nf4 DiT option the route always 400s.
This commit is contained in:
Daniel Han
parent
cc6d7c96a6
commit
aa54a062ec
4 changed files with 109 additions and 23 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue