ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge diffusion-train-precision (review fixes + formatting) into diffusion-train-tab-2

Browse source
This commit is contained in:
Daniel Han 2026-07-03 09:49:23 +00:00
commit 23f55bbcf4
5 changed files with 155 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

92
studio/backend/core/training/diffusion_dit_trainer.py
View file

@ -267,7 +267,6 @@ def _apply_fp8_training(transformer, on_event) -> bool:
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,
@ -285,25 +284,41 @@ def _pick_auto_precision(prequant, device, free_gb, dense_gb, capability, has_fp
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."""
the same hardware. int8 must still materialise the full bf16 transformer before
``quantize_`` shrinks it module-by-module, so its band requires the dense-load
transient (1.15x dense) to fit -- what int8 buys in that band is steady-state
headroom for activations and the latent cache, not load-time memory.
``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:
if free_gb > dense_gb * 1.5:
return "bf16"
if free_gb > dense_gb * 0.55 * headroom:
if free_gb > dense_gb * 1.15:
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)."""
explicit modes pass through (normalized() already validated them against the repo and
compute dtype) but are re-checked against the live device here: the dense modes are
CUDA-only, and /info never advertises them on a host without a GPU, so an explicit
request from a stale or direct client fails fast instead of loading a full dense
transformer onto the CPU."""
mode = (cfg.base_precision or "nf4").strip().lower()
if mode != "auto":
if mode in ("bf16", "int8", "fp8") 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'."
)
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.
if getattr(cfg, "mixed_precision", "bf16") != "bf16":
return "nf4"
prequant = repo_is_prequantized(cfg.base_model)
free_gb = None
capability = None
@ -367,7 +382,12 @@ def _flux_encode_latent_stats(vae, pixel_values):
return (dist.mean - vae.config.shift_factor) * scale, dist.std * scale
def _flux_collate(entries, device, weight_dtype, pad_to = None):
def _flux_collate(
entries,
device,
weight_dtype,
pad_to = None,
):
import torch
# FLUX embeds are fixed-length (encode_prompt pads to max_sequence_length), so a plain
@ -490,7 +510,12 @@ def _qwen_encode_latent_stats(vae, pixel_values):
return (dist.mean - mean) / std, dist.std / std
def _qwen_collate(entries, device, weight_dtype, pad_to = None):
def _qwen_collate(
entries,
device,
weight_dtype,
pad_to = None,
):
import torch
import torch.nn.functional as F
@ -591,7 +616,12 @@ def _zimage_encode_latent_stats(vae, pixel_values):
return _zimage_encode_latents(vae, pixel_values), None
def _zimage_collate(entries, device, weight_dtype, pad_to = None):
def _zimage_collate(
entries,
device,
weight_dtype,
pad_to = None,
):
caps = [e[0].to(device = device, dtype = weight_dtype) for e in entries]
return (caps,)
@ -759,7 +789,7 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev
total = len(image_paths)
for i, path in enumerate(image_paths):
variants = []
for (u_left, u_top, flip) in plan[i]:
for u_left, u_top, flip in plan[i]:
px = (
_load_pixel_tensor_planned(
path, cfg.resolution, cfg.center_crop, u_left, u_top, flip
@ -796,7 +826,12 @@ 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, base_precision = "nf4") -> 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
@ -813,7 +848,12 @@ def _should_compile(cfg, base_is_bnb, device, base_precision = "nf4") -> bool:
def _maybe_compile_transformer(
transformer, cfg, base_is_bnb, device, on_event, base_precision = "nf4"
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
@ -831,7 +871,9 @@ def _maybe_compile_transformer(
fn = getattr(transformer, "compile_repeated_blocks", None)
if not callable(fn):
_emit(on_event, "warning", message = "torch.compile unavailable for this model; running eager.")
_emit(
on_event, "warning", message = "torch.compile unavailable for this model; running eager."
)
return False
try:
dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None)
@ -919,7 +961,14 @@ def run_dit_lora_training(
perf_snap = _apply_perf_flags(cfg, device)
try:
return _train_dit(
cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_stop,
cfg,
spec,
pairs,
rng,
device,
weight_dtype,
on_event,
_check_stop,
lambda: save_on_stop,
)
finally:
@ -968,8 +1017,12 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
)
if latent_cache is None: # stopped during the cache build; nothing trained yet
_emit(
on_event, "complete", output_dir = str(out_dir), lora_path = None,
stopped = True, steps_run = 0,
on_event,
"complete",
output_dir = str(out_dir),
lora_path = None,
stopped = True,
steps_run = 0,
)
return str(out_dir)
try:
@ -1090,7 +1143,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
noisy = (1.0 - sigmas) * latents + sigmas * noise
embeds = spec.collate(
[caption_embeds[captions[i]] for i in idxs], device, weight_dtype,
[caption_embeds[captions[i]] for i in idxs],
device,
weight_dtype,
pad_to = qwen_pad_to,
)
with autocast:
@ -1165,6 +1220,7 @@ def _make_optimizer(params, lr):
regression for LoRA -- else torch AdamW, fused on CUDA (with a fallback when this
build/device lacks the fused kernel)."""
import torch
try:
import bitsandbytes as bnb
return bnb.optim.AdamW8bit(params, lr = lr)

10
studio/backend/core/training/diffusion_lora_trainer.py
View file

@ -201,7 +201,7 @@ def _build_sdxl_latent_cache(
total = len(image_paths)
for i, path in enumerate(image_paths):
variants = []
for (u_left, u_top, flip) in plan[i]:
for u_left, u_top, flip in plan[i]:
tensor, time_ids = _load_image_tensor_planned(
path, cfg.resolution, cfg.center_crop, u_left, u_top, flip
)
@ -390,7 +390,13 @@ def run_diffusion_lora_training(
latent_cache = None
if use_cache:
latent_cache = _build_sdxl_latent_cache(
vae, vae_scale, [p for p, _ in pairs], cfg, device, weight_dtype, on_event,
vae,
vae_scale,
[p for p, _ in pairs],
cfg,
device,
weight_dtype,
on_event,
_check_stop,
)
if latent_cache is None: # stopped during the cache build; nothing trained yet

31
studio/backend/core/training/diffusion_train_common.py
View file

@ -140,7 +140,6 @@ def train_precision_modes() -> tuple[list[str], str]:
recommended = "nf4"
try:
import torch
if torch.cuda.is_available():
modes += ["bf16", "int8"]
major, minor = torch.cuda.get_device_capability()
@ -434,11 +433,7 @@ def _emit(on_event: Optional[EventCb], type_: str, **kw: Any) -> None:
def _plan_cache_variants(
num_images: int,
cache_variants: int,
center_crop: bool,
random_flip: bool,
seed: int,
num_images: int, cache_variants: int, center_crop: bool, random_flip: bool, seed: int
) -> list[list[tuple[float, float, bool]]]:
"""Seed-deterministic crop/flip plan for the latent cache: per image, up to
``cache_variants`` draws of (u_left, u_top, flip) with the crop as unit fractions the
@ -463,10 +458,13 @@ def _plan_cache_variants(
def _apply_perf_flags(
cfg: "DiffusionLoraConfig", device: str, cudnn_benchmark: bool = False
cfg: "DiffusionLoraConfig",
device: str,
cudnn_benchmark: bool = False,
) -> dict:
"""Set the run-scoped torch backend knobs: TF32 matmuls + high fp32 matmul precision
(under ``cfg.enable_tf32``), plus cudnn autotuning when the caller opts in. Autotune is
when ``cfg.enable_tf32`` is on, strict fp32 (all TF32 flags cleared) when it is off,
plus cudnn autotuning when the caller opts in. Autotune is
for the conv-heavy SDXL U-Net only: measured on B200, it DOUBLES peak VRAM (fp32 VAE
conv workspaces) while the DiT loop -- pure matmuls once the latent cache is built --
gains nothing from it. Returns a snapshot for ``_restore_perf_flags``. Best-effort:
@ -484,6 +482,12 @@ def _apply_perf_flags(
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision("high")
else:
# The opt-out is a strict-fp32 A/B mode, so actively clear the flags rather
# than inherit ambient state (cudnn TF32 defaults to ON in torch).
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
torch.set_float32_matmul_precision("highest")
if cudnn_benchmark:
torch.backends.cudnn.benchmark = True
# The cuDNN SDPA backend's TRAINING graph is broken for the FLUX attention shapes
@ -516,8 +520,15 @@ def _restore_perf_flags(snap: Optional[dict]) -> None:
if snap.get("matmul_precision"):
torch.set_float32_matmul_precision(snap["matmul_precision"])
if snap.get("cudnn_sdp"):
torch.backends.cuda.enable_cudnn_sdp(True)
# Restore the exact pre-run cudnn SDPA state; None means the flag was unreadable
# (or absent) at apply time and was never touched.
cuda_backends = getattr(torch.backends, "cuda", None)
if (
snap.get("cudnn_sdp") is not None
and cuda_backends is not None
and hasattr(cuda_backends, "enable_cudnn_sdp")
):
cuda_backends.enable_cudnn_sdp(bool(snap["cudnn_sdp"]))
except Exception: # noqa: BLE001 -- best-effort restore
pass

29
studio/backend/tests/test_diffusion_base_precision.py
View file

@ -105,9 +105,13 @@ def test_pick_auto_precision_policy_table():
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.
# Middle band (30 > 23.8 * 1.15 = 27.4, but not > 23.8 * 1.5 = 35.7) -> int8.
assert p(False, "cuda", 30, 23.8, (10, 0), True) == "int8"
# int8 still materialises the full bf16 transformer before quantize_ shrinks it, so
# free VRAM below the dense-load transient (25 < 27.4) must fall back to nf4 even
# though the QUANTIZED weights would have fit.
assert p(False, "cuda", 25, 23.8, (10, 0), True) == "nf4"
# Too little free VRAM for any dense load -> nf4.
assert p(False, "cuda", 10, 23.8, (10, 0), True) == "nf4"
@ -118,8 +122,22 @@ def test_resolve_base_precision_passes_explicit_through():
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"
# The dense modes are CUDA-only: an explicit request on a GPU-less host fails fast
# (before any model load) instead of silently proceeding; /info never advertised it.
with pytest.raises(ValueError, match = "CUDA"):
dit._resolve_base_precision(cfg, spec, "cpu")
# nf4 stays a passthrough on any device (the bnb load path owns its own errors).
assert dit._resolve_base_precision(_cfg(base_precision = "nf4"), spec, "cpu") == "nf4"
def test_resolve_auto_requires_bf16_compute():
# auto may resolve to bf16/int8 which train in bf16 compute, so a non-bf16
# mixed_precision pins auto to the nf4 floor BEFORE any GPU probe (pure, no CUDA
# needed here) -- mirroring the normalized() rule for explicit dense modes.
spec = dit._SPECS["flux.1"]
cfg = _cfg(base_precision = "auto", mixed_precision = "fp16")
assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4"
# ── _fp8_module_filter ────────────────────────────────────────────────────────
@ -154,7 +172,6 @@ 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")

35
studio/backend/tests/test_diffusion_train_perf.py
View file

@ -82,9 +82,7 @@ def test_plan_cache_variants_deterministic_and_deduped():
# ── per-family collate fns ────────────────────────────────────────────────────
def test_flux_collate_shapes():
# FLUX embeds are fixed length: 3 entries batch by a plain cat; text_ids are shared.
entries = [
(torch.randn(1, 512, 32), torch.randn(1, 16), torch.randn(512, 3)) for _ in range(3)
]
entries = [(torch.randn(1, 512, 32), torch.randn(1, 16), torch.randn(512, 3)) for _ in range(3)]
pe, pooled, text_ids = _flux_collate(entries, "cpu", torch.float32)
assert pe.shape == (3, 512, 32)
assert pooled.shape == (3, 16)
@ -242,9 +240,7 @@ def test_service_stop_save_flag():
# ── preparing / warning events + stopped completion messages ──────────────────
def test_apply_event_preparing_and_warning():
svc = DiffusionTrainingService()
svc._apply_event(
{"type": "preparing", "stage": "cache_latents", "done": 4, "total": 8}
)
svc._apply_event({"type": "preparing", "stage": "cache_latents", "done": 4, "total": 8})
st = svc.status()
assert st["status"] == "running"
assert st["in_model_load"] is True
@ -335,3 +331,30 @@ def test_perf_flags_cpu_roundtrip():
snap = _apply_perf_flags(_cfg(), "cpu")
assert isinstance(snap, dict)
_restore_perf_flags(snap) # no exception
def test_perf_flags_tf32_off_clears_flags():
# enable_tf32=False is the strict-fp32 A/B mode: it must actively clear the TF32 flags
# (cudnn TF32 defaults ON in torch) rather than inherit ambient state, and restore must
# put the ambient values back. The flag attributes are plain Python state, present and
# settable on CPU-only torch builds, so this runs without a GPU.
import torch
before = (
torch.backends.cuda.matmul.allow_tf32,
torch.backends.cudnn.allow_tf32,
torch.get_float32_matmul_precision(),
)
snap = _apply_perf_flags(_cfg(enable_tf32 = False), "cuda")
try:
assert torch.backends.cuda.matmul.allow_tf32 is False
assert torch.backends.cudnn.allow_tf32 is False
assert torch.get_float32_matmul_precision() == "highest"
finally:
_restore_perf_flags(snap)
after = (
torch.backends.cuda.matmul.allow_tf32,
torch.backends.cudnn.allow_tf32,
torch.get_float32_matmul_precision(),
)
assert after == before