Add mxfp8 training base precision and SDXL U-Net regional compile
- base_precision="mxfp8": torchao MX block-scaled float8 compute on the frozen base linears (Blackwell sm100+, cuBLAS kernels). Applied after add_adapter like fp8, never fatal, weights stay bf16 in memory. Measured 1.16x over compiled bf16 on Z-Image at 1024px batch 4 (16k tokens/step); a wash at small token counts, so it stays an explicit opt-in and auto never picks it. - SDXL: regionally compile the U-Net's BasicTransformerBlocks through the same never-fatal wrapper the DiT trainer uses. 1.35x steady state at 1024px batch 4 with same-seed loss parity (~1e-5 per step) and unchanged peak VRAM; ~30 s one-time warmup. Steady-state samples/sec now excludes step 1, matching the DiT trainer. - /info: mxfp8 advertised only on sm100+; supports_compile now true for sdxl. - NVFP4 training: not available in torchao 0.16 (no autograd path, no training recipe), so NVFP4 stays an inference-only quant for now. 193 diffusion backend tests green; frontend build clean.
This commit is contained in:
Daniel Han
parent
11a00d0238
commit
cb9247e537
9 changed files with 294 additions and 41 deletions
|
|
@ -218,8 +218,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
|
||||
|
|
@ -238,7 +238,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).
|
||||
|
|
@ -299,6 +299,45 @@ 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
|
||||
return mod.in_features % 32 == 0 and mod.out_features % 32 == 0
|
||||
|
||||
|
||||
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.prototype.mx_formats import MXLinearConfig
|
||||
from torchao.quantization import quantize_
|
||||
quantize_(
|
||||
transformer,
|
||||
MXLinearConfig.from_recipe_name("mxfp8_cublas"),
|
||||
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, free_gb, dense_gb, capability, has_fp8) -> str:
|
||||
"""Pure policy for base_precision="auto": nf4 for a prequant base or no CUDA; else the
|
||||
fastest dense mode whose weights + headroom (activations, optimizer, cache) fit the
|
||||
|
|
@ -330,7 +369,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'."
|
||||
|
|
@ -977,8 +1016,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(
|
||||
|
|
@ -994,11 +1034,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
|
||||
|
|
@ -1177,8 +1220,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"
|
||||
|
|
@ -1207,12 +1250,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
|
||||
|
|
@ -353,6 +355,15 @@ 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
|
||||
|
|
@ -429,7 +440,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)
|
||||
|
||||
def _next_batch() -> tuple[list[int], list[str], list[str]]:
|
||||
idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs)))
|
||||
|
|
@ -442,6 +453,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)
|
||||
|
|
@ -524,17 +536,25 @@ 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,15 +128,17 @@ 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
|
||||
|
||||
|
||||
def train_precision_modes() -> tuple[list[str], str]:
|
||||
"""(supported base_precision modes, recommended pick) for the current machine: nf4
|
||||
always works; bf16/int8/auto need CUDA; fp8 needs an fp8-capable GPU (sm89+). Used by
|
||||
the /info endpoint so the UI can gate the precision selector. Never raises."""
|
||||
always works; bf16/int8/auto need CUDA; fp8 needs an fp8-capable GPU (sm89+); mxfp8
|
||||
(block-scaled fp8 compute) needs the Blackwell tensor cores (sm100+) its cuBLAS
|
||||
kernels target. Used by the /info endpoint so the UI can gate the precision
|
||||
selector. Never raises."""
|
||||
modes = ["nf4"]
|
||||
recommended = "nf4"
|
||||
try:
|
||||
|
|
@ -146,6 +148,8 @@ def train_precision_modes() -> tuple[list[str], str]:
|
|||
major, minor = torch.cuda.get_device_capability()
|
||||
if (major, minor) >= (8, 9) and hasattr(torch, "float8_e4m3fn"):
|
||||
modes.append("fp8")
|
||||
if (major, minor) >= (10, 0):
|
||||
modes.append("mxfp8")
|
||||
modes.append("auto")
|
||||
recommended = "auto"
|
||||
except Exception: # noqa: BLE001 -- no torch / probe failure -> nf4 only
|
||||
|
|
@ -219,8 +223,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 ("flux.1", "qwen-image", "z-image", "krea-2")
|
||||
infos.append(
|
||||
{
|
||||
|
|
@ -232,7 +237,7 @@ def family_train_infos() -> list[dict[str, Any]]:
|
|||
"vram_note": _FAMILY_VRAM_NOTES.get(name, ""),
|
||||
"precision_modes": dit_modes if is_dit else [],
|
||||
"recommended_precision": dit_recommended if is_dit else "nf4",
|
||||
"supports_compile": is_dit,
|
||||
"supports_compile": True,
|
||||
}
|
||||
)
|
||||
return infos
|
||||
|
|
@ -329,9 +334,11 @@ 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")
|
||||
if 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"
|
||||
)
|
||||
if 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 "
|
||||
|
|
|
|||
|
|
@ -736,11 +736,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."
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -9,16 +9,25 @@ exercised by the live GPU smokes, not here."""
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from core.training.diffusion_dit_trainer import (
|
||||
_GATED_TRAIN_REPOS,
|
||||
_SPECS,
|
||||
_apply_mxfp8_training,
|
||||
_assert_gated_access,
|
||||
_mx_module_filter,
|
||||
_repo_is_prequantized,
|
||||
_should_compile,
|
||||
run_dit_lora_training,
|
||||
)
|
||||
from core.training.diffusion_train_common import DiffusionLoraConfig, family_train_infos
|
||||
from core.training.diffusion_train_common import (
|
||||
DiffusionLoraConfig,
|
||||
family_train_infos,
|
||||
train_precision_modes,
|
||||
)
|
||||
|
||||
|
||||
def test_specs_cover_the_dit_families():
|
||||
|
|
@ -85,3 +94,109 @@ 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):
|
||||
import torch.nn as nn
|
||||
|
||||
return nn.Linear(in_features, out_features)
|
||||
|
||||
|
||||
def test_mx_module_filter_accepts_dense_block_linear():
|
||||
# A 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_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 the import to raise, then
|
||||
# assert the helper returns False and emits exactly one warning naming mxfp8.
|
||||
monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", 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 _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+).
|
||||
import torch
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
|
||||
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -103,6 +103,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -339,6 +339,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_route_start_rejects_uncontained_paths(client):
|
||||
# An absolute path outside the Studio dataset roots is a 400, not silently accepted.
|
||||
r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "/etc"})
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -182,17 +182,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("");
|
||||
|
|
@ -235,7 +240,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">(
|
||||
|
|
@ -737,11 +742,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)";
|
||||
};
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue