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

Merge diffusion-train-perf (pre-commit formatting + strict TF32 opt-out) into diffusion-train-precision

Browse source 
# Conflicts:
#	studio/backend/core/training/diffusion_train_common.py
This commit is contained in:
Daniel Han 2026-07-03 09:40:23 +00:00
commit f2c2ff9a2b
4 changed files with 98 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -367,7 +367,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 +495,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 +601,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 +774,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
@ -831,7 +846,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 +936,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 +992,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 +1118,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 +1195,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

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

@ -434,11 +434,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 +459,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 +483,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 +521,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

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