Add base_precision speed modes to DiT training: bf16 2.3-2.6x, int8, fp8
New base_precision config for the DiT trainers: nf4 (unchanged default) | bf16 | int8 | fp8 | auto, advertised per family + per machine through /api/train/diffusion/info (precision_modes, recommended_precision, supports_compile) so the UI can gate the selector. - bf16: dense transformer + regional torch.compile (auto-armed). The measured speed mode: 2.3x nf4 on FLUX (1.81 -> 4.12 steps/s), 2.6x on Z-Image (2.5 -> 6.38 steps/s) on B200, at dense-weight VRAM (FLUX 24.7 GB / Z-Image 13.6 GB peak vs 10.4 / 4.7 for nf4). - int8: torchao weight-only int8 on the frozen base, quantized AFTER add_adapter (quantizing first trips peft 0.18's TorchaoLoraLinear, which is incompatible with the torchao 0.16 config API). Runs eager: inductor rejects the int8 subclass training graph (aliased subclass outputs), so compile is force-disabled for it. - fp8: torchao convert_to_float8_training on the frozen linears (filter skips lora_ modules, proj_out, non-divisible-by-16 dims, pad_inner_dim), applied after add_adapter, compile auto-armed. Works and round-trips, but measured SLOWER than compiled bf16 at LoRA-training shapes (FLUX 3.15 vs 4.12 steps/s; Z-Image similar), so it is an explicit opt-in and auto never picks it. - auto: free VRAM (measured before load) + dense-size table -> bf16 when it fits with headroom, int8 in the middle band, else nf4. Prequant bnb repos always resolve to nf4; dense modes on them are rejected at validation with a pointer to the family's dense base. Two crashes found and fixed along the way: - The cuDNN SDPA backend's training graph fails on the FLUX attention shapes (torch 2.10 + cu130, B200): mha_graph.execute errors, then the context degrades into illegal memory accesses. The perf-flag guard now pins flash/mem-efficient SDPA for the run (mathematically equivalent, snapshot/restored). nf4 escaped it by routing attention differently. - Regional compile now uses dynamic=True (the inference layer's proven default): dynamic=False specialisation fused a gemm_and_bias epilogue that failed with CUBLAS_STATUS_EXECUTION_FAILED on the FLUX training graph; dynamic=True is also faster (Z-Image 3.84 -> 6.38 steps/s). Verified: 98 backend tests green (new test_diffusion_base_precision.py: validation, auto policy table, fp8 filter, compile gating, /info fields); per-mode 40-step runs on FLUX + Z-Image with loss means inside the nf4 envelope and adapter round-trip generation through the normal LoRA path for bf16-, fp8-, and int8-trained adapters.
This commit is contained in:
Daniel Han
parent
9d967d2c5b
commit
5f725bacf2
5 changed files with 486 additions and 50 deletions
|
|
@ -47,6 +47,7 @@ from core.training.diffusion_train_common import (
|
|||
_publish_to_lora_catalog,
|
||||
_restore_perf_flags,
|
||||
discover_image_caption_pairs,
|
||||
repo_is_prequantized,
|
||||
)
|
||||
|
||||
# Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks
|
||||
|
|
@ -74,6 +75,9 @@ class _FamilySpec:
|
|||
lora_targets: tuple[str, ...]
|
||||
# bf16 only (Z-Image overflows fp16 and its RoPE/embedder run in fp32).
|
||||
force_bf16: bool
|
||||
# Approximate dense-bf16 transformer weight size, used by base_precision="auto" to
|
||||
# decide which mode fits the free VRAM (with headroom for activations + optimizer).
|
||||
dense_bf16_gb: float
|
||||
# Phased load, so the (multi-GB) transformer never has to coexist with the text
|
||||
# encoders + VAE: ``load_conditioners`` builds the pipeline WITHOUT its transformer
|
||||
# (encode_prompt + VAE only) and returns (pipe, vae); ``load_transformer`` loads the
|
||||
|
|
@ -155,12 +159,9 @@ def _bnb_4bit_config():
|
|||
)
|
||||
|
||||
|
||||
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 we load it as-is rather than re-quantizing on the fly. A
|
||||
dense (bf16) base instead gets on-the-fly nf4 quantization for QLoRA."""
|
||||
name = str(base_model or "").lower()
|
||||
return "bnb-4bit" in name or "-4bit" in name or "int4" in name or "nf4" in name
|
||||
# Kept as a module name for existing callers/tests; the heuristic itself moved to
|
||||
# diffusion_train_common so config validation can use it without importing this module.
|
||||
_repo_is_prequantized = repo_is_prequantized
|
||||
|
||||
|
||||
def _load_quantized_transformer(transformer_cls, cfg):
|
||||
|
|
@ -191,24 +192,132 @@ def _load_pipe_without_transformer(pipe_cls, cfg, device):
|
|||
return pipe, pipe.vae
|
||||
|
||||
|
||||
def _load_dit_transformer(transformer_cls, cfg, device, qlora):
|
||||
"""Load the transformer alone. 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 when
|
||||
``qlora`` (the default) so QLoRA still fits."""
|
||||
def _load_dit_transformer(transformer_cls, cfg, device, base_precision):
|
||||
"""Load the transformer alone in the resolved ``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).
|
||||
- int8: the dense transformer quantized in place to torchao weight-only int8 (the
|
||||
PEFT-attachable scheme), roughly halving the bf16 weight footprint."""
|
||||
import torch
|
||||
|
||||
if qlora and not _repo_is_prequantized(cfg.base_model):
|
||||
return _load_quantized_transformer(transformer_cls, cfg)
|
||||
transformer = transformer_cls.from_pretrained(
|
||||
if base_precision == "nf4":
|
||||
if not repo_is_prequantized(cfg.base_model):
|
||||
return _load_quantized_transformer(transformer_cls, cfg)
|
||||
transformer = transformer_cls.from_pretrained(
|
||||
cfg.base_model,
|
||||
subfolder = "transformer",
|
||||
torch_dtype = torch.bfloat16,
|
||||
token = cfg.hf_token,
|
||||
)
|
||||
# A prequant load is already device-placed by bitsandbytes.
|
||||
if not getattr(transformer, "is_loaded_in_4bit", False):
|
||||
transformer = transformer.to(device)
|
||||
return transformer
|
||||
|
||||
# Dense load for bf16 / fp8 / 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).
|
||||
return transformer_cls.from_pretrained(
|
||||
cfg.base_model,
|
||||
subfolder = "transformer",
|
||||
torch_dtype = torch.bfloat16,
|
||||
token = cfg.hf_token,
|
||||
).to(device)
|
||||
|
||||
|
||||
def _int8_quantize_base(transformer) -> None:
|
||||
"""torchao weight-only int8 on the big frozen linears, applied after add_adapter so
|
||||
the base_layer inside each LoRA wrapper quantizes while the adapters stay high
|
||||
precision. ``make_filter_fn`` (shared with the inference quant layer) keeps only
|
||||
Linears with >= 512 features -- which also naturally skips the rank-sized LoRA
|
||||
matrices -- and drops the M=1 modulation projections int8 kernels reject."""
|
||||
from core.inference.diffusion_transformer_quant import exclude_tokens_for_scheme, make_filter_fn
|
||||
from torchao.quantization import Int8WeightOnlyConfig, quantize_
|
||||
|
||||
quantize_(
|
||||
transformer,
|
||||
Int8WeightOnlyConfig(),
|
||||
filter_fn = make_filter_fn(512, exclude_name_tokens = exclude_tokens_for_scheme("int8")),
|
||||
)
|
||||
# A bnb-quantized load is already device-placed; a dense one still sits on CPU.
|
||||
if not getattr(transformer, "is_loaded_in_4bit", False):
|
||||
transformer = transformer.to(device)
|
||||
return transformer
|
||||
|
||||
|
||||
def _fp8_module_filter(mod, fqn: str) -> bool:
|
||||
"""Which frozen linears get float8 training compute: skip anything LoRA-owned (the
|
||||
adapters must stay high precision -- PEFT has no float8 base support), the output
|
||||
projection, and shapes float8 kernels reject (dims not divisible by 16), matching the
|
||||
diffusers FLUX2 reference filter."""
|
||||
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 % 16 == 0 and mod.out_features % 16 == 0
|
||||
|
||||
|
||||
def _apply_fp8_training(transformer, on_event) -> bool:
|
||||
"""Convert the frozen base linears to torchao float8 training compute (dynamic scaling;
|
||||
weights stay bf16 in memory). Applied AFTER add_adapter so the filter can exclude the
|
||||
LoRA modules. Never fatal: on any failure the run continues in bf16 with a warning."""
|
||||
try:
|
||||
from torchao.float8 import Float8LinearConfig, convert_to_float8_training
|
||||
|
||||
convert_to_float8_training(
|
||||
transformer,
|
||||
module_filter_fn = _fp8_module_filter,
|
||||
config = Float8LinearConfig(pad_inner_dim = True),
|
||||
)
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 -- fp8 is an optimisation, never fatal
|
||||
_emit(on_event, "warning", message = f"fp8 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
|
||||
free VRAM at decision time. bf16 + regional compile is the measured speed winner
|
||||
(2.3-2.6x over nf4 on B200); fp8 stays an explicit opt-in because torchao float8's
|
||||
dynamic-scaling overhead made it SLOWER than compiled bf16 at LoRA-training shapes on
|
||||
the same hardware. ``capability``/``has_fp8`` remain parameters so the policy can be
|
||||
revisited per GPU generation without changing callers."""
|
||||
_ = capability, has_fp8
|
||||
if prequant or device != "cuda" or not free_gb or not dense_gb:
|
||||
return "nf4"
|
||||
headroom = 1.5
|
||||
if free_gb > dense_gb * headroom:
|
||||
return "bf16"
|
||||
if free_gb > dense_gb * 0.55 * headroom:
|
||||
return "int8"
|
||||
return "nf4"
|
||||
|
||||
|
||||
def _resolve_base_precision(cfg, spec, device) -> str:
|
||||
"""Resolve "auto" against the live GPU (free VRAM measured BEFORE anything loads);
|
||||
explicit modes pass through (normalized() already validated them)."""
|
||||
mode = (cfg.base_precision or "nf4").strip().lower()
|
||||
if mode != "auto":
|
||||
return mode
|
||||
prequant = repo_is_prequantized(cfg.base_model)
|
||||
free_gb = None
|
||||
capability = None
|
||||
has_fp8 = False
|
||||
if device == "cuda":
|
||||
try:
|
||||
import torch
|
||||
|
||||
free_gb = torch.cuda.mem_get_info()[0] / 1e9
|
||||
capability = torch.cuda.get_device_capability()
|
||||
has_fp8 = hasattr(torch, "float8_e4m3fn")
|
||||
except Exception: # noqa: BLE001 -- probe failure -> the safe mode
|
||||
pass
|
||||
return _pick_auto_precision(prequant, device, free_gb, spec.dense_bf16_gb, capability, has_fp8)
|
||||
|
||||
|
||||
# ── FLUX.1-dev ────────────────────────────────────────────────────────────────
|
||||
|
|
@ -217,9 +326,9 @@ def _flux_load_conditioners(cfg, device, weight_dtype):
|
|||
return _load_pipe_without_transformer(FluxPipeline, cfg, device)
|
||||
|
||||
|
||||
def _flux_load_transformer(cfg, device, weight_dtype, qlora):
|
||||
def _flux_load_transformer(cfg, device, weight_dtype, base_precision):
|
||||
from diffusers import FluxTransformer2DModel
|
||||
return _load_dit_transformer(FluxTransformer2DModel, cfg, device, qlora)
|
||||
return _load_dit_transformer(FluxTransformer2DModel, cfg, device, base_precision)
|
||||
|
||||
|
||||
def _flux_encode_prompts(pipe, captions, device):
|
||||
|
|
@ -325,12 +434,12 @@ def _qwen_load_conditioners(cfg, device, weight_dtype):
|
|||
return _load_pipe_without_transformer(QwenImagePipeline, cfg, device)
|
||||
|
||||
|
||||
def _qwen_load_transformer(cfg, device, weight_dtype, qlora):
|
||||
def _qwen_load_transformer(cfg, device, weight_dtype, base_precision):
|
||||
# The prequant default (unsloth/Qwen-Image-2512-unsloth-bnb-4bit) ships the transformer
|
||||
# 4-bit and loads trainable as-is; the dense 20B Qwen/Qwen-Image is quantized to nf4 on
|
||||
# the fly so QLoRA still fits.
|
||||
# 4-bit and loads trainable as-is under nf4; the dense modes need the 20B
|
||||
# Qwen/Qwen-Image base.
|
||||
from diffusers import QwenImageTransformer2DModel
|
||||
return _load_dit_transformer(QwenImageTransformer2DModel, cfg, device, qlora)
|
||||
return _load_dit_transformer(QwenImageTransformer2DModel, cfg, device, base_precision)
|
||||
|
||||
|
||||
def _qwen_encode_prompts(pipe, captions, device):
|
||||
|
|
@ -443,11 +552,11 @@ def _zimage_load_conditioners(cfg, device, weight_dtype):
|
|||
return _load_pipe_without_transformer(ZImagePipeline, cfg, device)
|
||||
|
||||
|
||||
def _zimage_load_transformer(cfg, device, weight_dtype, qlora):
|
||||
# Prequant default loads 4-bit as-is; the dense bf16 Tongyi-MAI base is quantized to nf4
|
||||
# on the fly. Z-Image is bf16 only (its RoPE/embedder run fp32; fp16 overflows).
|
||||
def _zimage_load_transformer(cfg, device, weight_dtype, base_precision):
|
||||
# Prequant default loads 4-bit as-is under nf4; the dense modes use the bf16 Tongyi-MAI
|
||||
# base. Z-Image is bf16 only (its RoPE/embedder run fp32; fp16 overflows).
|
||||
from diffusers import ZImageTransformer2DModel
|
||||
return _load_dit_transformer(ZImageTransformer2DModel, cfg, device, qlora)
|
||||
return _load_dit_transformer(ZImageTransformer2DModel, cfg, device, base_precision)
|
||||
|
||||
|
||||
def _zimage_encode_prompts(pipe, captions, device):
|
||||
|
|
@ -513,6 +622,7 @@ _SPECS: dict[str, _FamilySpec] = {
|
|||
family = "flux.1",
|
||||
lora_targets = _FLUX_TARGETS,
|
||||
force_bf16 = False,
|
||||
dense_bf16_gb = 23.8,
|
||||
load_conditioners = _flux_load_conditioners,
|
||||
load_transformer = _flux_load_transformer,
|
||||
encode_prompts = _flux_encode_prompts,
|
||||
|
|
@ -526,6 +636,7 @@ _SPECS: dict[str, _FamilySpec] = {
|
|||
family = "qwen-image",
|
||||
lora_targets = _QWEN_TARGETS,
|
||||
force_bf16 = True,
|
||||
dense_bf16_gb = 41.0,
|
||||
load_conditioners = _qwen_load_conditioners,
|
||||
load_transformer = _qwen_load_transformer,
|
||||
encode_prompts = _qwen_encode_prompts,
|
||||
|
|
@ -539,6 +650,7 @@ _SPECS: dict[str, _FamilySpec] = {
|
|||
family = "z-image",
|
||||
lora_targets = _ZIMAGE_TARGETS,
|
||||
force_bf16 = True,
|
||||
dense_bf16_gb = 12.3,
|
||||
load_conditioners = _zimage_load_conditioners,
|
||||
load_transformer = _zimage_load_transformer,
|
||||
encode_prompts = _zimage_encode_prompts,
|
||||
|
|
@ -684,23 +796,36 @@ def _sample_cached_latents(cache, idxs, variant_rng, device):
|
|||
return lat_a + lat_b * torch.randn_like(lat_a)
|
||||
|
||||
|
||||
def _should_compile(cfg, base_is_bnb, device) -> bool:
|
||||
def _should_compile(cfg, base_is_bnb, device, base_precision = "nf4") -> bool:
|
||||
mode = (cfg.compile_transformer or "auto").strip().lower()
|
||||
if device != "cuda" or mode == "off":
|
||||
return False
|
||||
# torch.compile cannot trace the torchao int8 subclass in training (inductor rejects
|
||||
# the aliased subclass graph outputs), so int8 always runs eager.
|
||||
if base_precision == "int8":
|
||||
return False
|
||||
if mode == "on":
|
||||
return True
|
||||
# auto: regional compile is a clean win on a dense base but fragile over bitsandbytes
|
||||
# 4-bit modules (graph breaks in the dequant path), so it stays off for QLoRA.
|
||||
return not base_is_bnb
|
||||
# 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")
|
||||
|
||||
|
||||
def _maybe_compile_transformer(transformer, cfg, base_is_bnb, device, on_event) -> bool:
|
||||
def _maybe_compile_transformer(
|
||||
transformer, cfg, base_is_bnb, device, on_event, base_precision = "nf4"
|
||||
) -> bool:
|
||||
"""Regionally compile the transformer blocks (diffusers compile_repeated_blocks) after
|
||||
the LoRA is attached. Never fatal: a wrap failure falls back to eager with a warning
|
||||
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):
|
||||
if not _should_compile(cfg, base_is_bnb, device, base_precision):
|
||||
if base_precision == "fp8":
|
||||
_emit(
|
||||
on_event,
|
||||
"warning",
|
||||
message = "fp8 training without torch.compile is slow; enable compile for the speedup.",
|
||||
)
|
||||
return False
|
||||
import torch
|
||||
|
||||
|
|
@ -719,10 +844,12 @@ def _maybe_compile_transformer(transformer, cfg, base_is_bnb, device, on_event)
|
|||
setattr(dynamo_cfg, attr, max(getattr(dynamo_cfg, attr) or 0, 64))
|
||||
if hasattr(dynamo_cfg, "suppress_errors"):
|
||||
dynamo_cfg.suppress_errors = True
|
||||
# Shapes are static under the latent cache (fixed resolution/batch/text bucket), so
|
||||
# dynamic=False lets inductor specialise. fullgraph only on a dense base: bnb 4-bit
|
||||
# layers graph-break by design.
|
||||
fn(fullgraph = not base_is_bnb, dynamic = False)
|
||||
# dynamic=True matches the inference speed layer's proven default: on torch 2.10 /
|
||||
# B200 the dynamic=False specialisation fused a gemm_and_bias epilogue that failed
|
||||
# with CUBLAS_STATUS_EXECUTION_FAILED then an illegal memory access on the FLUX
|
||||
# training graph. fullgraph only on a dense base: bnb 4-bit layers graph-break by
|
||||
# design.
|
||||
fn(fullgraph = not base_is_bnb, dynamic = True)
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 -- optimisation only, never fatal
|
||||
_emit(on_event, "warning", message = f"torch.compile disabled (eager fallback): {exc}")
|
||||
|
|
@ -858,10 +985,12 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
# the same seed-deterministic sequence whether or not the cache is enabled.
|
||||
variant_rng = random.Random(cfg.seed + 1)
|
||||
|
||||
# Phase 3: only now load the transformer (QLoRA nf4 by default: the prequant Qwen /
|
||||
# Z-Image repos are already 4-bit; FLUX quantizes on the fly).
|
||||
transformer = spec.load_transformer(cfg, device, weight_dtype, qlora = True)
|
||||
base_is_bnb = True # the DiT base is always a bitsandbytes 4-bit QLoRA load today
|
||||
# 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).
|
||||
base_precision = _resolve_base_precision(cfg, spec, device)
|
||||
transformer = spec.load_transformer(cfg, device, weight_dtype, base_precision)
|
||||
base_is_bnb = base_precision == "nf4"
|
||||
|
||||
# Freeze the base; attach the trainable LoRA to the transformer.
|
||||
transformer.requires_grad_(False)
|
||||
|
|
@ -887,7 +1016,16 @@ 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]
|
||||
|
||||
compiled = _maybe_compile_transformer(transformer, cfg, base_is_bnb, device, on_event)
|
||||
# int8 / fp8 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"
|
||||
|
||||
compiled = _maybe_compile_transformer(
|
||||
transformer, cfg, base_is_bnb, device, on_event, base_precision
|
||||
)
|
||||
# Compiled Qwen graphs need one fixed text length across steps: pin the pad bucket to
|
||||
# the dataset's longest caption (encode_prompt already caps it at 1024 tokens).
|
||||
qwen_pad_to = None
|
||||
|
|
@ -908,7 +1046,7 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
num_training_steps = cfg.train_steps,
|
||||
)
|
||||
|
||||
_emit(on_event, "model_load_completed", compiled = compiled)
|
||||
_emit(on_event, "model_load_completed", compiled = compiled, base_precision = base_precision)
|
||||
|
||||
transformer.train()
|
||||
n_images = len(image_paths)
|
||||
|
|
|
|||
|
|
@ -124,6 +124,35 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None
|
|||
return "sdxl"
|
||||
|
||||
|
||||
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."""
|
||||
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."""
|
||||
modes = ["nf4"]
|
||||
recommended = "nf4"
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
modes += ["bf16", "int8"]
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
if (major, minor) >= (8, 9) and hasattr(torch, "float8_e4m3fn"):
|
||||
modes.append("fp8")
|
||||
modes.append("auto")
|
||||
recommended = "auto"
|
||||
except Exception: # noqa: BLE001 -- no torch / probe failure -> nf4 only
|
||||
pass
|
||||
return modes, recommended
|
||||
|
||||
|
||||
def get_trainer(family: str) -> Callable[..., str]:
|
||||
"""Return the training entrypoint for ``family``. Imports the trainer module lazily so
|
||||
this shared module stays free of the heavy trainer imports (and any import cycle)."""
|
||||
|
|
@ -178,12 +207,16 @@ def family_train_infos() -> list[dict[str, Any]]:
|
|||
the family registry so it stays in sync with what the trainers actually support."""
|
||||
from core.inference.diffusion_families import detect_family
|
||||
|
||||
dit_modes, dit_recommended = train_precision_modes()
|
||||
infos: list[dict[str, Any]] = []
|
||||
for name in trainable_family_names():
|
||||
fam = detect_family("", override = name)
|
||||
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.
|
||||
is_dit = name in ("flux.1", "qwen-image", "z-image")
|
||||
infos.append(
|
||||
{
|
||||
"name": name,
|
||||
|
|
@ -192,6 +225,9 @@ 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 [],
|
||||
"recommended_precision": dit_recommended if is_dit else "nf4",
|
||||
"supports_compile": is_dit,
|
||||
}
|
||||
)
|
||||
return infos
|
||||
|
|
@ -240,6 +276,12 @@ class DiffusionLoraConfig:
|
|||
# TF32 matmuls + high fp32 matmul precision + cudnn autotuning for the run. Near-lossless;
|
||||
# disable for strict bit-reproducibility A/Bs.
|
||||
enable_tf32: bool = True
|
||||
# DiT base transformer precision: "nf4" (bitsandbytes QLoRA, the memory floor and the
|
||||
# default), "bf16" (dense, fastest eager, compile-friendly), "int8" (torchao
|
||||
# weight-only, half of bf16), "fp8" (torchao float8 training compute on the frozen
|
||||
# linears, Ada/Hopper/Blackwell + compile), or "auto" (pick by free VRAM + GPU class).
|
||||
# Non-nf4 modes need a dense base repo (not a prequant bnb-4bit one). SDXL ignores it.
|
||||
base_precision: str = "nf4"
|
||||
# How often to emit a progress event (in optimizer steps).
|
||||
log_every: int = 1
|
||||
# Optional explicit family override ("sdxl" / "flux.1" / ...); None = detect from
|
||||
|
|
@ -276,6 +318,21 @@ class DiffusionLoraConfig:
|
|||
compile_transformer = str(self.compile_transformer or "auto").strip().lower()
|
||||
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 repo_is_prequantized(self.base_model):
|
||||
raise ValueError(
|
||||
f"base_precision={base_precision!r} needs a dense base repo, but "
|
||||
f"'{self.base_model}' is already bitsandbytes-quantized. Pick the "
|
||||
f"family's dense (bf16) base repo for this mode, or use nf4/auto."
|
||||
)
|
||||
if self.mixed_precision != "bf16":
|
||||
raise ValueError(
|
||||
f"base_precision={base_precision!r} trains in bf16 compute; set "
|
||||
f"mixed_precision to bf16."
|
||||
)
|
||||
# learning_rate can arrive as a string ("1e-4") from the Studio config path, which
|
||||
# preserves it as a string after validation; coerce so AdamW receives a float.
|
||||
try:
|
||||
|
|
@ -298,6 +355,7 @@ class DiffusionLoraConfig:
|
|||
hf_token = token or None,
|
||||
cache_variants = int(self.cache_variants),
|
||||
compile_transformer = compile_transformer,
|
||||
base_precision = base_precision,
|
||||
resolved_family = resolved_family,
|
||||
)
|
||||
|
||||
|
|
@ -428,6 +486,18 @@ def _apply_perf_flags(
|
|||
torch.set_float32_matmul_precision("high")
|
||||
if cudnn_benchmark:
|
||||
torch.backends.cudnn.benchmark = True
|
||||
# The cuDNN SDPA backend's TRAINING graph is broken for the FLUX attention shapes
|
||||
# on torch 2.10 + cu130 (B200): mha_graph.execute fails, then poisons the context
|
||||
# into illegal memory accesses. Flash / mem-efficient SDPA are mathematically
|
||||
# equivalent, so pin those for the run (restored on exit).
|
||||
cuda_backends = getattr(torch.backends, "cuda", None)
|
||||
if cuda_backends is not None and hasattr(cuda_backends, "enable_cudnn_sdp"):
|
||||
try:
|
||||
snap["cudnn_sdp"] = bool(cuda_backends.cudnn_sdp_enabled())
|
||||
except Exception: # noqa: BLE001 -- flag unreadable: skip the tweak entirely
|
||||
snap["cudnn_sdp"] = None
|
||||
if snap["cudnn_sdp"]:
|
||||
cuda_backends.enable_cudnn_sdp(False)
|
||||
except Exception: # noqa: BLE001 -- perf flags are never fatal
|
||||
pass
|
||||
return snap
|
||||
|
|
@ -441,13 +511,15 @@ def _restore_perf_flags(snap: Optional[dict]) -> None:
|
|||
from core.inference.diffusion_speed import restore_backend_flags
|
||||
|
||||
restore_backend_flags(snap.get("flags"))
|
||||
if snap.get("matmul_precision"):
|
||||
try:
|
||||
import torch
|
||||
try:
|
||||
import torch
|
||||
|
||||
if snap.get("matmul_precision"):
|
||||
torch.set_float32_matmul_precision(snap["matmul_precision"])
|
||||
except Exception: # noqa: BLE001 -- best-effort restore
|
||||
pass
|
||||
if snap.get("cudnn_sdp"):
|
||||
torch.backends.cuda.enable_cudnn_sdp(True)
|
||||
except Exception: # noqa: BLE001 -- best-effort restore
|
||||
pass
|
||||
|
||||
|
||||
def _assert_trusted_base_model(base_model: str) -> None:
|
||||
|
|
|
|||
|
|
@ -727,6 +727,14 @@ 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(
|
||||
"nf4",
|
||||
description = (
|
||||
"DiT base transformer precision: nf4 QLoRA (memory floor, default), bf16 dense, "
|
||||
"int8 torchao weight-only, fp8 float8 training compute (Ada/Hopper/Blackwell), "
|
||||
"or auto (pick by free VRAM + GPU class). Dense modes need a non-prequant base."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class DiffusionTrainingStopRequest(BaseModel):
|
||||
|
|
@ -801,6 +809,12 @@ class DiffusionTrainableFamily(BaseModel):
|
|||
base_repos: List[str] = Field(default_factory = list)
|
||||
defaults: dict = Field(default_factory = dict)
|
||||
vram_note: str = ""
|
||||
# base_precision modes this machine supports for the family (empty = the family has no
|
||||
# precision selector, e.g. SDXL), plus the recommended pick and whether regional
|
||||
# torch.compile applies. Defaults keep older backends' payloads valid.
|
||||
precision_modes: List[str] = Field(default_factory = list)
|
||||
recommended_precision: str = "nf4"
|
||||
supports_compile: bool = False
|
||||
|
||||
|
||||
class DiffusionTrainingInfoResponse(BaseModel):
|
||||
|
|
|
|||
209
studio/backend/tests/test_diffusion_base_precision.py
Normal file
209
studio/backend/tests/test_diffusion_base_precision.py
Normal file
|
|
@ -0,0 +1,209 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""CPU-only unit tests for the DiT base_precision work.
|
||||
|
||||
Covers the new precision plumbing the precision PR adds: the ``base_precision``
|
||||
config validation (dense-vs-prequant + mixed-precision gating), the prequant-repo
|
||||
heuristic and its trainer alias, the pure ``auto`` precision policy table, the
|
||||
explicit-mode passthrough of ``_resolve_base_precision``, the fp8 module filter, the
|
||||
fp8 branch of the compile policy, the ``train_precision_modes`` machine probe, the
|
||||
family-info precision fields, and the request-model ``base_precision`` field. No GPU /
|
||||
model load: every helper here is pure or name-based, so the config validation runs on
|
||||
name matching (``resolve_trainable_family`` is offline) and the torch probe is monkeypatched.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch.nn as nn
|
||||
|
||||
import core.training.diffusion_train_common as common
|
||||
from core.training import diffusion_dit_trainer as dit
|
||||
from core.training.diffusion_train_common import (
|
||||
DiffusionLoraConfig,
|
||||
_config_from_dict,
|
||||
repo_is_prequantized,
|
||||
train_precision_modes,
|
||||
)
|
||||
from models.training import DiffusionTrainingStartRequest
|
||||
|
||||
# A dense (non-prequant) DiT base and a prequant bnb-4bit base. Both resolve a trainer
|
||||
# family from their names alone, so normalized() runs without a network call.
|
||||
_FLUX_DENSE = "black-forest-labs/FLUX.1-dev"
|
||||
_Z_PREQUANT = "unsloth/Z-Image-Turbo-unsloth-bnb-4bit"
|
||||
|
||||
|
||||
def _cfg(base_model = _FLUX_DENSE, **kw) -> DiffusionLoraConfig:
|
||||
return DiffusionLoraConfig(base_model = base_model, data_dir = "d", output_dir = "o", **kw)
|
||||
|
||||
|
||||
# ── base_precision validation ─────────────────────────────────────────────────
|
||||
def test_base_precision_validation():
|
||||
# Default normalizes to the nf4 memory floor.
|
||||
assert _cfg().normalized().base_precision == "nf4"
|
||||
|
||||
# An unknown mode is rejected by name.
|
||||
with pytest.raises(ValueError, match = "base_precision"):
|
||||
_cfg(base_precision = "banana").normalized()
|
||||
|
||||
# A dense mode is case/space-insensitive and stored lowered: " FP8 " on a dense base
|
||||
# with bf16 compute normalizes cleanly to "fp8".
|
||||
norm = _cfg(base_precision = " FP8 ", mixed_precision = "bf16").normalized()
|
||||
assert norm.base_precision == "fp8"
|
||||
|
||||
# A dense mode against a prequant (bnb-4bit) base is refused: the repo already ships a
|
||||
# 4-bit transformer and cannot serve the dense precisions.
|
||||
with pytest.raises(ValueError, match = "dense base repo"):
|
||||
_cfg(base_model = _Z_PREQUANT, base_precision = "bf16").normalized()
|
||||
|
||||
# A dense mode with non-bf16 compute is refused: these modes train in bf16 compute.
|
||||
with pytest.raises(ValueError, match = "bf16 compute"):
|
||||
_cfg(base_precision = "int8", mixed_precision = "fp16").normalized()
|
||||
|
||||
# "auto" is ACCEPTED by normalized() even on a prequant base: the concrete mode is
|
||||
# resolved at runtime against the live GPU, not at config validation.
|
||||
assert _cfg(base_model = _Z_PREQUANT, base_precision = "auto").normalized().base_precision == "auto"
|
||||
|
||||
|
||||
# ── repo_is_prequantized heuristic + trainer alias ────────────────────────────
|
||||
@pytest.mark.parametrize(
|
||||
"repo, expected",
|
||||
[
|
||||
("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", True),
|
||||
("some/model-4bit", True),
|
||||
("some/model-int4", True),
|
||||
("some/model-nf4", True),
|
||||
("black-forest-labs/FLUX.1-dev", False),
|
||||
("Tongyi-MAI/Z-Image-Turbo", False),
|
||||
],
|
||||
)
|
||||
def test_repo_is_prequantized_cases(repo, expected):
|
||||
assert repo_is_prequantized(repo) is expected
|
||||
|
||||
|
||||
def test_repo_is_prequantized_alias_is_same_object():
|
||||
# The trainer keeps a module-level alias for callers/tests; it must be the exact same
|
||||
# function object as the common heuristic (moved there for config validation).
|
||||
assert dit._repo_is_prequantized is repo_is_prequantized
|
||||
|
||||
|
||||
# ── _pick_auto_precision policy table (pure) ──────────────────────────────────
|
||||
def test_pick_auto_precision_policy_table():
|
||||
p = dit._pick_auto_precision
|
||||
|
||||
# A prequant base always resolves to nf4 (it can only serve 4-bit).
|
||||
assert p(True, "cuda", 140, 23.8, (10, 0), True) == "nf4"
|
||||
# No CUDA -> nf4 (the dense modes need a GPU).
|
||||
assert p(False, "cpu", 140, 23.8, (10, 0), True) == "nf4"
|
||||
# Missing free-VRAM number -> the safe nf4 mode.
|
||||
assert p(False, "cuda", None, 23.8, (10, 0), True) == "nf4"
|
||||
|
||||
# Plenty of free VRAM -> bf16 regardless of fp8 capability: compiled bf16 measured
|
||||
# FASTER than torchao float8 at LoRA-training shapes, so fp8 is opt-in only.
|
||||
assert p(False, "cuda", 140, 23.8, (10, 0), True) == "bf16"
|
||||
assert p(False, "cuda", 140, 23.8, (8, 0), True) == "bf16"
|
||||
assert p(False, "cuda", 140, 23.8, (10, 0), False) == "bf16"
|
||||
|
||||
# Middle band (25 > 23.8 * 0.55 * 1.5 = 19.6, but not > 23.8 * 1.5 = 35.7) -> int8.
|
||||
assert p(False, "cuda", 25, 23.8, (10, 0), True) == "int8"
|
||||
# Too little free VRAM for even int8 -> nf4.
|
||||
assert p(False, "cuda", 10, 23.8, (10, 0), True) == "nf4"
|
||||
|
||||
|
||||
# ── _resolve_base_precision passthrough ───────────────────────────────────────
|
||||
def test_resolve_base_precision_passes_explicit_through():
|
||||
# An explicit mode passes straight through without probing the GPU (normalized() already
|
||||
# validated it); the spec is only consulted for "auto".
|
||||
spec = dit._SPECS["flux.1"]
|
||||
cfg = _cfg(base_precision = "bf16")
|
||||
assert dit._resolve_base_precision(cfg, spec, "cuda") == "bf16"
|
||||
# Device is irrelevant for an explicit mode: still bf16 on cpu.
|
||||
assert dit._resolve_base_precision(cfg, spec, "cpu") == "bf16"
|
||||
|
||||
|
||||
# ── _fp8_module_filter ────────────────────────────────────────────────────────
|
||||
def test_fp8_module_filter():
|
||||
lin = nn.Linear(64, 64)
|
||||
# A plain feed-forward Linear with divisible dims gets float8 training compute.
|
||||
assert dit._fp8_module_filter(lin, "transformer_blocks.0.ff.net.0") is True
|
||||
# A LoRA-owned module is skipped (adapters stay high precision).
|
||||
assert dit._fp8_module_filter(lin, "transformer_blocks.0.attn.to_q.lora_A.default") is False
|
||||
# The output projection is skipped.
|
||||
assert dit._fp8_module_filter(lin, "proj_out") is False
|
||||
# An in_features not divisible by 16 is rejected (float8 kernels reject the shape).
|
||||
assert dit._fp8_module_filter(nn.Linear(30, 64), "transformer_blocks.0.ff.net.0") is False
|
||||
# A non-Linear module is never float8.
|
||||
assert dit._fp8_module_filter(nn.LayerNorm(64), "transformer_blocks.0.norm") is False
|
||||
|
||||
|
||||
# ── _should_compile fp8 branch ────────────────────────────────────────────────
|
||||
def test_should_compile_fp8_branch():
|
||||
# fp8 is only competitive compiled, so auto arms compile for it on a dense (non-bnb)
|
||||
# cuda base.
|
||||
cfg = _cfg(compile_transformer = "auto")
|
||||
assert dit._should_compile(cfg, False, "cuda", "fp8") is True
|
||||
# fp8 forces compile under auto even when the base is (hypothetically) reported as bnb.
|
||||
assert dit._should_compile(cfg, True, "cuda", "fp8") is True
|
||||
# An explicit "off" still wins over fp8 -- compile stays off.
|
||||
assert dit._should_compile(_cfg(compile_transformer = "off"), False, "cuda", "fp8") is False
|
||||
|
||||
|
||||
# ── train_precision_modes machine probe ───────────────────────────────────────
|
||||
def test_train_precision_modes_no_cuda(monkeypatch):
|
||||
# Patch the torch module attribute the function imports so it observes a CPU-only box:
|
||||
# no CUDA -> the nf4-only floor with nf4 recommended, and it never raises.
|
||||
import torch
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
||||
assert train_precision_modes() == (["nf4"], "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
|
||||
# (no precision selector) stays empty regardless of the probe.
|
||||
monkeypatch.setattr(common, "train_precision_modes", lambda: (["nf4", "bf16"], "auto"))
|
||||
infos = {i["name"]: i for i in common.family_train_infos()}
|
||||
|
||||
flux = infos["flux.1"]
|
||||
assert flux["precision_modes"] == ["nf4", "bf16"]
|
||||
assert flux["recommended_precision"] == "auto"
|
||||
assert flux["supports_compile"] is True
|
||||
|
||||
sdxl = infos["sdxl"]
|
||||
assert sdxl["precision_modes"] == []
|
||||
assert sdxl["recommended_precision"] == "nf4"
|
||||
assert sdxl["supports_compile"] is False
|
||||
|
||||
|
||||
# ── request model base_precision field ────────────────────────────────────────
|
||||
def test_request_model_base_precision():
|
||||
# The request defaults to the nf4 memory floor.
|
||||
req = DiffusionTrainingStartRequest(base_model = "x", data_dir = "d", output_dir = "o")
|
||||
assert req.base_precision == "nf4"
|
||||
|
||||
# An allowed dense mode is accepted.
|
||||
assert (
|
||||
DiffusionTrainingStartRequest(
|
||||
base_model = "x", data_dir = "d", output_dir = "o", base_precision = "fp8"
|
||||
).base_precision
|
||||
== "fp8"
|
||||
)
|
||||
|
||||
# An out-of-Literal value is rejected by pydantic.
|
||||
with pytest.raises(Exception):
|
||||
DiffusionTrainingStartRequest(
|
||||
base_model = "x", data_dir = "d", output_dir = "o", base_precision = "int4"
|
||||
)
|
||||
|
||||
# The generic Studio dict path carries base_precision through onto DiffusionLoraConfig.
|
||||
cfg = _config_from_dict(
|
||||
{
|
||||
"base_model": _FLUX_DENSE,
|
||||
"data_dir": "d",
|
||||
"output_dir": "o",
|
||||
"base_precision": "bf16",
|
||||
}
|
||||
)
|
||||
assert cfg.base_precision == "bf16"
|
||||
|
|
@ -194,8 +194,11 @@ def test_should_compile_policy():
|
|||
assert _should_compile(_cfg(compile_transformer = "on"), False, "cuda") is True
|
||||
# auto stays off over a bitsandbytes base (graph breaks in the dequant path).
|
||||
assert _should_compile(_cfg(compile_transformer = "auto"), True, "cuda") is False
|
||||
# auto turns on for a dense (non-bnb) base on cuda.
|
||||
assert _should_compile(_cfg(compile_transformer = "auto"), False, "cuda") is True
|
||||
# auto turns on for the dense bf16 base precision on cuda.
|
||||
assert (
|
||||
_should_compile(_cfg(compile_transformer = "auto"), False, "cuda", base_precision = "bf16")
|
||||
is True
|
||||
)
|
||||
# Any mode is a no-op on cpu.
|
||||
for mode in ("off", "on", "auto"):
|
||||
assert _should_compile(_cfg(compile_transformer = mode), False, "cpu") is False
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue