Merge remote-tracking branch 'origin/diffusion-train-perf2' into fold-integration
# Conflicts: # studio/backend/core/training/diffusion_dit_trainer.py # studio/backend/core/training/diffusion_train_common.py
This commit is contained in:
Daniel Han
commit
e1dd2dda6b
10 changed files with 457 additions and 106 deletions
|
|
@ -238,8 +238,8 @@ def _load_dit_transformer(transformer_cls, cfg, device, base_precision):
|
|||
|
||||
- nf4: a prequant (bnb-4bit) repo carries its quantization config and loads 4-bit
|
||||
as-is; a dense base is quantized to nf4 on the fly. The memory floor.
|
||||
- bf16 / fp8: the dense transformer (fp8 converts its frozen linears to float8
|
||||
training compute AFTER the LoRA attaches; storage stays bf16).
|
||||
- bf16 / fp8 / mxfp8: the dense transformer (fp8/mxfp8 convert its frozen linears to
|
||||
float8 training compute AFTER the LoRA attaches; storage stays bf16).
|
||||
- int8: the dense transformer quantized in place to torchao weight-only int8 (the
|
||||
PEFT-attachable scheme), roughly halving the bf16 weight footprint."""
|
||||
import torch
|
||||
|
|
@ -258,7 +258,7 @@ def _load_dit_transformer(transformer_cls, cfg, device, base_precision):
|
|||
transformer = transformer.to(device)
|
||||
return transformer
|
||||
|
||||
# Dense load for bf16 / fp8 / int8. int8 quantizes AFTER the LoRA attaches (see
|
||||
# Dense load for bf16 / fp8 / mxfp8 / int8. int8 quantizes AFTER the LoRA attaches (see
|
||||
# _int8_quantize_base): quantizing first makes peft dispatch its TorchaoLoraLinear
|
||||
# wrapper, whose peft-0.18 constructor is incompatible with the torchao-0.16 config API
|
||||
# (missing get_apply_tensor_subclass).
|
||||
|
|
@ -319,6 +319,64 @@ def _apply_fp8_training(transformer, on_event) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _mx_module_filter(mod, fqn: str) -> bool:
|
||||
"""Which frozen linears get mxfp8 training compute: skip anything LoRA-owned (the
|
||||
adapters must stay high precision), the output projection (same guard as fp8), and
|
||||
shapes the 32-wide MX block scaling cannot tile (dims not divisible by 32)."""
|
||||
import torch.nn as nn
|
||||
|
||||
if not isinstance(mod, nn.Linear):
|
||||
return False
|
||||
if "lora_" in fqn:
|
||||
return False
|
||||
if fqn.endswith("proj_out") or ".proj_out." in fqn:
|
||||
return False
|
||||
# Skip biased linears: the torchao 0.17 MX training path swaps the weight for a wrapper tensor
|
||||
# whose linear override computes input @ weight_t and drops the bias entirely, so an mxfp8'd
|
||||
# FROZEN base linear would silently lose its bias and change the output the LoRA regresses
|
||||
# against (verified on Blackwell: the bias term is fully dropped). Keep biased linears in bf16.
|
||||
if getattr(mod, "bias", None) is not None:
|
||||
return False
|
||||
return mod.in_features % 32 == 0 and mod.out_features % 32 == 0
|
||||
|
||||
|
||||
def _mxfp8_training_config():
|
||||
"""The torchao MX training config across the prototype API's revisions: torchao 0.16
|
||||
ships ``MXLinearConfig`` in ``prototype.mx_formats``; 0.17 removed it in favour of the
|
||||
``MXFP8TrainingOpConfig`` recipe API shared with MoE training. Both feed ``quantize_``.
|
||||
Raises ImportError when neither API exists (mxfp8 then falls back to bf16)."""
|
||||
try:
|
||||
from torchao.prototype.mx_formats import MXLinearConfig
|
||||
return MXLinearConfig.from_recipe_name("mxfp8_cublas")
|
||||
except ImportError:
|
||||
from torchao.prototype.moe_training.config import (
|
||||
MXFP8TrainingOpConfig,
|
||||
MXFP8TrainingRecipe,
|
||||
)
|
||||
return MXFP8TrainingOpConfig.from_recipe(MXFP8TrainingRecipe.MXFP8_RCEIL)
|
||||
|
||||
|
||||
def _apply_mxfp8_training(transformer, on_event) -> bool:
|
||||
"""Swap the frozen base linears to torchao MX float8 training compute (mxfp8, the
|
||||
Blackwell-native block-scaled format; the swap is in place and the weights stay bf16
|
||||
in memory, so like fp8 this is a speed mode, not a memory mode). Applied AFTER
|
||||
add_adapter so the filter can exclude the LoRA modules. Only competitive under
|
||||
torch.compile and only ahead of compiled bf16 at large token counts (high resolution
|
||||
or batch), which is why it stays an explicit opt-in rather than an "auto" pick.
|
||||
Never fatal: on any failure the run continues in bf16 with a warning."""
|
||||
try:
|
||||
from torchao.quantization import quantize_
|
||||
quantize_(
|
||||
transformer,
|
||||
_mxfp8_training_config(),
|
||||
filter_fn = _mx_module_filter,
|
||||
)
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 -- mxfp8 is an optimisation, never fatal
|
||||
_emit(on_event, "warning", message = f"mxfp8 training unavailable, using bf16 compute: {exc}")
|
||||
return False
|
||||
|
||||
|
||||
def _pick_auto_precision(
|
||||
prequant,
|
||||
device,
|
||||
|
|
@ -360,7 +418,7 @@ def _resolve_base_precision(cfg, spec, device) -> str:
|
|||
transformer onto the CPU."""
|
||||
mode = (cfg.base_precision or "nf4").strip().lower()
|
||||
if mode != "auto":
|
||||
if mode in ("bf16", "int8", "fp8") and device != "cuda":
|
||||
if mode in ("bf16", "int8", "fp8", "mxfp8") and device != "cuda":
|
||||
raise ValueError(
|
||||
f"base_precision={mode!r} needs a CUDA GPU; this host has none. "
|
||||
f"Use base_precision='nf4' or 'auto'."
|
||||
|
|
@ -377,6 +435,21 @@ def _resolve_base_precision(cfg, spec, device) -> str:
|
|||
"torchao is missing or the non-functional Windows-ROCm stub. Use "
|
||||
"base_precision='nf4', 'bf16', or 'auto'."
|
||||
)
|
||||
# mxfp8 needs Blackwell (sm100+): its MX GEMM has no kernel below sm100 and raises at the
|
||||
# first training step, AFTER a full dense-transformer load. /info only advertises mxfp8 on
|
||||
# sm100+ (train_precision_modes), so re-check it here to fail fast for a stale or direct
|
||||
# client on an older CUDA GPU instead of crashing mid-run.
|
||||
if mode == "mxfp8" and device == "cuda":
|
||||
try:
|
||||
import torch
|
||||
blackwell = torch.cuda.get_device_capability() >= (10, 0)
|
||||
except Exception: # noqa: BLE001 -- probe failure -> treat as unsupported, fail fast
|
||||
blackwell = False
|
||||
if not blackwell:
|
||||
raise ValueError(
|
||||
"base_precision='mxfp8' needs a Blackwell (sm100+) GPU; this GPU is older. "
|
||||
"Use base_precision='bf16', 'int8', 'nf4', 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.
|
||||
|
|
@ -1060,8 +1133,9 @@ def _should_compile(
|
|||
return True
|
||||
# auto: regional compile is the whole point of the dense modes (measured 2.6x on
|
||||
# Z-Image bf16) but fragile over bitsandbytes 4-bit modules (graph breaks in the
|
||||
# dequant path), so it stays off for QLoRA. fp8 is only competitive compiled.
|
||||
return base_precision in ("bf16", "fp8")
|
||||
# dequant path), so it stays off for QLoRA. fp8/mxfp8 are only competitive compiled
|
||||
# (eager, their per-matmul dynamic casts run 4-5x slower than bf16).
|
||||
return base_precision in ("bf16", "fp8", "mxfp8")
|
||||
|
||||
|
||||
def _maybe_compile_transformer(
|
||||
|
|
@ -1077,11 +1151,14 @@ def _maybe_compile_transformer(
|
|||
event, and dynamo's suppress_errors keeps a frame that fails to COMPILE at the first
|
||||
step running eager instead of raising mid-run."""
|
||||
if not _should_compile(cfg, base_is_bnb, device, base_precision):
|
||||
if base_precision == "fp8":
|
||||
if base_precision in ("fp8", "mxfp8"):
|
||||
_emit(
|
||||
on_event,
|
||||
"warning",
|
||||
message = "fp8 training without torch.compile is slow; enable compile for the speedup.",
|
||||
message = (
|
||||
f"{base_precision} training without torch.compile is slow; "
|
||||
f"enable compile for the speedup."
|
||||
),
|
||||
)
|
||||
return False
|
||||
import torch
|
||||
|
|
@ -1273,8 +1350,8 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
variant_rng = random.Random(cfg.seed + 1)
|
||||
|
||||
# Phase 3: only now load the transformer, in the resolved base precision (nf4 QLoRA by
|
||||
# default; bf16 / int8 / fp8 are the dense speed modes; "auto" picks from free VRAM
|
||||
# measured before the load).
|
||||
# default; bf16 / int8 / fp8 / mxfp8 are the dense speed modes; "auto" picks from
|
||||
# free VRAM measured before the load).
|
||||
base_precision = _resolve_base_precision(cfg, spec, device)
|
||||
transformer = spec.load_transformer(cfg, device, weight_dtype, base_precision)
|
||||
base_is_bnb = base_precision == "nf4"
|
||||
|
|
@ -1303,12 +1380,14 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
cast_training_params(transformer, dtype = torch.float32)
|
||||
lora_params = [p for p in transformer.parameters() if p.requires_grad]
|
||||
|
||||
# int8 / fp8 convert the frozen base linears AFTER the LoRA attaches, so the adapter
|
||||
# modules are excluded and stay high precision.
|
||||
# int8 / fp8 / mxfp8 convert the frozen base linears AFTER the LoRA attaches, so the
|
||||
# adapter modules are excluded and stay high precision.
|
||||
if base_precision == "int8":
|
||||
_int8_quantize_base(transformer)
|
||||
if base_precision == "fp8" and not _apply_fp8_training(transformer, on_event):
|
||||
base_precision = "bf16"
|
||||
if base_precision == "mxfp8" and not _apply_mxfp8_training(transformer, on_event):
|
||||
base_precision = "bf16"
|
||||
|
||||
compiled = _maybe_compile_transformer(
|
||||
transformer, cfg, base_is_bnb, device, on_event, base_precision
|
||||
|
|
|
|||
|
|
@ -26,7 +26,9 @@ latents are likewise precomputed into a small CPU cache (``cache_latents``) and
|
|||
freed. The cache stores the posterior's affine pair (mean/std, scale folded in), so every
|
||||
step still draws a fresh VAE sample -- distribution-identical to encoding in the loop,
|
||||
without keeping the VAE resident or paying a per-step encode. TF32 matmuls + cudnn
|
||||
autotuning are enabled for the run under ``cfg.enable_tf32``.
|
||||
autotuning are enabled for the run under ``cfg.enable_tf32``, and the U-Net's repeated
|
||||
transformer blocks are regionally torch.compiled (``cfg.compile_transformer``, never
|
||||
fatal -- any failure falls back to eager with a warning event).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -384,6 +386,16 @@ def run_diffusion_lora_training(
|
|||
if weight_dtype != torch.float32:
|
||||
cast_training_params(unet, dtype = torch.float32)
|
||||
|
||||
# Regionally torch.compile the U-Net's repeated BasicTransformerBlocks through the
|
||||
# DiT trainer's never-fatal wrapper (a wrap/compile failure falls back to eager
|
||||
# with a warning event). The U-Net is a dense bf16 base here, the combination that
|
||||
# wrapper compiles under "auto".
|
||||
from core.training.diffusion_dit_trainer import _maybe_compile_transformer
|
||||
|
||||
compiled = _maybe_compile_transformer(
|
||||
unet, cfg, False, device, on_event, base_precision = "bf16"
|
||||
)
|
||||
|
||||
lora_params = [p for p in unet.parameters() if p.requires_grad]
|
||||
optimizer = _make_lora_optimizer(lora_params, cfg.learning_rate)
|
||||
# The scheduler advances once per optimizer update: lr_sched.step() runs a single
|
||||
|
|
@ -465,7 +477,7 @@ def run_diffusion_lora_training(
|
|||
# seed-deterministic sequence whether or not the cache is enabled.
|
||||
variant_rng = random.Random(cfg.seed + 1)
|
||||
|
||||
_emit(on_event, "model_load_completed")
|
||||
_emit(on_event, "model_load_completed", compiled = compiled)
|
||||
|
||||
# Permutation-cycle index sampler (shared with the DiT trainer): each dataset image is
|
||||
# visited once per cycle before any repeat, so a short run does not leave part of a
|
||||
|
|
@ -484,6 +496,7 @@ def run_diffusion_lora_training(
|
|||
running_loss = 0.0
|
||||
peak_gb = 0.0
|
||||
t_start = time.time()
|
||||
t_steady = None
|
||||
done = 0
|
||||
for opt_step in range(cfg.train_steps):
|
||||
optimizer.zero_grad(set_to_none = True)
|
||||
|
|
@ -567,17 +580,23 @@ def run_diffusion_lora_training(
|
|||
|
||||
running_loss += step_loss
|
||||
done = opt_step + 1
|
||||
now = time.time()
|
||||
if done == 1:
|
||||
# Step 1 pays the one-time costs (cudnn autotune, torch.compile warmup), so
|
||||
# the reported rate starts after it and reflects the steady state (the DiT
|
||||
# trainer does the same).
|
||||
t_steady = now
|
||||
if done % cfg.log_every == 0 or done == cfg.train_steps:
|
||||
# ``learning_rate`` (not ``lr``) is the field the Studio training pump reads, so
|
||||
# these progress events are directly consumable by the existing training
|
||||
# status/SSE machinery when the diffusion trainer is wired into the worker.
|
||||
if device == "cuda":
|
||||
peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2)
|
||||
samples_per_second = round(
|
||||
(done * cfg.train_batch_size * cfg.gradient_accumulation_steps)
|
||||
/ max(time.time() - t_start, 1e-6),
|
||||
3,
|
||||
)
|
||||
per_step = cfg.train_batch_size * cfg.gradient_accumulation_steps
|
||||
if t_steady is not None and done > 1:
|
||||
samples_per_second = round((done - 1) * per_step / max(now - t_steady, 1e-6), 3)
|
||||
else:
|
||||
samples_per_second = round(done * per_step / max(now - t_start, 1e-6), 3)
|
||||
_emit(
|
||||
on_event,
|
||||
"progress",
|
||||
|
|
|
|||
|
|
@ -128,7 +128,7 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None
|
|||
def repo_is_prequantized(base_model: str) -> bool:
|
||||
"""Heuristic: a repo whose name marks a bitsandbytes 4-bit build already ships a
|
||||
quantized transformer, so it loads as-is for nf4 and cannot serve the dense
|
||||
(bf16/int8/fp8) base precisions."""
|
||||
(bf16/int8/fp8/mxfp8) base precisions."""
|
||||
name = str(base_model or "").lower()
|
||||
return "bnb-4bit" in name or "-4bit" in name or "int4" in name or "nf4" in name
|
||||
|
||||
|
|
@ -175,13 +175,14 @@ 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 a bf16-capable CUDA GPU (Ampere+); int8/fp8 additionally
|
||||
always works; bf16/auto need a bf16-capable CUDA GPU (Ampere+); int8/fp8/mxfp8 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."""
|
||||
Windows-ROCm stub only looks installed). fp8 also needs an fp8-capable GPU (sm89+); mxfp8
|
||||
(block-scaled fp8 compute) needs the Blackwell tensor cores (sm100+) its cuBLAS kernels
|
||||
target. 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:
|
||||
|
|
@ -194,6 +195,8 @@ def train_precision_modes() -> tuple[list[str], str]:
|
|||
major, minor = torch.cuda.get_device_capability()
|
||||
if torchao_ok and (major, minor) >= (8, 9) and hasattr(torch, "float8_e4m3fn"):
|
||||
modes.append("fp8")
|
||||
if torchao_ok and (major, minor) >= (10, 0):
|
||||
modes.append("mxfp8")
|
||||
modes.append("auto")
|
||||
recommended = "auto"
|
||||
except Exception: # noqa: BLE001 -- no torch / probe failure -> nf4 only
|
||||
|
|
@ -332,8 +335,9 @@ def family_train_infos() -> list[dict[str, Any]]:
|
|||
if fam is None:
|
||||
continue
|
||||
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.
|
||||
# base_precision applies to the DiT trainer only; SDXL keeps its mixed_precision
|
||||
# lever, so the UI hides the precision selector for it. compile applies everywhere:
|
||||
# the SDXL trainer regionally compiles the U-Net's transformer blocks too.
|
||||
is_dit = name in _DIT_TRAIN_FAMILIES
|
||||
# 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
|
||||
|
|
@ -356,7 +360,9 @@ def family_train_infos() -> list[dict[str, Any]]:
|
|||
"vram_note": dit_block or _FAMILY_VRAM_NOTES.get(name, ""),
|
||||
"precision_modes": fam_modes,
|
||||
"recommended_precision": "nf4" if (not is_dit or dit_block) else dit_recommended,
|
||||
"supports_compile": bool(is_dit and not dit_block),
|
||||
# compile is offered everywhere (SDXL regional U-Net + DiT), except a DiT family
|
||||
# the GPU can't train in bf16 (dit_block), where training is refused outright.
|
||||
"supports_compile": bool(not dit_block),
|
||||
# Krea trains on Raw but previews adapters on Turbo; None elsewhere.
|
||||
"deploy_base": fam.deploy_base_repo,
|
||||
}
|
||||
|
|
@ -455,13 +461,13 @@ class DiffusionLoraConfig:
|
|||
if compile_transformer not in ("off", "on", "auto"):
|
||||
raise ValueError("compile_transformer must be one of off / on / auto")
|
||||
base_precision = str(self.base_precision or "nf4").strip().lower()
|
||||
if base_precision not in ("nf4", "bf16", "int8", "fp8", "auto"):
|
||||
raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / auto")
|
||||
# base_precision is a DiT-only lever (nf4/bf16/int8/fp8/auto for the transformer
|
||||
# load); SDXL uses its own mixed_precision path and ignores base_precision entirely,
|
||||
# so the dense-mode gates (prequant base / non-bf16 compute) apply only to the DiT
|
||||
# families. The mode-name validity check above still runs for every family.
|
||||
if resolved_family != "sdxl" and base_precision in ("bf16", "int8", "fp8"):
|
||||
if base_precision not in ("nf4", "bf16", "int8", "fp8", "mxfp8", "auto"):
|
||||
raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto")
|
||||
# base_precision is a DiT-only lever (the transformer load precision); SDXL uses its
|
||||
# own mixed_precision path and ignores base_precision entirely, so the dense-mode
|
||||
# gates (prequant base / non-bf16 compute) apply only to the DiT families. The
|
||||
# mode-name validity check above still runs for every family.
|
||||
if resolved_family != "sdxl" and base_precision in ("bf16", "int8", "fp8", "mxfp8"):
|
||||
if repo_is_prequantized(self.base_model):
|
||||
raise ValueError(
|
||||
f"base_precision={base_precision!r} needs a dense base repo, but "
|
||||
|
|
|
|||
|
|
@ -740,11 +740,12 @@ class DiffusionTrainingStartRequest(BaseModel):
|
|||
enable_tf32: bool = Field(
|
||||
True, description = "TF32 matmuls + cudnn autotuning (near-lossless speedup)"
|
||||
)
|
||||
base_precision: Literal["nf4", "bf16", "int8", "fp8", "auto"] = Field(
|
||||
base_precision: Literal["nf4", "bf16", "int8", "fp8", "mxfp8", "auto"] = Field(
|
||||
"nf4",
|
||||
description = (
|
||||
"DiT base transformer precision: nf4 QLoRA (memory floor, default), bf16 dense, "
|
||||
"int8 torchao weight-only, fp8 float8 training compute (Ada/Hopper/Blackwell), "
|
||||
"mxfp8 block-scaled float8 compute (Blackwell, best at high resolution/batch), "
|
||||
"or auto (pick by free VRAM + GPU class). Dense modes need a non-prequant base."
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -512,7 +512,9 @@ def test_family_train_infos_carries_precision_fields(monkeypatch):
|
|||
sdxl = infos["sdxl"]
|
||||
assert sdxl["precision_modes"] == []
|
||||
assert sdxl["recommended_precision"] == "nf4"
|
||||
assert sdxl["supports_compile"] is False
|
||||
# The SDXL trainer regionally compiles its U-Net blocks too, so compile is advertised
|
||||
# for every family; only the precision selector stays DiT-only.
|
||||
assert sdxl["supports_compile"] is True
|
||||
|
||||
|
||||
# ── request model base_precision field ────────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -9,6 +9,9 @@ exercised by the live GPU smokes, not here."""
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
from core.training.diffusion_dit_trainer import (
|
||||
|
|
@ -17,15 +20,20 @@ from core.training.diffusion_dit_trainer import (
|
|||
_QWEN_TARGETS,
|
||||
_SPECS,
|
||||
_ZIMAGE_TARGETS,
|
||||
_apply_mxfp8_training,
|
||||
_assert_gated_access,
|
||||
_mx_module_filter,
|
||||
_repo_is_prequantized,
|
||||
_resolve_base_precision,
|
||||
_select_lora_targets,
|
||||
_should_compile,
|
||||
run_dit_lora_training,
|
||||
)
|
||||
from core.training.diffusion_train_common import (
|
||||
DEFAULT_LORA_TARGETS,
|
||||
DiffusionLoraConfig,
|
||||
family_train_infos,
|
||||
train_precision_modes,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -117,3 +125,176 @@ def test_family_train_infos_lists_dit_families():
|
|||
assert "gated" in infos["flux.1"]["vram_note"].lower()
|
||||
# Z-Image defaults to the prequant nf4 repo for QLoRA.
|
||||
assert "4bit" in infos["z-image"]["default_base"].lower()
|
||||
|
||||
|
||||
def test_family_train_infos_sdxl_supports_compile_without_precision_modes(monkeypatch):
|
||||
# Regional compile now applies to every family (the SDXL trainer compiles its U-Net
|
||||
# blocks too), but base_precision stays DiT-only, so SDXL advertises no precision modes
|
||||
# while a DiT family (z-image) keeps its own. Pin the precision list so the assertion
|
||||
# holds regardless of the test host's GPU capability.
|
||||
import core.training.diffusion_train_common as dtc
|
||||
|
||||
monkeypatch.setattr(dtc, "train_precision_modes", lambda: (["nf4", "bf16", "auto"], "auto"))
|
||||
infos = {i["name"]: i for i in family_train_infos()}
|
||||
assert infos["sdxl"]["supports_compile"] is True
|
||||
assert infos["sdxl"]["precision_modes"] == []
|
||||
assert infos["z-image"]["supports_compile"] is True
|
||||
assert infos["z-image"]["precision_modes"] == ["nf4", "bf16", "auto"]
|
||||
|
||||
|
||||
# ── mxfp8 base precision (DiT dense speed mode) ───────────────────────────────
|
||||
def _linear(
|
||||
in_features,
|
||||
out_features,
|
||||
bias = False,
|
||||
):
|
||||
import torch.nn as nn
|
||||
return nn.Linear(in_features, out_features, bias = bias)
|
||||
|
||||
|
||||
def test_mx_module_filter_accepts_dense_block_linear():
|
||||
# A bias-free 3072x3072 attention/FFN linear at a normal block fqn is a valid mxfp8 target.
|
||||
assert _mx_module_filter(_linear(3072, 3072), "blocks.0.ff.up") is True
|
||||
|
||||
|
||||
def test_mx_module_filter_skips_biased_linear():
|
||||
# The torchao 0.17 MX training path drops the bias term (its linear override computes
|
||||
# input @ weight_t only), so an mxfp8'd biased FROZEN linear would silently lose its bias and
|
||||
# corrupt the base output the LoRA regresses against. Biased linears must stay bf16.
|
||||
assert _mx_module_filter(_linear(3072, 3072, bias = True), "blocks.0.ff.up") is False
|
||||
|
||||
|
||||
def test_resolve_base_precision_explicit_mxfp8_requires_blackwell(monkeypatch):
|
||||
# An explicit mxfp8 request on a non-Blackwell CUDA GPU must fail fast: its MX GEMM has no
|
||||
# kernel below sm100 and would otherwise crash at the first training step, after a full dense
|
||||
# transformer load. /info only advertises mxfp8 on sm100+, so this mirrors that gate.
|
||||
import torch
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 9))
|
||||
cfg = types.SimpleNamespace(base_precision = "mxfp8", mixed_precision = "bf16", base_model = "x")
|
||||
with pytest.raises(ValueError, match = "Blackwell"):
|
||||
_resolve_base_precision(cfg, None, "cuda")
|
||||
|
||||
|
||||
def test_resolve_base_precision_explicit_mxfp8_ok_on_blackwell(monkeypatch):
|
||||
import torch
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 0))
|
||||
cfg = types.SimpleNamespace(base_precision = "mxfp8", mixed_precision = "bf16", base_model = "x")
|
||||
assert _resolve_base_precision(cfg, None, "cuda") == "mxfp8"
|
||||
|
||||
|
||||
def test_mx_module_filter_skips_lora_and_proj_out():
|
||||
# LoRA-owned modules (adapters stay high precision) and the output projection are
|
||||
# excluded, mirroring the fp8 filter's guards.
|
||||
lin = _linear(3072, 3072)
|
||||
assert _mx_module_filter(lin, "blocks.0.attn.to_q.lora_A.default") is False
|
||||
assert _mx_module_filter(lin, "proj_out") is False
|
||||
assert _mx_module_filter(lin, "x.proj_out.y") is False
|
||||
|
||||
|
||||
def test_mx_module_filter_rejects_non_block_aligned_dims():
|
||||
# MX block scaling tiles 32-wide, so a dim not divisible by 32 (3000) is rejected.
|
||||
assert _mx_module_filter(_linear(3000, 3072), "blocks.0.ff.up") is False
|
||||
|
||||
|
||||
def test_mx_module_filter_rejects_non_linear():
|
||||
import torch.nn as nn
|
||||
|
||||
# A non-Linear module is never a target even if it exposes matching feature counts.
|
||||
assert _mx_module_filter(nn.LayerNorm(3072), "blocks.0.norm") is False
|
||||
|
||||
|
||||
def test_should_compile_auto_mxfp8_on_cuda():
|
||||
# auto compiles the dense speed modes on cuda; int8 stays eager (torchao subclass);
|
||||
# an explicit "off" wins over the mode.
|
||||
cfg = DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o")
|
||||
assert _should_compile(cfg, False, "cuda", base_precision = "mxfp8") is True
|
||||
assert _should_compile(cfg, False, "cuda", base_precision = "int8") is False
|
||||
off = DiffusionLoraConfig(
|
||||
base_model = "b", data_dir = "d", output_dir = "o", compile_transformer = "off"
|
||||
)
|
||||
assert _should_compile(off, False, "cuda", base_precision = "mxfp8") is False
|
||||
|
||||
|
||||
def test_apply_mxfp8_training_failure_falls_back_with_warning(monkeypatch):
|
||||
# An unavailable torchao MX path must never be fatal: force both API revisions'
|
||||
# imports to raise, then assert the helper returns False and emits exactly one
|
||||
# warning naming mxfp8.
|
||||
monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", None)
|
||||
monkeypatch.setitem(sys.modules, "torchao.prototype.moe_training.config", None)
|
||||
events = []
|
||||
ok = _apply_mxfp8_training(object(), lambda e: events.append(e))
|
||||
assert ok is False
|
||||
warnings = [e for e in events if e["type"] == "warning"]
|
||||
assert len(warnings) == 1
|
||||
assert "mxfp8" in warnings[0]["message"]
|
||||
|
||||
|
||||
def test_mxfp8_training_config_falls_back_to_the_torchao_0_17_api(monkeypatch):
|
||||
# torchao 0.17 removed prototype.mx_formats.MXLinearConfig in favour of the
|
||||
# MXFP8TrainingOpConfig recipe API; the config helper must fall back to it so the
|
||||
# advertised mxfp8 mode keeps engaging on those installs instead of silently
|
||||
# training dense bf16.
|
||||
from types import SimpleNamespace
|
||||
|
||||
from core.training.diffusion_dit_trainer import _mxfp8_training_config
|
||||
|
||||
calls = {}
|
||||
|
||||
class _Recipe:
|
||||
MXFP8_RCEIL = "mxfp8_rceil"
|
||||
|
||||
class _OpConfig:
|
||||
@staticmethod
|
||||
def from_recipe(recipe):
|
||||
calls["recipe"] = recipe
|
||||
return "cfg-0.17"
|
||||
|
||||
fake_config = SimpleNamespace(MXFP8TrainingOpConfig = _OpConfig, MXFP8TrainingRecipe = _Recipe)
|
||||
monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", None)
|
||||
monkeypatch.setitem(
|
||||
sys.modules, "torchao.prototype.moe_training", SimpleNamespace(config = fake_config)
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "torchao.prototype.moe_training.config", fake_config)
|
||||
assert _mxfp8_training_config() == "cfg-0.17"
|
||||
assert calls["recipe"] == _Recipe.MXFP8_RCEIL
|
||||
|
||||
|
||||
def _patch_capability(monkeypatch, capability):
|
||||
# Drive train_precision_modes' GPU probe: pretend CUDA is present at the given tensor
|
||||
# core capability (fp8 needs sm89+, mxfp8 needs sm100+). The torchao probe is stubbed
|
||||
# functional so these tests exercise the CAPABILITY gate on hosts without torchao
|
||||
# (the CPU-only CI runner does not install it).
|
||||
import torch
|
||||
|
||||
import core.training.diffusion_train_common as dtc
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
|
||||
monkeypatch.setattr(dtc, "has_functional_torchao", lambda: True)
|
||||
|
||||
|
||||
def test_train_precision_modes_blackwell_lists_mxfp8(monkeypatch):
|
||||
# sm100 (Blackwell) exposes both fp8 and mxfp8, ordered before the "auto" pick.
|
||||
_patch_capability(monkeypatch, (10, 0))
|
||||
modes, recommended = train_precision_modes()
|
||||
assert "mxfp8" in modes and "fp8" in modes
|
||||
assert modes.index("mxfp8") < modes.index("auto")
|
||||
assert modes.index("fp8") < modes.index("auto")
|
||||
assert recommended == "auto"
|
||||
|
||||
|
||||
def test_train_precision_modes_ada_has_fp8_without_mxfp8(monkeypatch):
|
||||
# sm89 (Ada) is fp8-capable but not block-scaled mxfp8-capable.
|
||||
_patch_capability(monkeypatch, (8, 9))
|
||||
modes, _ = train_precision_modes()
|
||||
assert "fp8" in modes
|
||||
assert "mxfp8" not in modes
|
||||
|
||||
|
||||
def test_train_precision_modes_newer_blackwell_has_mxfp8(monkeypatch):
|
||||
# Any capability >= sm100 keeps mxfp8 (sm120 here).
|
||||
_patch_capability(monkeypatch, (12, 0))
|
||||
modes, _ = train_precision_modes()
|
||||
assert "mxfp8" in modes
|
||||
|
|
|
|||
|
|
@ -115,6 +115,48 @@ def test_config_normalized_validation(kw):
|
|||
DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw).normalized()
|
||||
|
||||
|
||||
def test_config_normalized_accepts_mxfp8_dense_base():
|
||||
# mxfp8 is a dense speed mode: a dense base + bf16 compute normalises through.
|
||||
cfg = DiffusionLoraConfig(
|
||||
base_model = "black-forest-labs/FLUX.1-dev",
|
||||
data_dir = "d",
|
||||
output_dir = "o",
|
||||
base_precision = "mxfp8",
|
||||
).normalized()
|
||||
assert cfg.base_precision == "mxfp8"
|
||||
|
||||
|
||||
def test_config_normalized_mxfp8_rejects_prequant_base():
|
||||
# A prequant (bnb-4bit) base cannot serve the dense mxfp8 base precision.
|
||||
with pytest.raises(ValueError, match = "mxfp8"):
|
||||
DiffusionLoraConfig(
|
||||
base_model = "unsloth/Qwen-Image-2512-unsloth-bnb-4bit",
|
||||
data_dir = "d",
|
||||
output_dir = "o",
|
||||
base_precision = "mxfp8",
|
||||
).normalized()
|
||||
|
||||
|
||||
def test_config_normalized_mxfp8_requires_bf16_compute():
|
||||
# Like the other dense modes, mxfp8 trains in bf16 compute; fp16 is refused.
|
||||
with pytest.raises(ValueError, match = "mxfp8"):
|
||||
DiffusionLoraConfig(
|
||||
base_model = "black-forest-labs/FLUX.1-dev",
|
||||
data_dir = "d",
|
||||
output_dir = "o",
|
||||
base_precision = "mxfp8",
|
||||
mixed_precision = "fp16",
|
||||
).normalized()
|
||||
|
||||
|
||||
def test_config_normalized_lists_mxfp8_in_invalid_mode_error():
|
||||
# The invalid-base_precision message enumerates the allowed modes, including mxfp8.
|
||||
with pytest.raises(ValueError, match = "mxfp8"):
|
||||
DiffusionLoraConfig(
|
||||
base_model = "b", data_dir = "d", output_dir = "o", base_precision = "bogus"
|
||||
).normalized()
|
||||
|
||||
|
||||
def _cfg(**kw):
|
||||
return DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw)
|
||||
|
||||
|
|
|
|||
|
|
@ -393,6 +393,20 @@ def test_request_model_num_epochs_bounds():
|
|||
DiffusionTrainingStartRequest(**base, num_epochs = bad)
|
||||
|
||||
|
||||
def test_request_model_base_precision_accepts_mxfp8():
|
||||
# The base_precision Literal now includes mxfp8 (the DiT dense speed mode); a bogus
|
||||
# mode is still rejected.
|
||||
from pydantic import ValidationError
|
||||
|
||||
from models.training import DiffusionTrainingStartRequest
|
||||
|
||||
base = {"base_model": "b", "data_dir": "d", "output_dir": "o"}
|
||||
assert DiffusionTrainingStartRequest(**base).base_precision == "nf4" # default
|
||||
assert DiffusionTrainingStartRequest(**base, base_precision = "mxfp8").base_precision == "mxfp8"
|
||||
with pytest.raises(ValidationError):
|
||||
DiffusionTrainingStartRequest(**base, base_precision = "bogus")
|
||||
|
||||
|
||||
def test_config_from_dict_epoch_mode_drops_max_steps_sentinel():
|
||||
# The generic Studio epoch-mode payload sends max_steps: 0 as the "use epochs" sentinel.
|
||||
# The max_steps -> train_steps alias would copy that 0 and normalized() would reject
|
||||
|
|
|
|||
|
|
@ -287,9 +287,10 @@ export interface DiffusionTrainingStartRequest {
|
|||
lr_warmup_steps?: number;
|
||||
// DiT-family quantised base precision (nf4 QLoRA by default). Ignored for sdxl, which
|
||||
// uses mixed_precision instead. "auto" lets the backend pick per family.
|
||||
base_precision?: "nf4" | "bf16" | "int8" | "fp8" | "auto";
|
||||
// Whether to torch.compile the transformer (DiT families that support it). "auto" lets
|
||||
// the backend decide; "off"/"on" force it.
|
||||
base_precision?: "nf4" | "bf16" | "int8" | "fp8" | "mxfp8" | "auto";
|
||||
// Whether to torch.compile the transformer (any family whose /info reports
|
||||
// supports_compile; that includes the SDXL U-Net). "auto" lets the backend decide;
|
||||
// "off"/"on" force it.
|
||||
compile_transformer?: "off" | "on" | "auto";
|
||||
// Precompute + cache the VAE latents before the loop (skips re-encoding each epoch).
|
||||
cache_latents?: boolean;
|
||||
|
|
|
|||
|
|
@ -94,7 +94,7 @@ const CUSTOM_BASE = "__custom__";
|
|||
const UPLOAD_DATASET = "__upload__";
|
||||
// The dense DiT base precisions: they load a dense (bf16) base and quantise/cast it, so the
|
||||
// backend rejects them for an already-quantised bnb-4bit repo. "nf4"/"auto" stay valid.
|
||||
const DENSE_PRECISIONS = new Set(["bf16", "int8", "fp8"]);
|
||||
const DENSE_PRECISIONS = new Set(["bf16", "int8", "fp8", "mxfp8"]);
|
||||
// Mirror the backend's repo_is_prequantized heuristic: a repo whose name marks a
|
||||
// bitsandbytes 4-bit build already ships a quantised transformer and cannot serve the dense
|
||||
// base precisions. Kept in sync with diffusion_train_common.repo_is_prequantized.
|
||||
|
|
@ -197,17 +197,22 @@ export function DiffusionTrainPanel({
|
|||
const isDiT = familyName !== "sdxl";
|
||||
// The quantised base precisions this family can train in, with a stable fallback when the
|
||||
// backend does not report them (older backend, or a preset-only family).
|
||||
const precisionModes = useMemo<Array<"nf4" | "bf16" | "int8" | "fp8" | "auto">>(() => {
|
||||
const precisionModes = useMemo<
|
||||
Array<"nf4" | "bf16" | "int8" | "fp8" | "mxfp8" | "auto">
|
||||
>(() => {
|
||||
const reported = reportedFamily?.precision_modes?.filter(
|
||||
(m): m is "nf4" | "bf16" | "int8" | "fp8" =>
|
||||
m === "nf4" || m === "bf16" || m === "int8" || m === "fp8",
|
||||
(m): m is "nf4" | "bf16" | "int8" | "fp8" | "mxfp8" =>
|
||||
m === "nf4" || m === "bf16" || m === "int8" || m === "fp8" || m === "mxfp8",
|
||||
);
|
||||
if (reported && reported.length > 0) return ["auto", ...reported];
|
||||
// Fallback without a backend report: the GPU-independent modes only (mxfp8 needs a
|
||||
// Blackwell probe, so it is offered strictly when the backend advertises it).
|
||||
return ["auto", "nf4", "bf16", "int8", "fp8"];
|
||||
}, [reportedFamily?.precision_modes]);
|
||||
// Whether to show the torch.compile control. Default on for DiT families when the backend
|
||||
// does not say otherwise; sdxl's U-Net path does not expose it here.
|
||||
const supportsCompile = isDiT && (reportedFamily?.supports_compile ?? true);
|
||||
// Whether to show the torch.compile control. The backend advertises this per family
|
||||
// (the SDXL U-Net path compiles regionally too now); default on for DiT families when
|
||||
// an older backend does not report it.
|
||||
const supportsCompile = reportedFamily?.supports_compile ?? isDiT;
|
||||
|
||||
const [baseChoice, setBaseChoice] = useState<string>(family?.base_repos[0] ?? "");
|
||||
const [customBase, setCustomBase] = useState("");
|
||||
|
|
@ -250,7 +255,7 @@ export function DiffusionTrainPanel({
|
|||
// lets the backend pick the family's recommended mode. Re-seeded to the family's
|
||||
// recommendation on family change (unless the user picked one).
|
||||
const [basePrecision, setBasePrecision] = useState<
|
||||
"nf4" | "bf16" | "int8" | "fp8" | "auto"
|
||||
"nf4" | "bf16" | "int8" | "fp8" | "mxfp8" | "auto"
|
||||
>("auto");
|
||||
// Whether to torch.compile the DiT transformer. "auto" defers to the backend.
|
||||
const [compileTransformer, setCompileTransformer] = useState<"off" | "on" | "auto">(
|
||||
|
|
@ -816,11 +821,14 @@ export function DiffusionTrainPanel({
|
|||
</div>
|
||||
);
|
||||
|
||||
const precisionLabel = (m: "nf4" | "bf16" | "int8" | "fp8" | "auto"): string => {
|
||||
const precisionLabel = (
|
||||
m: "nf4" | "bf16" | "int8" | "fp8" | "mxfp8" | "auto",
|
||||
): string => {
|
||||
if (m === "auto") return "Auto (recommended)";
|
||||
if (m === "nf4") return "nf4 (4-bit QLoRA, lowest VRAM)";
|
||||
if (m === "bf16") return "bf16 (fastest, most VRAM)";
|
||||
if (m === "int8") return "int8 (8-bit)";
|
||||
if (m === "mxfp8") return "mxfp8 (Blackwell, best at high res/batch)";
|
||||
return "fp8 (experimental)";
|
||||
};
|
||||
|
||||
|
|
@ -883,63 +891,40 @@ export function DiffusionTrainPanel({
|
|||
</div>
|
||||
|
||||
{isDiT ? (
|
||||
<>
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Base precision</Label>
|
||||
<select
|
||||
value={basePrecision}
|
||||
onChange={(e) => {
|
||||
precisionDirty.current = true;
|
||||
setBasePrecision(e.target.value as typeof basePrecision);
|
||||
}}
|
||||
className={selectClass}
|
||||
aria-label="Base precision"
|
||||
>
|
||||
{precisionModes.map((m) => (
|
||||
<option
|
||||
key={m}
|
||||
value={m}
|
||||
disabled={basePrequantized && DENSE_PRECISIONS.has(m)}
|
||||
>
|
||||
{precisionLabel(m)}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
<p className="text-[11px] leading-snug text-muted-foreground">
|
||||
How the frozen base weights are quantised. nf4 (4-bit) uses the least VRAM;
|
||||
bf16 is fastest but needs the most. Auto picks this family's recommended
|
||||
mode.
|
||||
{basePrequantized && (
|
||||
<>
|
||||
{" "}
|
||||
This base is already 4-bit quantised, so only nf4/auto apply; pick a dense
|
||||
(bf16) base repo for the other modes.
|
||||
</>
|
||||
)}
|
||||
</p>
|
||||
</div>
|
||||
{supportsCompile && (
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Compile transformer</Label>
|
||||
<select
|
||||
value={compileTransformer}
|
||||
onChange={(e) =>
|
||||
setCompileTransformer(e.target.value as typeof compileTransformer)
|
||||
}
|
||||
className={selectClass}
|
||||
aria-label="Compile transformer"
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Base precision</Label>
|
||||
<select
|
||||
value={basePrecision}
|
||||
onChange={(e) => {
|
||||
precisionDirty.current = true;
|
||||
setBasePrecision(e.target.value as typeof basePrecision);
|
||||
}}
|
||||
className={selectClass}
|
||||
aria-label="Base precision"
|
||||
>
|
||||
{precisionModes.map((m) => (
|
||||
<option
|
||||
key={m}
|
||||
value={m}
|
||||
disabled={basePrequantized && DENSE_PRECISIONS.has(m)}
|
||||
>
|
||||
<option value="auto">Auto</option>
|
||||
<option value="on">On (faster after warmup)</option>
|
||||
<option value="off">Off</option>
|
||||
</select>
|
||||
<p className="text-[11px] leading-snug text-muted-foreground">
|
||||
torch.compile the transformer. Adds a one-time warmup, then speeds up each
|
||||
step.
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
</>
|
||||
{precisionLabel(m)}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
<p className="text-[11px] leading-snug text-muted-foreground">
|
||||
How the frozen base weights are quantised. nf4 (4-bit) uses the least VRAM;
|
||||
bf16 is fastest but needs the most. Auto picks this family's recommended
|
||||
mode.
|
||||
{basePrequantized && (
|
||||
<>
|
||||
{" "}
|
||||
This base is already 4-bit quantised, so only nf4/auto apply; pick a dense
|
||||
(bf16) base repo for the other modes.
|
||||
</>
|
||||
)}
|
||||
</p>
|
||||
</div>
|
||||
) : (
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Precision</Label>
|
||||
|
|
@ -958,6 +943,27 @@ export function DiffusionTrainPanel({
|
|||
</p>
|
||||
</div>
|
||||
)}
|
||||
{supportsCompile && (
|
||||
<div className="grid gap-1.5">
|
||||
<Label className="text-xs">Compile transformer</Label>
|
||||
<select
|
||||
value={compileTransformer}
|
||||
onChange={(e) =>
|
||||
setCompileTransformer(e.target.value as typeof compileTransformer)
|
||||
}
|
||||
className={selectClass}
|
||||
aria-label="Compile transformer"
|
||||
>
|
||||
<option value="auto">Auto</option>
|
||||
<option value="on">On (faster after warmup)</option>
|
||||
<option value="off">Off</option>
|
||||
</select>
|
||||
<p className="text-[11px] leading-snug text-muted-foreground">
|
||||
torch.compile the transformer blocks. Adds a one-time warmup, then speeds up
|
||||
each step.
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue