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

Speed up diffusion LoRA training and cut DiT peak VRAM by a third

Browse source 
Perf core for the diffusion trainers, defaults preserving the training math:

- Phased model loading: the pipeline now loads without its transformer
  (conditioning only), captions are encoded and the text encoders freed,
  the VAE latent cache is built and the VAE freed, and only then does the
  transformer load. The multi-GB denoiser never shares VRAM with the
  encoders, cutting measured peak VRAM on B200: FLUX 17.1 -> 10.4 GB,
  Qwen-Image 19.1 -> 12.8 GB, Z-Image 7.3 -> 4.7 GB.
- Latent cache (cache_latents, default on): per-image crop/flip variants
  (cache_variants, default 4 vs the single frozen variant of the diffusers
  --cache_latents) store the VAE posterior's affine parameters, so every
  step still draws a fresh VAE sample; a cached center-crop Z-Image run
  matches the uncached one at the bf16 nondeterminism floor.
- True batching: train_batch_size now actually batches the transformer
  forward (it was silently 1). nf4 dequant dominates the step cost, so
  batch 4 lands near batch-1 step time: 4.0x samples/s on Qwen-Image,
  3.1x on FLUX, 2.1x on Z-Image, with multi-seed loss envelopes
  overlapping batch-1.
- LR scheduler support in the DiT loop (lr_scheduler / lr_warmup_steps
  were accepted but ignored); progress events now report the real
  per-step LR.
- TF32 + high fp32 matmul precision under enable_tf32 (default on),
  snapshot/restored around the run. cudnn.benchmark is scoped to a
  caller opt-in only: autotuning the fp32 VAE convs doubled peak VRAM
  on the DiT families for zero steady-state gain.
- Vectorized sigma gathering (drops a per-step Python search loop),
  cached FLUX img_ids/guidance, fused torch AdamW fallback, steady-state
  samples_per_second (excludes the first-step warmup).
- Regional torch.compile plumbing (compile_transformer off/on/auto with
  eager fallback): auto stays off over a bitsandbytes base where compile
  is a net loss (27 s warmup, slightly slower steady on Z-Image); it
  arms automatically for the dense/quantized speed modes that follow.
- Stop parity with the LLM trainer: /api/train/diffusion/stop accepts an
  optional {save} body and the service forwards save=False as a
  no-save cancel; a new preparing event surfaces cache-build progress.
- SDXL trainer gets the same latent cache, perf flags, and fused
  fallback; its batching, LR schedule, and min-SNR stay as they were.

Verified: 83 backend tests green; per-family 30-40 step runs with
adapter round-trip generation through the normal LoRA path (FLUX,
Qwen-Image, Z-Image all pass).
This commit is contained in:
Daniel Han 2026-07-03 08:44:57 +00:00
commit 9d967d2c5b
8 changed files with 1389 additions and 375 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -10,19 +10,24 @@ the diffusers dreambooth scripts, form ``noisy = (1 - sigma) * latents + sigma *
predict the velocity, and regress it onto ``target = noise - latents``.
The per-family differences (latent normalisation + packing, the transformer forward
signature, and the LoRA save entrypoint) live in small ``_FamilySpec`` objects; the loop
itself is family-agnostic. Verified against diffusers 0.38.0.
signature, embedding collation, and the LoRA save entrypoint) live in small ``_FamilySpec``
objects; the loop itself is family-agnostic. Verified against diffusers 0.38.0.
Memory: the text encoder(s) are the largest module (T5-XXL ~9 GB for FLUX, Qwen2.5-VL ~7 GB
for Qwen-Image, Qwen3 for Z-Image), so captions are encoded ONCE up front and the encoders
are freed before the loop. The transformer trains as a QLoRA (nf4) adapter by default with
gradient checkpointing and 8-bit AdamW, so only the (small) LoRA params + optimizer state
and the frozen 4-bit base sit in VRAM during the loop.
are freed before the loop. VAE latents are likewise precomputed into a small CPU cache
(``cache_latents``) and the VAE freed: the cache stores the posterior's affine parameters
(mean/std folded through the family's latent normalisation), 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. The transformer trains as a QLoRA (nf4) adapter by
default with gradient checkpointing and 8-bit AdamW, so only the (small) LoRA params +
optimizer state and the frozen 4-bit base sit in VRAM during the loop.
"""
from __future__ import annotations
import gc
import os
import random
import time
from contextlib import nullcontext
@ -35,9 +40,12 @@ from core.training.diffusion_train_common import (
DiffusionLoraConfig,
EventCb,
StopCb,
_apply_perf_flags,
_assert_trusted_base_model,
_emit,
_plan_cache_variants,
_publish_to_lora_catalog,
_restore_perf_flags,
discover_image_caption_pairs,
)
@ -66,14 +74,24 @@ class _FamilySpec:
lora_targets: tuple[str, ...]
# bf16 only (Z-Image overflows fp16 and its RoPE/embedder run in fp32).
force_bf16: bool
# Builds (pipe, transformer, vae) with the transformer loaded as a trainable nf4 QLoRA
# when qlora=True. Returns the pipeline (for save_lora_weights + encode_prompt), the
# transformer to attach LoRA to, and the VAE (kept resident for latent encoding).
load: Callable[..., tuple[Any, Any, Any]]
# 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
# transformer alone (as a trainable nf4 QLoRA when qlora=True) once the conditioning
# modules are freed. Roughly halves peak VRAM for the big DiTs.
load_conditioners: Callable[..., tuple[Any, Any]]
load_transformer: Callable[..., Any]
# Encode a list of captions -> a per-caption tuple of CPU tensors (the family's embeds).
encode_prompts: Callable[..., list[tuple]]
# Encode a pixel tensor [B,3,H,W] in [-1,1] -> latents (family-normalised, on device).
encode_latents: Callable[..., Any]
# Encode a pixel tensor -> (A, B) affine posterior parameters so a per-step sample is
# A + B * randn (family normalisation folded in). B is None for a deterministic
# (mode-based) family. Used by the latent cache.
encode_latent_stats: Callable[..., tuple]
# Collate a list of per-caption embed tuples -> one batched tuple on device. ``pad_to``
# pins a fixed text length for families with variable-length embeds (compile).
collate: Callable[..., tuple]
# One transformer forward: (transformer, noisy, timesteps, sigmas, embeds_batch, cfg,
# device, weight_dtype) -> model_pred aligned with target = noise - latents.
forward: Callable[..., Any]
@ -82,15 +100,12 @@ class _FamilySpec:
# ── shared flow-matching helpers ──────────────────────────────────────────────
def _get_sigmas(scheduler, timesteps, device, dtype, n_dim):
"""Gather per-sample sigmas for ``timesteps`` and broadcast to ``n_dim`` (matches the
diffusers dreambooth get_sigmas helper)."""
import torch
sigmas = scheduler.sigmas.to(device = device, dtype = dtype)
schedule_timesteps = scheduler.timesteps.to(device)
step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
sigma = sigmas[step_indices].flatten()
def _gather_sigmas(scheduler, indices, device, dtype, n_dim):
"""Gather per-sample sigmas for schedule ``indices`` and broadcast to ``n_dim``.
Index-based (no per-item search): ``indices`` are the positions ``_sample_timesteps``
drew from ``scheduler.timesteps``, and ``scheduler.sigmas`` is aligned with it, so this
returns exactly what the diffusers ``get_sigmas`` timestep-matching helper would."""
sigma = scheduler.sigmas[indices].to(device = device, dtype = dtype).flatten()
while sigma.ndim < n_dim:
sigma = sigma.unsqueeze(-1)
return sigma
@ -99,7 +114,6 @@ def _get_sigmas(scheduler, timesteps, device, dtype, n_dim):
def _sample_timesteps(scheduler, batch_size, device):
"""Logit-normal density timestep sampling (weighting_scheme='logit_normal'), returning
(timesteps, indices) into the scheduler's schedule."""
import torch
from diffusers.training_utils import compute_density_for_timestep_sampling
u = compute_density_for_timestep_sampling(
@ -112,7 +126,7 @@ def _sample_timesteps(scheduler, batch_size, device):
num_train = scheduler.config.num_train_timesteps
indices = (u * num_train).long().clamp(0, num_train - 1)
timesteps = scheduler.timesteps.to(device)[indices].to(device)
return timesteps
return timesteps, indices
def _encoders_to_device(pipe, device) -> None:
@ -161,32 +175,51 @@ def _load_quantized_transformer(transformer_cls, cfg):
)
# ── FLUX.1-dev ────────────────────────────────────────────────────────────────
def _flux_load(cfg, device, weight_dtype, qlora):
def _load_pipe_without_transformer(pipe_cls, cfg, device):
"""Load a pipeline for conditioning only: ``transformer = None`` skips the multi-GB
denoiser entirely (the documented diffusers pattern), leaving just the text encoders +
tokenizers + VAE + scheduler. The transformer loads later, after these are freed."""
import torch
from diffusers import FluxPipeline, FluxTransformer2DModel
if qlora:
transformer = FluxTransformer2DModel.from_pretrained(
cfg.base_model,
subfolder = "transformer",
quantization_config = _bnb_4bit_config(),
torch_dtype = torch.bfloat16,
token = cfg.hf_token,
)
pipe = FluxPipeline.from_pretrained(
cfg.base_model,
transformer = transformer,
torch_dtype = torch.bfloat16,
token = cfg.hf_token,
)
else:
pipe = FluxPipeline.from_pretrained(
cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token
)
transformer = pipe.transformer
pipe = pipe_cls.from_pretrained(
cfg.base_model,
transformer = None,
torch_dtype = torch.bfloat16,
token = cfg.hf_token,
)
pipe.vae.to(device, dtype = torch.float32)
return pipe, transformer, pipe.vae
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."""
import torch
if qlora and 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 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
# ── FLUX.1-dev ────────────────────────────────────────────────────────────────
def _flux_load_conditioners(cfg, device, weight_dtype):
from diffusers import FluxPipeline
return _load_pipe_without_transformer(FluxPipeline, cfg, device)
def _flux_load_transformer(cfg, device, weight_dtype, qlora):
from diffusers import FluxTransformer2DModel
return _load_dit_transformer(FluxTransformer2DModel, cfg, device, qlora)
def _flux_encode_prompts(pipe, captions, device):
@ -216,24 +249,61 @@ def _flux_encode_latents(vae, pixel_values):
return lat
def _flux_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype):
def _flux_encode_latent_stats(vae, pixel_values):
import torch
with torch.no_grad():
dist = vae.encode(pixel_values.to(torch.float32)).latent_dist
scale = vae.config.scaling_factor
return (dist.mean - vae.config.shift_factor) * scale, dist.std * scale
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
# cat batches them; text_ids are shared position ids, identical across prompts.
pe = torch.cat([e[0] for e in entries]).to(device = device, dtype = weight_dtype)
pooled = torch.cat([e[1] for e in entries]).to(device = device, dtype = weight_dtype)
text_ids = entries[0][2].to(device = device, dtype = torch.float32)
return (pe, pooled, text_ids)
# Per-run cache of the step-invariant FLUX conditioning tensors (RoPE image ids + the
# guidance vector): their shapes are fixed once resolution/batch are, so rebuilding them
# every step is pure allocator churn. Cleared at run start (subprocess-local anyway).
_FLUX_STATIC: dict[tuple, tuple] = {}
def _flux_static_inputs(bsz, h, w, device):
import torch
from diffusers import FluxPipeline
key = (bsz, h, w, str(device))
hit = _FLUX_STATIC.get(key)
if hit is None:
# Position ids drive RoPE and are indices, not activations -- keep them float32 (the
# dtype diffusers' own pipeline builds) regardless of the bf16 training dtype.
img_ids = FluxPipeline._prepare_latent_image_ids(bsz, h // 2, w // 2, device, torch.float32)
guidance = torch.full((bsz,), 1.0, device = device, dtype = torch.float32)
hit = _FLUX_STATIC[key] = (img_ids, guidance)
return hit
def _flux_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype):
from diffusers import FluxPipeline
pe, pooled, text_ids = embeds_batch
bsz, c, h, w = noisy.shape
packed = FluxPipeline._pack_latents(noisy, bsz, c, h, w)
# Position ids drive RoPE and are indices, not activations -- keep them float32 (the
# dtype diffusers' own pipeline builds) regardless of the bf16 training dtype.
img_ids = FluxPipeline._prepare_latent_image_ids(bsz, h // 2, w // 2, device, torch.float32)
guidance = torch.full((bsz,), 1.0, device = device, dtype = torch.float32)
img_ids, guidance = _flux_static_inputs(bsz, h, w, device)
model_pred = transformer(
hidden_states = packed,
timestep = timesteps / 1000,
guidance = guidance,
pooled_projections = pooled.to(weight_dtype),
encoder_hidden_states = pe.to(weight_dtype),
txt_ids = text_ids.to(torch.float32),
pooled_projections = pooled,
encoder_hidden_states = pe,
txt_ids = text_ids,
img_ids = img_ids,
return_dict = False,
)[0]
@ -250,19 +320,17 @@ def _flux_save(pipe_cls, out_dir, transformer_lora_layers):
# ── Qwen-Image ────────────────────────────────────────────────────────────────
def _qwen_load(cfg, device, weight_dtype, qlora):
import torch
from diffusers import QwenImagePipeline, QwenImageTransformer2DModel
def _qwen_load_conditioners(cfg, device, weight_dtype):
from diffusers import QwenImagePipeline
return _load_pipe_without_transformer(QwenImagePipeline, cfg, device)
def _qwen_load_transformer(cfg, device, weight_dtype, qlora):
# The prequant default (unsloth/Qwen-Image-2512-unsloth-bnb-4bit) ships the transformer
# 4-bit, so from_pretrained loads it trainable as-is. A dense (bf16) base -- the 20B
# Qwen/Qwen-Image -- is quantized to nf4 on the fly so QLoRA still fits.
kwargs = {"torch_dtype": torch.bfloat16, "token": cfg.hf_token}
if qlora and not _repo_is_prequantized(cfg.base_model):
kwargs["transformer"] = _load_quantized_transformer(QwenImageTransformer2DModel, cfg)
pipe = QwenImagePipeline.from_pretrained(cfg.base_model, **kwargs)
pipe.vae.to(device, dtype = torch.float32)
return pipe, pipe.transformer, pipe.vae
# 4-bit and loads trainable as-is; the dense 20B Qwen/Qwen-Image is quantized to nf4 on
# the fly so QLoRA still fits.
from diffusers import QwenImageTransformer2DModel
return _load_dit_transformer(QwenImageTransformer2DModel, cfg, device, qlora)
def _qwen_encode_prompts(pipe, captions, device):
@ -282,6 +350,15 @@ def _qwen_encode_prompts(pipe, captions, device):
return out
def _qwen_latent_affine(vae, ref):
import torch
z = vae.config.z_dim
mean = torch.tensor(vae.config.latents_mean, device = ref.device, dtype = ref.dtype)
std = torch.tensor(vae.config.latents_std, device = ref.device, dtype = ref.dtype)
return mean.view(1, z, 1, 1, 1), std.view(1, z, 1, 1, 1)
def _qwen_encode_latents(vae, pixel_values):
import torch
@ -290,16 +367,48 @@ def _qwen_encode_latents(vae, pixel_values):
px = pixel_values.to(torch.float32).unsqueeze(2) # [B,3,1,H,W]
with torch.no_grad():
lat = vae.encode(px).latent_dist.sample() # [B,16,1,h,w]
z = vae.config.z_dim
mean = torch.tensor(vae.config.latents_mean, device = lat.device, dtype = lat.dtype)
std = torch.tensor(vae.config.latents_std, device = lat.device, dtype = lat.dtype)
mean = mean.view(1, z, 1, 1, 1)
std = std.view(1, z, 1, 1, 1)
mean, std = _qwen_latent_affine(vae, lat)
return (lat - mean) / std
def _qwen_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype):
def _qwen_encode_latent_stats(vae, pixel_values):
import torch
px = pixel_values.to(torch.float32).unsqueeze(2)
with torch.no_grad():
dist = vae.encode(px).latent_dist
mean, std = _qwen_latent_affine(vae, dist.mean)
return (dist.mean - mean) / std, dist.std / std
def _qwen_collate(entries, device, weight_dtype, pad_to = None):
import torch
import torch.nn.functional as F
# Qwen embeds are variable-length: pad to the batch max (or a pinned ``pad_to`` bucket
# under compile so the graph shape stays fixed) and batch the validity mask with them.
seqs = [e[0].shape[1] for e in entries]
target = max(pad_to or 0, max(seqs))
pes, masks = [], []
for pe, mask in entries:
s = pe.shape[1]
if mask is None:
mask = torch.ones((1, s), dtype = torch.int64)
if s < target:
pe = F.pad(pe, (0, 0, 0, target - s))
mask = F.pad(mask, (0, target - s))
pes.append(pe)
masks.append(mask)
pe_b = torch.cat(pes).to(device = device, dtype = weight_dtype)
mask_b = torch.cat(masks).to(device)
# A single unpadded sample keeps the legacy None mask (identical math; avoids any
# behaviour delta for existing single-image runs whose pipeline returned None).
if len(entries) == 1 and entries[0][1] is None and target == seqs[0]:
mask_b = None
return (pe_b, mask_b)
def _qwen_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype):
from diffusers import QwenImagePipeline
pe, mask = embeds_batch
@ -310,8 +419,8 @@ def _qwen_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, devi
img_shapes = [[(1, h // 2, w // 2)]] * bsz
pred = transformer(
hidden_states = packed,
encoder_hidden_states = pe.to(weight_dtype),
encoder_hidden_states_mask = mask.to(device) if mask is not None else None,
encoder_hidden_states = pe,
encoder_hidden_states_mask = mask,
timestep = timesteps / 1000,
img_shapes = img_shapes,
return_dict = False,
@ -329,18 +438,16 @@ def _qwen_save(pipe_cls, out_dir, transformer_lora_layers):
# ── Z-Image ───────────────────────────────────────────────────────────────────
def _zimage_load(cfg, device, weight_dtype, qlora):
import torch
from diffusers import ZImagePipeline, ZImageTransformer2DModel
def _zimage_load_conditioners(cfg, device, weight_dtype):
from diffusers import ZImagePipeline
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).
kwargs = {"torch_dtype": torch.bfloat16, "token": cfg.hf_token}
if qlora and not _repo_is_prequantized(cfg.base_model):
kwargs["transformer"] = _load_quantized_transformer(ZImageTransformer2DModel, cfg)
pipe = ZImagePipeline.from_pretrained(cfg.base_model, **kwargs)
pipe.vae.to(device, dtype = torch.float32)
return pipe, pipe.transformer, pipe.vae
from diffusers import ZImageTransformer2DModel
return _load_dit_transformer(ZImageTransformer2DModel, cfg, device, qlora)
def _zimage_encode_prompts(pipe, captions, device):
@ -369,16 +476,26 @@ def _zimage_encode_latents(vae, pixel_values):
return (lat - vae.config.shift_factor) * vae.config.scaling_factor
def _zimage_encode_latent_stats(vae, pixel_values):
# Z-Image trains from the posterior mode (deterministic), so the cached entry is the
# final latent itself: B is None and the loop skips the per-step sampling draw.
return _zimage_encode_latents(vae, pixel_values), 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,)
def _zimage_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype):
import torch
(emb,) = embeds_batch
(caps,) = embeds_batch
# List I/O: one [C,1,H,W] latent + one [seq,2560] caption per sample. The timestep
# convention is REVERSED ((1000 - t) / 1000) and the prediction is NEGATED.
x_list = list(noisy.unsqueeze(2).unbind(dim = 0))
cap_list = [emb.to(device = device, dtype = weight_dtype)]
t_norm = (1000 - timesteps) / 1000
out = transformer(x_list, t_norm, cap_list, return_dict = False)[0]
out = transformer(x_list, t_norm, list(caps), return_dict = False)[0]
return -torch.stack(out, dim = 0).squeeze(2)
@ -396,9 +513,12 @@ _SPECS: dict[str, _FamilySpec] = {
family = "flux.1",
lora_targets = _FLUX_TARGETS,
force_bf16 = False,
load = _flux_load,
load_conditioners = _flux_load_conditioners,
load_transformer = _flux_load_transformer,
encode_prompts = _flux_encode_prompts,
encode_latents = _flux_encode_latents,
encode_latent_stats = _flux_encode_latent_stats,
collate = _flux_collate,
forward = _flux_forward,
save = _flux_save,
),
@ -406,9 +526,12 @@ _SPECS: dict[str, _FamilySpec] = {
family = "qwen-image",
lora_targets = _QWEN_TARGETS,
force_bf16 = True,
load = _qwen_load,
load_conditioners = _qwen_load_conditioners,
load_transformer = _qwen_load_transformer,
encode_prompts = _qwen_encode_prompts,
encode_latents = _qwen_encode_latents,
encode_latent_stats = _qwen_encode_latent_stats,
collate = _qwen_collate,
forward = _qwen_forward,
save = _qwen_save,
),
@ -416,9 +539,12 @@ _SPECS: dict[str, _FamilySpec] = {
family = "z-image",
lora_targets = _ZIMAGE_TARGETS,
force_bf16 = True,
load = _zimage_load,
load_conditioners = _zimage_load_conditioners,
load_transformer = _zimage_load_transformer,
encode_prompts = _zimage_encode_prompts,
encode_latents = _zimage_encode_latents,
encode_latent_stats = _zimage_encode_latent_stats,
collate = _zimage_collate,
forward = _zimage_forward,
save = _zimage_save,
),
@ -441,18 +567,32 @@ def _assert_gated_access(base_model: str, hf_token: Optional[str]) -> None:
)
def _load_pixel_tensor(path, resolution, center_crop, random_flip, rng):
"""Load an image -> a normalised [3,H,W] tensor in [-1,1]. Same geometry as the SDXL
loader but without the SDXL time-ids (DiT families don't use them)."""
import numpy as np
import torch
def _open_resized(path, resolution):
"""Open + EXIF-orient + short-side resize to ``resolution`` (same geometry as the SDXL
loader). Returns the resized PIL image and its (rw, rh)."""
from PIL import Image, ImageOps
img = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
w0, h0 = img.size
scale = resolution / min(w0, h0)
rw, rh = max(resolution, round(w0 * scale)), max(resolution, round(h0 * scale))
img = img.resize((rw, rh), Image.LANCZOS)
return img.resize((rw, rh), Image.LANCZOS), rw, rh
def _to_unit_tensor(img):
import numpy as np
import torch
arr = np.asarray(img, dtype = np.float32) / 255.0
return torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0
def _load_pixel_tensor(path, resolution, center_crop, random_flip, rng):
"""Load an image -> a normalised [3,H,W] tensor in [-1,1]. Same geometry as the SDXL
loader but without the SDXL time-ids (DiT families don't use them)."""
from PIL import Image
img, rw, rh = _open_resized(path, resolution)
if center_crop:
left, top = (rw - resolution) // 2, (rh - resolution) // 2
else:
@ -461,8 +601,132 @@ def _load_pixel_tensor(path, resolution, center_crop, random_flip, rng):
img = img.crop((left, top, left + resolution, top + resolution))
if random_flip and rng.random() < 0.5:
img = img.transpose(Image.FLIP_LEFT_RIGHT)
arr = np.asarray(img, dtype = np.float32) / 255.0
return torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0
return _to_unit_tensor(img)
def _load_pixel_tensor_planned(path, resolution, center_crop, u_left, u_top, flip):
"""Deterministic variant of ``_load_pixel_tensor`` for the latent cache: the crop comes
as unit fractions (mapped uniformly over the same inclusive integer range ``randint``
draws from) and the flip as a bool. ``center_crop`` reproduces the exact legacy
floor-div center so a cached center-crop run matches the uncached one bit-for-bit."""
from PIL import Image
img, rw, rh = _open_resized(path, resolution)
if center_crop:
left, top = (rw - resolution) // 2, (rh - resolution) // 2
else:
left = min(int(u_left * (rw - resolution + 1)), max(0, rw - resolution))
top = min(int(u_top * (rh - resolution + 1)), max(0, rh - resolution))
img = img.crop((left, top, left + resolution, top + resolution))
if flip:
img = img.transpose(Image.FLIP_LEFT_RIGHT)
return _to_unit_tensor(img)
def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_event, check_stop):
"""Precompute the per-image latent posterior cache: for each planned crop/flip variant,
encode once and store the affine (A, B) pair on CPU (pinned when possible) in the
training dtype. Returns None if the build was interrupted by a stop request."""
plan = _plan_cache_variants(
len(image_paths), cfg.cache_variants, cfg.center_crop, cfg.random_flip, cfg.seed
)
def _hold(t):
if t is None:
return None
t = t.to(weight_dtype).cpu()
if device == "cuda":
try:
t = t.pin_memory()
except RuntimeError:
pass
return t
cache: list[list[tuple]] = []
total = len(image_paths)
for i, path in enumerate(image_paths):
variants = []
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
)
.unsqueeze(0)
.to(device)
)
a, b = spec.encode_latent_stats(vae, px)
variants.append((_hold(a), _hold(b)))
cache.append(variants)
if (i + 1) % 4 == 0 or i + 1 == total:
_emit(on_event, "preparing", stage = "cache_latents", done = i + 1, total = total)
if check_stop():
return None
return cache
def _sample_cached_latents(cache, idxs, variant_rng, device):
"""Draw one latent per index from the cache: pick a variant, then sample the posterior
(A + B * randn) when the family is stochastic. Fresh noise per step, exactly like an
in-loop ``latent_dist.sample()``."""
import torch
parts_a, parts_b = [], []
for i in idxs:
variants = cache[i]
a, b = variants[variant_rng.randrange(len(variants))] if len(variants) > 1 else variants[0]
parts_a.append(a)
parts_b.append(b)
lat_a = torch.cat(parts_a).to(device, non_blocking = True)
if parts_b[0] is None:
return lat_a
lat_b = torch.cat(parts_b).to(device, non_blocking = True)
return lat_a + lat_b * torch.randn_like(lat_a)
def _should_compile(cfg, base_is_bnb, device) -> bool:
mode = (cfg.compile_transformer or "auto").strip().lower()
if device != "cuda" or mode == "off":
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
def _maybe_compile_transformer(transformer, cfg, base_is_bnb, device, on_event) -> 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):
return False
import torch
fn = getattr(transformer, "compile_repeated_blocks", None)
if not callable(fn):
_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)
if dynamo_cfg is not None:
# Heterogeneous-block DiTs (Z-Image: ~11 distinct block shapes) exceed dynamo's
# default recompile limit of 8; bump it the same way the inference speed layer
# does (diffusers' documented regional-compile fix).
for attr in ("recompile_limit", "cache_size_limit"):
if hasattr(dynamo_cfg, attr):
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)
return True
except Exception as exc: # noqa: BLE001 -- optimisation only, never fatal
_emit(on_event, "warning", message = f"torch.compile disabled (eager fallback): {exc}")
return False
def run_dit_lora_training(
@ -473,11 +737,6 @@ def run_dit_lora_training(
) -> str:
"""Train a flow-matching DiT LoRA (FLUX.1-dev / Qwen-Image / Z-Image) and export it."""
import torch
import torch.nn.functional as F
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.training_utils import cast_training_params
from peft import LoraConfig
from peft.utils import get_peft_model_state_dict
cfg = config.normalized()
spec = _SPECS.get(cfg.resolved_family)
@ -486,6 +745,7 @@ def run_dit_lora_training(
rng = random.Random(cfg.seed)
torch.manual_seed(cfg.seed)
_FLUX_STATIC.clear()
save_on_stop = True
@ -513,7 +773,6 @@ def run_dit_lora_training(
# The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is
# unsupported for real runs but keeps import/unit tests architecture-agnostic).
weight_dtype = torch.bfloat16 if device == "cuda" else torch.float32
use_lora_targets = tuple(cfg.lora_target_modules) or spec.lora_targets
_assert_trusted_base_model(cfg.base_model)
_assert_gated_access(cfg.base_model, cfg.hf_token)
@ -528,9 +787,35 @@ def run_dit_lora_training(
)
return str(out_dir)
# QLoRA by default for the big DiTs (nf4 transformer). The prequant Qwen/Z-Image repos
# are already 4-bit; FLUX quantizes its transformer on the fly.
pipe, transformer, vae = spec.load(cfg, device, weight_dtype, qlora = True)
# TF32 / cudnn.benchmark for the run, restored on the way out (the trainer subprocess is
# disposable, but restoring keeps in-process callers -- tests, notebooks -- clean).
perf_snap = _apply_perf_flags(cfg, device)
try:
return _train_dit(
cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_stop,
lambda: save_on_stop,
)
finally:
_restore_perf_flags(perf_snap)
def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_stop, _save_on_stop):
"""The body of ``run_dit_lora_training``, split out so the backend perf flags are
snapshot/restored around it in exactly one place."""
import torch
import torch.nn.functional as F
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.optimization import get_scheduler
from diffusers.training_utils import cast_training_params
from peft import LoraConfig
from peft.utils import get_peft_model_state_dict
use_lora_targets = tuple(cfg.lora_target_modules) or spec.lora_targets
out_dir = Path(cfg.output_dir).expanduser()
# Phase 1: conditioning only. The pipeline loads WITHOUT its transformer, so the text
# encoders + VAE never share VRAM with the multi-GB denoiser.
pipe, vae = spec.load_conditioners(cfg, device, weight_dtype)
# Precompute all caption embeddings, then free the (large) text encoder(s): captions are
# constant and the encoders are frozen, so this is exact and the biggest memory win.
@ -544,6 +829,40 @@ def run_dit_lora_training(
if device == "cuda":
torch.cuda.empty_cache()
# Phase 2: the VAE latent cache, then free the VAE too (see module docstring: the cache
# keeps the posterior affine parameters, so per-step sampling noise is preserved).
use_cache = cfg.cache_latents and os.environ.get(
"UNSLOTH_DIFFUSION_NO_LATENT_CACHE", ""
) not in ("1", "true")
latent_cache = None
if use_cache:
latent_cache = _build_latent_cache(
spec, vae, image_paths, cfg, device, weight_dtype, on_event, _check_stop
)
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,
)
return str(out_dir)
try:
pipe.vae = None
except Exception: # noqa: BLE001 -- a pipeline without a settable vae keeps it
pass
del vae
vae = None
gc.collect()
if device == "cuda":
torch.cuda.empty_cache()
# Variant picks use their own stream so the training loop's index/noise draws stay on
# 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
# Freeze the base; attach the trainable LoRA to the transformer.
transformer.requires_grad_(False)
transformer.add_adapter(
@ -568,57 +887,77 @@ def run_dit_lora_training(
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)
# 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
if compiled and spec.family == "qwen-image":
qwen_pad_to = max(e[0].shape[1] for e in caption_embeds.values())
optimizer = _make_optimizer(lora_params, cfg.learning_rate)
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
cfg.base_model, subfolder = "scheduler", token = cfg.hf_token
)
# The LR schedule advances once per optimizer update, so warmup/decay are counted in
# optimizer steps (matching the SDXL trainer; multiplying by the accumulation factor
# would stretch warmup past the run and never reach the decay).
lr_sched = get_scheduler(
cfg.lr_scheduler,
optimizer = optimizer,
num_warmup_steps = cfg.lr_warmup_steps,
num_training_steps = cfg.train_steps,
)
_emit(on_event, "model_load_completed")
_emit(on_event, "model_load_completed", compiled = compiled)
transformer.train()
n_images = len(image_paths)
batch_size = cfg.train_batch_size
stopped = False
running_loss = 0.0
peak_gb = 0.0
t_start = time.time()
t_steady = None
done = 0
# bf16 autocast around the forward + loss, matching the diffusers dreambooth scripts'
# accelerator.autocast: it reconciles the fp32 LoRA params with the bnb 4-bit base
# matmuls in one compute dtype. Without it the 4-bit backward on FLUX dies with an
# illegal-address / CUBLAS failure.
autocast = (
torch.autocast(device_type = "cuda", dtype = torch.bfloat16)
if device == "cuda"
else nullcontext()
)
for opt_step in range(cfg.train_steps):
optimizer.zero_grad(set_to_none = True)
step_loss = 0.0
for _ in range(cfg.gradient_accumulation_steps):
i = rng.randrange(len(image_paths))
px = (
_load_pixel_tensor(
image_paths[i], cfg.resolution, cfg.center_crop, cfg.random_flip, rng
)
.unsqueeze(0)
.to(device)
)
latents = spec.encode_latents(vae, px).to(weight_dtype)
idxs = [rng.randrange(n_images) for _ in range(batch_size)]
if latent_cache is not None:
latents = _sample_cached_latents(latent_cache, idxs, variant_rng, device)
else:
px = torch.stack(
[
_load_pixel_tensor(
image_paths[i], cfg.resolution, cfg.center_crop, cfg.random_flip, rng
)
for i in idxs
]
).to(device)
latents = spec.encode_latents(vae, px).to(weight_dtype)
noise = torch.randn_like(latents)
timesteps = _sample_timesteps(scheduler, latents.shape[0], device)
sigmas = _get_sigmas(scheduler, timesteps, device, weight_dtype, latents.ndim)
timesteps, t_indices = _sample_timesteps(scheduler, latents.shape[0], device)
sigmas = _gather_sigmas(scheduler, t_indices, device, weight_dtype, latents.ndim)
noisy = (1.0 - sigmas) * latents + sigmas * noise
emb = caption_embeds[captions[i]]
emb_dev = tuple(
t.to(device = device, dtype = weight_dtype)
if (t is not None and t.is_floating_point())
else (t.to(device) if t is not None else None)
for t in emb
)
# bf16 autocast around the forward + loss, matching the diffusers dreambooth
# scripts' accelerator.autocast: it reconciles the fp32 LoRA params with the
# bnb 4-bit base matmuls in one compute dtype. Without it the 4-bit backward
# on FLUX dies with an illegal-address / CUBLAS failure.
autocast = (
torch.autocast(device_type = "cuda", dtype = torch.bfloat16)
if device == "cuda"
else nullcontext()
embeds = spec.collate(
[caption_embeds[captions[i]] for i in idxs], device, weight_dtype,
pad_to = qwen_pad_to,
)
with autocast:
model_pred = spec.forward(
transformer, noisy, timesteps, sigmas, emb_dev, cfg, device, weight_dtype
transformer, noisy, timesteps, sigmas, embeds, cfg, device, weight_dtype
)
target = noise - latents
loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean")
@ -628,17 +967,23 @@ def run_dit_lora_training(
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)
optimizer.step()
lr_sched.step()
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.
t_steady = now
if done % cfg.log_every == 0 or done == cfg.train_steps:
if device == "cuda":
peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2)
sps = round(
(done * cfg.train_batch_size * cfg.gradient_accumulation_steps)
/ max(time.time() - t_start, 1e-6),
3,
)
per_step = batch_size * cfg.gradient_accumulation_steps
if t_steady is not None and done > 1:
sps = round((done - 1) * per_step / max(now - t_steady, 1e-6), 3)
else:
sps = round(done * per_step / max(now - t_start, 1e-6), 3)
_emit(
on_event,
"progress",
@ -646,7 +991,7 @@ def run_dit_lora_training(
total_steps = cfg.train_steps,
loss = round(step_loss, 5),
avg_loss = round(running_loss / done, 5),
learning_rate = cfg.learning_rate,
learning_rate = lr_sched.get_last_lr()[0],
samples_per_second = sps,
peak_memory_gb = peak_gb or None,
)
@ -654,10 +999,9 @@ def run_dit_lora_training(
stopped = True
break
out_dir = Path(cfg.output_dir).expanduser()
lora_path: Optional[str] = None
catalog_path: Optional[str] = None
if not (stopped and not save_on_stop):
if not (stopped and not _save_on_stop()):
out_dir.mkdir(parents = True, exist_ok = True)
layers = get_peft_model_state_dict(transformer)
spec.save(pipe, str(out_dir), layers)
@ -673,19 +1017,27 @@ def run_dit_lora_training(
base_model = cfg.base_model,
stopped = stopped,
steps_run = done if cfg.train_steps else 0,
wall_seconds = round(time.time() - t_start, 1),
)
return str(out_dir)
def _make_optimizer(params, lr):
"""8-bit AdamW (bitsandbytes) when available -- half the optimizer state, no accuracy
regression for LoRA -- else the torch AdamW fallback."""
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)
except Exception: # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW
return torch.optim.AdamW(params, lr = lr)
pass
if torch.cuda.is_available():
try:
return torch.optim.AdamW(params, lr = lr, fused = True)
except Exception: # noqa: BLE001 -- fused unsupported on this build/device
pass
return torch.optim.AdamW(params, lr = lr)
def _free_text_encoders(pipe) -> None:

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

@ -20,6 +20,13 @@ Design:
- ``run_diffusion_training_process`` is the thin mp.Queue adapter; it dispatches to the
trainer registered for the resolved family (SDXL here, DiT families in a follow-up).
``main`` is a CLI.
Memory/perf: captions are encoded once up front and the CLIP text encoders freed; VAE
latents are likewise precomputed into a small CPU cache (``cache_latents``) and the VAE
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``.
"""
from __future__ import annotations
@ -40,12 +47,15 @@ from core.training.diffusion_train_common import ( # noqa: F401
EventCb,
StopCb,
DiffusionLoraConfig,
_apply_perf_flags,
_assert_trusted_base_model,
_coerce_gradient_checkpointing,
_config_from_dict,
_CONFIG_ALIASES,
_emit,
_plan_cache_variants,
_publish_to_lora_catalog,
_restore_perf_flags,
discover_image_caption_pairs,
get_trainer,
)
@ -99,6 +109,43 @@ def _load_image_tensor(
return tensor, time_ids
def _load_image_tensor_planned(
path: str, resolution: int, center_crop: bool, u_left: float, u_top: float, flip: bool
) -> tuple[Any, tuple[int, int, int, int, int, int]]:
"""Deterministic variant of ``_load_image_tensor`` for the latent cache: the crop comes
as unit fractions (mapped uniformly over the same inclusive integer range ``randint``
draws from) and the flip as a bool. Geometry (EXIF transpose, LANCZOS short-side resize,
the SDXL ``add_time_ids`` from the original size + actual crop offset) matches
``_load_image_tensor`` exactly; ``center_crop`` reproduces the legacy floor-div center
bit-for-bit. The flip does not change time_ids (only the mirrored crop_left does)."""
import numpy as np
import torch
from PIL import Image, ImageOps
img = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
original_w, original_h = img.size
scale = resolution / min(original_w, original_h)
resized_w = max(resolution, round(original_w * scale))
resized_h = max(resolution, round(original_h * scale))
img = img.resize((resized_w, resized_h), Image.LANCZOS)
if center_crop:
left, top = (resized_w - resolution) // 2, (resized_h - resolution) // 2
else:
left = min(int(u_left * (resized_w - resolution + 1)), max(0, resized_w - resolution))
top = min(int(u_top * (resized_h - resolution + 1)), max(0, resized_h - resolution))
img = img.crop((left, top, left + resolution, top + resolution))
crop_left = left
if flip:
img = img.transpose(Image.FLIP_LEFT_RIGHT)
# Mirror the crop's left origin so the conditioning matches the flipped pixels, the
# same mirroring ``_load_image_tensor`` applies on a random flip.
crop_left = max(0, resized_w - resolution - left)
arr = np.asarray(img, dtype = np.float32) / 255.0
tensor = torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0
time_ids = (original_h, original_w, top, crop_left, resolution, resolution)
return tensor, time_ids
def _encode_sdxl_prompts(
prompts: list[str], tokenizers: list, text_encoders: list, device: Any
) -> tuple:
@ -126,6 +173,75 @@ def _encode_sdxl_prompts(
return prompt_embeds, pooled
def _build_sdxl_latent_cache(
vae, vae_scale, image_paths, cfg, device, weight_dtype, on_event, check_stop
):
"""Precompute the per-image latent posterior cache: for each planned crop/flip variant,
encode once and store ``(A, B, time_ids)`` on CPU in the training dtype. ``A`` and ``B``
are the affine posterior parameters (mean/std with the VAE scale folded in) so a per-step
sample is ``A + B * randn`` -- distribution-identical to an in-loop ``latent_dist.sample()``
-- and ``time_ids`` is the SDXL micro-conditioning for the crop. Returns None if the build
was interrupted by a stop request. ``vae_scale`` is read before the VAE is freed."""
import torch
plan = _plan_cache_variants(
len(image_paths), cfg.cache_variants, cfg.center_crop, cfg.random_flip, cfg.seed
)
def _hold(t):
t = t.to(weight_dtype).cpu()
if device == "cuda":
try:
t = t.pin_memory()
except RuntimeError:
pass
return t
cache: list[list[tuple]] = []
total = len(image_paths)
for i, path in enumerate(image_paths):
variants = []
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
)
pixel_values = tensor.unsqueeze(0).to(device, dtype = torch.float32)
with torch.no_grad():
dist = vae.encode(pixel_values).latent_dist
a = _hold(dist.mean * vae_scale)
b = _hold(dist.std * vae_scale)
variants.append((a, b, tuple(time_ids)))
cache.append(variants)
if (i + 1) % 4 == 0 or i + 1 == total:
_emit(on_event, "preparing", stage = "cache_latents", done = i + 1, total = total)
if check_stop():
return None
return cache
def _sample_sdxl_cached_latents(cache, idxs, variant_rng, device, weight_dtype):
"""Draw one latent + its time_ids per index from the cache: pick a variant, then sample
the posterior (A + B * randn) with fresh noise per step, exactly like an in-loop
``latent_dist.sample() * vae_scale``. Returns ``(latents, batch_time_ids)`` already on
``device`` in the training dtype (scale + dtype are folded into the cache)."""
import torch
parts_a, parts_b, tid_rows = [], [], []
for i in idxs:
variants = cache[i]
a, b, time_ids = (
variants[variant_rng.randrange(len(variants))] if len(variants) > 1 else variants[0]
)
parts_a.append(a)
parts_b.append(b)
tid_rows.append(time_ids)
lat_a = torch.cat(parts_a).to(device, non_blocking = True)
lat_b = torch.cat(parts_b).to(device, non_blocking = True)
latents = lat_a + lat_b * torch.randn_like(lat_a)
batch_time_ids = torch.tensor(tid_rows, device = device, dtype = weight_dtype)
return latents, batch_time_ids
def run_diffusion_lora_training(
config: DiffusionLoraConfig,
*,
@ -174,247 +290,305 @@ def run_diffusion_lora_training(
precision = "fp16"
weight_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "no": torch.float32}[precision]
# Preflight the base model against the same trust gate as inference, before any fetch.
_assert_trusted_base_model(cfg.base_model)
# TF32 / cudnn.benchmark for the run, restored on the way out (the trainer subprocess is
# disposable, but restoring keeps in-process callers -- tests, notebooks -- clean). Wraps
# the whole body so every return (early stop and normal) restores the backend flags.
snap = _apply_perf_flags(cfg, device)
try:
# Preflight the base model against the same trust gate as inference, before any fetch.
_assert_trusted_base_model(cfg.base_model)
pairs = discover_image_caption_pairs(
cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
)
_emit(on_event, "model_load_started", num_images = len(pairs))
pairs = discover_image_caption_pairs(
cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
)
_emit(on_event, "model_load_started", num_images = len(pairs))
# Honour a stop requested before the (potentially large / slow) base model loads, the
# same way the LLM training worker checks its stop thread around model load.
if _check_stop():
# Honour a stop requested before the (potentially large / slow) base model loads, the
# same way the LLM training worker checks its stop thread around model load.
if _check_stop():
out_dir = Path(cfg.output_dir).expanduser()
_emit(
on_event,
"complete",
output_dir = str(out_dir),
lora_path = None,
stopped = True,
steps_run = 0,
)
return str(out_dir)
pipe = StableDiffusionXLPipeline.from_pretrained(
cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token, add_watermarker = False
)
unet, vae = pipe.unet, pipe.vae
tokenizers = [pipe.tokenizer, pipe.tokenizer_2]
text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
# Freeze the base; only the LoRA trains. The SDXL VAE overflows fp16, so keep it fp32.
for m in (unet, vae, *text_encoders):
m.requires_grad_(False)
vae.to(device, dtype = torch.float32)
for m in (unet, *text_encoders):
m.to(device, dtype = weight_dtype)
unet.add_adapter(
LoraConfig(
r = cfg.lora_rank,
lora_alpha = cfg.lora_alpha,
lora_dropout = cfg.lora_dropout,
init_lora_weights = "gaussian",
target_modules = list(cfg.lora_target_modules),
)
)
if cfg.gradient_checkpointing:
unet.enable_gradient_checkpointing()
# LoRA params must be fp32 for a stable optimizer under mixed precision.
if weight_dtype != torch.float32:
cast_training_params(unet, dtype = torch.float32)
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
# time per outer opt_step (after the accumulation inner loop), for cfg.train_steps
# total. Count warmup/decay in those optimizer steps -- multiplying by the
# accumulation factor would stretch warmup past the run and never reach the decay.
lr_sched = get_scheduler(
cfg.lr_scheduler,
optimizer = optimizer,
num_warmup_steps = cfg.lr_warmup_steps,
num_training_steps = cfg.train_steps,
)
vae_scale = vae.config.scaling_factor
prediction_type = noise_scheduler.config.prediction_type
# Precompute text embeddings once per unique caption, then free the CLIP text encoders.
# SDXL re-encoded captions every step (pure waste: captions are constant) and kept both
# text encoders (~1.5 GB) resident. Embeddings are deterministic and this consumes no
# torch RNG, so the training math is bit-identical to in-loop encoding -- only faster and
# lighter. The env toggle exists purely so the accuracy guard can A/B the two paths.
precompute = os.environ.get("UNSLOTH_DIFFUSION_NO_PRECOMPUTE", "") not in ("1", "true")
caption_embeds: dict[str, tuple] = {}
if precompute:
for cap in sorted({c for _, c in pairs}):
pe, pooled_c = _encode_sdxl_prompts([cap], tokenizers, text_encoders, device)
caption_embeds[cap] = (pe.cpu(), pooled_c.cpu())
for te in text_encoders:
te.to("cpu")
text_encoders = []
gc.collect()
if device == "cuda":
torch.cuda.empty_cache()
# Precompute the VAE latent cache, then free the VAE: the cache holds the posterior
# affine pair (mean/std, scale folded in) so per-step sampling noise is preserved. The
# env toggle lets the accuracy guard A/B the cached vs in-loop encode paths.
use_cache = cfg.cache_latents and os.environ.get(
"UNSLOTH_DIFFUSION_NO_LATENT_CACHE", ""
) not in ("1", "true")
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,
_check_stop,
)
if latent_cache is None: # stopped during the cache build; nothing trained yet
out_dir = Path(cfg.output_dir).expanduser()
_emit(
on_event,
"complete",
output_dir = str(out_dir),
lora_path = None,
stopped = True,
steps_run = 0,
)
return str(out_dir)
try:
pipe.vae = None
except Exception: # noqa: BLE001 -- a pipeline without a settable vae keeps it
pass
del vae
vae = None
gc.collect()
if device == "cuda":
torch.cuda.empty_cache()
# Variant picks use their own stream so the loop's index/noise draws stay on the same
# seed-deterministic sequence whether or not the cache is enabled.
variant_rng = random.Random(cfg.seed + 1)
_emit(on_event, "model_load_completed")
def _next_batch() -> tuple[list[int], list[str], list[str]]:
idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs)))
chosen = [pairs[i] for i in idx]
return idx, [c[0] for c in chosen], [c[1] for c in chosen]
unet.train()
stopped = False
micro = 0
running_loss = 0.0
peak_gb = 0.0
t_start = time.time()
done = 0
for opt_step in range(cfg.train_steps):
optimizer.zero_grad(set_to_none = True)
step_loss = 0.0
for _ in range(cfg.gradient_accumulation_steps):
idx, img_paths, captions = _next_batch()
if latent_cache is not None:
# Scale + dtype are already folded into the cache, so do not re-apply.
latents, batch_time_ids = _sample_sdxl_cached_latents(
latent_cache, idx, variant_rng, device, weight_dtype
)
else:
loaded = [
_load_image_tensor(p, cfg.resolution, cfg.center_crop, cfg.random_flip, rng)
for p in img_paths
]
pixel_values = torch.stack([t for t, _ in loaded]).to(
device, dtype = torch.float32
)
# Per-sample SDXL micro-conditioning from the actual crop (original size + offset).
batch_time_ids = torch.tensor(
[tid for _, tid in loaded], device = device, dtype = weight_dtype
)
with torch.no_grad():
latents = vae.encode(pixel_values).latent_dist.sample() * vae_scale
latents = latents.to(dtype = weight_dtype)
noise = torch.randn_like(latents)
bsz = latents.shape[0]
timesteps = torch.randint(
0, noise_scheduler.config.num_train_timesteps, (bsz,), device = device
).long()
noisy = noise_scheduler.add_noise(latents, noise, timesteps)
if precompute:
prompt_embeds = torch.cat([caption_embeds[c][0] for c in captions]).to(device)
pooled = torch.cat([caption_embeds[c][1] for c in captions]).to(device)
else:
prompt_embeds, pooled = _encode_sdxl_prompts(
captions, tokenizers, text_encoders, device
)
prompt_embeds = prompt_embeds.to(dtype = weight_dtype)
pooled = pooled.to(dtype = weight_dtype)
added = {"text_embeds": pooled, "time_ids": batch_time_ids}
model_pred = unet(
noisy, timesteps, prompt_embeds, added_cond_kwargs = added, return_dict = False
)[0]
if prediction_type == "v_prediction":
target = noise_scheduler.get_velocity(latents, noise, timesteps)
else:
target = noise
if cfg.snr_gamma is not None:
snr = compute_snr(noise_scheduler, timesteps)
w = torch.stack([snr, cfg.snr_gamma * torch.ones_like(timesteps)], dim = 1).min(
dim = 1
)[0]
w = w / snr if prediction_type != "v_prediction" else w / (snr + 1)
loss = F.mse_loss(model_pred.float(), target.float(), reduction = "none")
loss = loss.mean(dim = list(range(1, loss.ndim))) * w
loss = loss.mean()
else:
loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean")
(loss / cfg.gradient_accumulation_steps).backward()
step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
micro += 1
# max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that);
# passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning).
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)
optimizer.step()
lr_sched.step()
running_loss += step_loss
done = opt_step + 1
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,
)
_emit(
on_event,
"progress",
step = done,
total_steps = cfg.train_steps,
loss = round(step_loss, 5),
avg_loss = round(running_loss / done, 5),
learning_rate = lr_sched.get_last_lr()[0],
samples_per_second = samples_per_second,
peak_memory_gb = peak_gb or None,
)
if _check_stop():
stopped = True
break
# Export the trained LoRA in diffusers format (loadable via load_lora_weights), unless
# the run was cancelled with save disabled -- then leave no partial adapter behind.
out_dir = Path(cfg.output_dir).expanduser()
lora_path: Optional[str] = None
catalog_path: Optional[str] = None
if not (stopped and not save_on_stop):
out_dir.mkdir(parents = True, exist_ok = True)
unet_lora = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
StableDiffusionXLPipeline.save_lora_weights(
save_directory = str(out_dir),
unet_lora_layers = unet_lora,
safe_serialization = True,
weight_name = DEFAULT_LORA_FILENAME,
)
lora_path = str(out_dir / DEFAULT_LORA_FILENAME)
# Mirror into the Studio diffusion LoRA directory so the Images picker discovers it
# (its scan lists only files directly under loras/diffusion, not subdirectories).
catalog_path = _publish_to_lora_catalog(lora_path, cfg)
_emit(
on_event,
"complete",
output_dir = str(out_dir),
lora_path = None,
stopped = True,
steps_run = 0,
lora_path = lora_path,
catalog_path = catalog_path,
family = cfg.resolved_family,
base_model = cfg.base_model,
stopped = stopped,
steps_run = done if cfg.train_steps else 0,
)
return str(out_dir)
pipe = StableDiffusionXLPipeline.from_pretrained(
cfg.base_model, torch_dtype = weight_dtype, token = cfg.hf_token, add_watermarker = False
)
unet, vae = pipe.unet, pipe.vae
tokenizers = [pipe.tokenizer, pipe.tokenizer_2]
text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
# Freeze the base; only the LoRA trains. The SDXL VAE overflows fp16, so keep it fp32.
for m in (unet, vae, *text_encoders):
m.requires_grad_(False)
vae.to(device, dtype = torch.float32)
for m in (unet, *text_encoders):
m.to(device, dtype = weight_dtype)
unet.add_adapter(
LoraConfig(
r = cfg.lora_rank,
lora_alpha = cfg.lora_alpha,
lora_dropout = cfg.lora_dropout,
init_lora_weights = "gaussian",
target_modules = list(cfg.lora_target_modules),
)
)
if cfg.gradient_checkpointing:
unet.enable_gradient_checkpointing()
# LoRA params must be fp32 for a stable optimizer under mixed precision.
if weight_dtype != torch.float32:
cast_training_params(unet, dtype = torch.float32)
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
# time per outer opt_step (after the accumulation inner loop), for cfg.train_steps
# total. Count warmup/decay in those optimizer steps -- multiplying by the
# accumulation factor would stretch warmup past the run and never reach the decay.
lr_sched = get_scheduler(
cfg.lr_scheduler,
optimizer = optimizer,
num_warmup_steps = cfg.lr_warmup_steps,
num_training_steps = cfg.train_steps,
)
vae_scale = vae.config.scaling_factor
prediction_type = noise_scheduler.config.prediction_type
# Precompute text embeddings once per unique caption, then free the CLIP text encoders.
# SDXL re-encoded captions every step (pure waste: captions are constant) and kept both
# text encoders (~1.5 GB) resident. Embeddings are deterministic and this consumes no
# torch RNG, so the training math is bit-identical to in-loop encoding -- only faster and
# lighter. The env toggle exists purely so the accuracy guard can A/B the two paths.
precompute = os.environ.get("UNSLOTH_DIFFUSION_NO_PRECOMPUTE", "") not in ("1", "true")
caption_embeds: dict[str, tuple] = {}
if precompute:
for cap in sorted({c for _, c in pairs}):
pe, pooled_c = _encode_sdxl_prompts([cap], tokenizers, text_encoders, device)
caption_embeds[cap] = (pe.cpu(), pooled_c.cpu())
for te in text_encoders:
te.to("cpu")
text_encoders = []
gc.collect()
if device == "cuda":
torch.cuda.empty_cache()
_emit(on_event, "model_load_completed")
def _next_batch() -> tuple[list[str], list[str]]:
idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs)))
chosen = [pairs[i] for i in idx]
return [c[0] for c in chosen], [c[1] for c in chosen]
unet.train()
stopped = False
micro = 0
running_loss = 0.0
peak_gb = 0.0
t_start = time.time()
done = 0
for opt_step in range(cfg.train_steps):
optimizer.zero_grad(set_to_none = True)
step_loss = 0.0
for _ in range(cfg.gradient_accumulation_steps):
img_paths, captions = _next_batch()
loaded = [
_load_image_tensor(p, cfg.resolution, cfg.center_crop, cfg.random_flip, rng)
for p in img_paths
]
pixel_values = torch.stack([t for t, _ in loaded]).to(device, dtype = torch.float32)
# Per-sample SDXL micro-conditioning from the actual crop (original size + offset).
batch_time_ids = torch.tensor(
[tid for _, tid in loaded], device = device, dtype = weight_dtype
)
with torch.no_grad():
latents = vae.encode(pixel_values).latent_dist.sample() * vae_scale
latents = latents.to(dtype = weight_dtype)
noise = torch.randn_like(latents)
bsz = latents.shape[0]
timesteps = torch.randint(
0, noise_scheduler.config.num_train_timesteps, (bsz,), device = device
).long()
noisy = noise_scheduler.add_noise(latents, noise, timesteps)
if precompute:
prompt_embeds = torch.cat([caption_embeds[c][0] for c in captions]).to(device)
pooled = torch.cat([caption_embeds[c][1] for c in captions]).to(device)
else:
prompt_embeds, pooled = _encode_sdxl_prompts(
captions, tokenizers, text_encoders, device
)
prompt_embeds = prompt_embeds.to(dtype = weight_dtype)
pooled = pooled.to(dtype = weight_dtype)
added = {"text_embeds": pooled, "time_ids": batch_time_ids}
model_pred = unet(
noisy, timesteps, prompt_embeds, added_cond_kwargs = added, return_dict = False
)[0]
if prediction_type == "v_prediction":
target = noise_scheduler.get_velocity(latents, noise, timesteps)
else:
target = noise
if cfg.snr_gamma is not None:
snr = compute_snr(noise_scheduler, timesteps)
w = torch.stack([snr, cfg.snr_gamma * torch.ones_like(timesteps)], dim = 1).min(
dim = 1
)[0]
w = w / snr if prediction_type != "v_prediction" else w / (snr + 1)
loss = F.mse_loss(model_pred.float(), target.float(), reduction = "none")
loss = loss.mean(dim = list(range(1, loss.ndim))) * w
loss = loss.mean()
else:
loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean")
(loss / cfg.gradient_accumulation_steps).backward()
step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
micro += 1
# max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that);
# passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning).
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)
optimizer.step()
lr_sched.step()
running_loss += step_loss
done = opt_step + 1
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,
)
_emit(
on_event,
"progress",
step = done,
total_steps = cfg.train_steps,
loss = round(step_loss, 5),
avg_loss = round(running_loss / done, 5),
learning_rate = lr_sched.get_last_lr()[0],
samples_per_second = samples_per_second,
peak_memory_gb = peak_gb or None,
)
if _check_stop():
stopped = True
break
# Export the trained LoRA in diffusers format (loadable via load_lora_weights), unless
# the run was cancelled with save disabled -- then leave no partial adapter behind.
out_dir = Path(cfg.output_dir).expanduser()
lora_path: Optional[str] = None
catalog_path: Optional[str] = None
if not (stopped and not save_on_stop):
out_dir.mkdir(parents = True, exist_ok = True)
unet_lora = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
StableDiffusionXLPipeline.save_lora_weights(
save_directory = str(out_dir),
unet_lora_layers = unet_lora,
safe_serialization = True,
weight_name = DEFAULT_LORA_FILENAME,
)
lora_path = str(out_dir / DEFAULT_LORA_FILENAME)
# Mirror into the Studio diffusion LoRA directory so the Images picker discovers it
# (its scan lists only files directly under loras/diffusion, not subdirectories).
catalog_path = _publish_to_lora_catalog(lora_path, cfg)
_emit(
on_event,
"complete",
output_dir = str(out_dir),
lora_path = lora_path,
catalog_path = catalog_path,
family = cfg.resolved_family,
base_model = cfg.base_model,
stopped = stopped,
steps_run = done if cfg.train_steps else 0,
)
return str(out_dir)
finally:
_restore_perf_flags(snap)
def _make_lora_optimizer(params: list, lr: float) -> Any:
"""8-bit AdamW (bitsandbytes) by default -- half the optimizer state, no meaningful
quality cost for LoRA -- falling back to fp32 AdamW when unavailable or when
UNSLOTH_DIFFUSION_FP32_OPTIM is set (used by the accuracy guard)."""
quality cost for LoRA -- falling back to torch AdamW (fused on CUDA) when unavailable.
UNSLOTH_DIFFUSION_FP32_OPTIM forces plain (non-fused) AdamW: the accuracy guard wants the
reference optimizer, so it must not take the fused path."""
import torch
if os.environ.get("UNSLOTH_DIFFUSION_FP32_OPTIM", "") not in ("1", "true"):
if os.environ.get("UNSLOTH_DIFFUSION_FP32_OPTIM", "") in ("1", "true"):
return torch.optim.AdamW(params, lr = lr)
try:
import bitsandbytes as bnb
return bnb.optim.AdamW8bit(params, lr = lr)
except Exception: # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW
pass
if torch.cuda.is_available():
try:
import bitsandbytes as bnb
return bnb.optim.AdamW8bit(params, lr = lr)
except Exception: # noqa: BLE001 -- bnb missing / no CUDA: fall back to torch AdamW
return torch.optim.AdamW(params, lr = lr, fused = True)
except Exception: # noqa: BLE001 -- fused unsupported on this build/device
pass
return torch.optim.AdamW(params, lr = lr)

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

@ -17,6 +17,7 @@ from __future__ import annotations
import json
import os
import random
import re
import time
from dataclasses import dataclass, field, replace
@ -228,6 +229,17 @@ class DiffusionLoraConfig:
caption_column: str = "text" # column in metadata.jsonl
adapter_name: str = "default"
hf_token: Optional[str] = None
# Precompute the VAE latents once (freeing the VAE for the whole run) instead of
# re-encoding every step. ``cache_variants`` crop/flip draws are frozen per image;
# the per-step VAE sampling noise itself is preserved (see the DiT trainer docstring).
cache_latents: bool = True
cache_variants: int = 4
# Regional torch.compile of the transformer blocks: "off" | "on" | "auto" (auto turns
# it on only for a dense, non-bitsandbytes base where it is a clean win).
compile_transformer: str = "auto"
# TF32 matmuls + high fp32 matmul precision + cudnn autotuning for the run. Near-lossless;
# disable for strict bit-reproducibility A/Bs.
enable_tf32: bool = True
# How often to emit a progress event (in optimizer steps).
log_every: int = 1
# Optional explicit family override ("sdxl" / "flux.1" / ...); None = detect from
@ -259,6 +271,11 @@ class DiffusionLoraConfig:
raise ValueError("resolution must be a multiple of 8 and >= 64")
if self.mixed_precision not in ("bf16", "fp16", "no"):
raise ValueError("mixed_precision must be one of bf16 / fp16 / no")
if not 1 <= int(self.cache_variants) <= 16:
raise ValueError("cache_variants must be between 1 and 16")
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")
# 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:
@ -279,6 +296,8 @@ class DiffusionLoraConfig:
lora_target_modules = targets,
max_grad_norm = float(self.max_grad_norm),
hf_token = token or None,
cache_variants = int(self.cache_variants),
compile_transformer = compile_transformer,
resolved_family = resolved_family,
)
@ -356,6 +375,81 @@ def _emit(on_event: Optional[EventCb], type_: str, **kw: Any) -> None:
on_event({"type": type_, "ts": time.time(), **kw})
def _plan_cache_variants(
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
loader maps onto its integer crop range. Uses its own rng stream so the training
loop's draws are untouched. Center-crop / no-flip collapse duplicate variants (a
center crop without flip is one variant no matter how many draws), so callers encode
each distinct variant exactly once. Pure (no torch) for CPU unit tests."""
crop_rng = random.Random(seed)
plan: list[list[tuple[float, float, bool]]] = []
for _ in range(max(0, num_images)):
variants: list[tuple[float, float, bool]] = []
for _ in range(max(1, cache_variants)):
u_left, u_top = crop_rng.random(), crop_rng.random()
flip = bool(random_flip and crop_rng.random() < 0.5)
if center_crop:
u_left = u_top = 0.5 # loader ignores the fractions for a center crop
key = (u_left, u_top, flip)
if key not in variants:
variants.append(key)
plan.append(variants)
return plan
def _apply_perf_flags(
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
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:
missing attributes on a CPU/other-vendor build are skipped."""
from core.inference.diffusion_speed import snapshot_backend_flags
snap: dict[str, Any] = {"flags": snapshot_backend_flags(), "matmul_precision": None}
if device != "cuda":
return snap
try:
import torch
snap["matmul_precision"] = torch.get_float32_matmul_precision()
if cfg.enable_tf32:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision("high")
if cudnn_benchmark:
torch.backends.cudnn.benchmark = True
except Exception: # noqa: BLE001 -- perf flags are never fatal
pass
return snap
def _restore_perf_flags(snap: Optional[dict]) -> None:
"""Undo ``_apply_perf_flags`` (the trainer subprocess is disposable, but in-process
callers -- tests, notebooks -- must not inherit mutated globals)."""
if not snap:
return
from core.inference.diffusion_speed import restore_backend_flags
restore_backend_flags(snap.get("flags"))
if snap.get("matmul_precision"):
try:
import torch
torch.set_float32_matmul_precision(snap["matmul_precision"])
except Exception: # noqa: BLE001 -- best-effort restore
pass
def _assert_trusted_base_model(base_model: str) -> None:
"""Gate the training base model the same way the inference backend gates non-GGUF loads:
a local path or a trusted repo (``unsloth/*`` or an allowlisted official base). This runs

43
studio/backend/core/training/diffusion_training_service.py
View file

@ -206,17 +206,25 @@ class DiffusionTrainingService:
self._pump.start()
return job_id
def stop(self) -> bool:
"""Request a clean stop (the trainer finishes the current step and saves a partial
adapter). Returns True if a stop was signalled, False if nothing was running."""
def stop(self, save: bool = True) -> bool:
"""Request a clean stop: the trainer finishes the current step, then either saves
a partial adapter (``save=True``, the default) or discards the run (``save=False``,
matching the LLM trainer's cancel). Returns True if a stop was signalled, False if
nothing was running."""
with self._lock:
if self._proc is None or not self._proc.is_alive() or self._stop_queue is None:
return False
try:
self._stop_queue.put(True)
# Bare True keeps the wire format older trainers expect; the dict form
# carries the no-save cancel flag the trainer's _check_stop understands.
self._stop_queue.put(True if save else {"save": False})
except Exception: # noqa: BLE001
return False
self._state["message"] = "Stop requested; finishing the current step..."
self._state["message"] = (
"Stop requested; finishing the current step and saving a partial adapter..."
if save
else "Cancel requested; finishing the current step (no adapter will be saved)..."
)
self._state["updated_at"] = time.time()
return True
@ -280,6 +288,25 @@ class DiffusionTrainingService:
s["num_images"] = ev.get("num_images")
elif etype == "model_load_completed":
s.update(in_model_load = False, message = "Training...")
elif etype == "preparing":
# A long precompute phase (e.g. the VAE latent cache) between model load and
# the first step; surfaced so the UI shows visible progress instead of a
# silent "Loading base model..." stall.
done, total = ev.get("done"), ev.get("total")
stage = str(ev.get("stage", "prepare")).replace("_", " ")
s.update(
status = "running",
in_model_load = True,
message = (
f"Preparing ({stage} {done}/{total})..."
if done is not None and total is not None
else f"Preparing ({stage})..."
),
)
elif etype == "warning":
# Non-fatal trainer notes (e.g. torch.compile falling back to eager); keep
# training state, surface the text.
s["message"] = str(ev.get("message", "warning"))
elif etype == "progress":
s.update(
status = "running",
@ -308,7 +335,11 @@ class DiffusionTrainingService:
status = "stopped" if ev.get("stopped") else "completed",
output_dir = ev.get("output_dir"),
lora_path = ev.get("lora_path"),
message = "Stopped (partial adapter saved)."
message = (
"Stopped (partial adapter saved)."
if ev.get("lora_path")
else "Stopped (no adapter saved)."
)
if ev.get("stopped")
else "Training complete.",
)

20
studio/backend/models/training.py
View file

@ -715,6 +715,26 @@ class DiffusionTrainingStartRequest(BaseModel):
random_flip: bool = Field(True)
caption_column: str = Field("text")
hf_token: Optional[str] = Field(None)
cache_latents: bool = Field(
True, description = "Precompute VAE latents once and free the VAE for the run"
)
cache_variants: int = Field(
4, ge = 1, le = 16, description = "Frozen crop/flip variants per image in the latent cache"
)
compile_transformer: Literal["off", "on", "auto"] = Field(
"auto", description = "Regional torch.compile of the transformer blocks"
)
enable_tf32: bool = Field(
True, description = "TF32 matmuls + cudnn autotuning (near-lossless speedup)"
)
class DiffusionTrainingStopRequest(BaseModel):
"""Optional body for stopping a diffusion training job. ``save`` mirrors the LLM
trainer's stop: True (default) exports the partial adapter, False cancels without
leaving one behind."""
save: bool = Field(True)
class DiffusionTrainingStartResponse(BaseModel):

13
studio/backend/routes/training.py
View file

@ -73,6 +73,7 @@ from models.training import (
DiffusionTrainingStartRequest,
DiffusionTrainingStartResponse,
DiffusionTrainingStatusResponse,
DiffusionTrainingStopRequest,
)
from models.responses import TrainingStopResponse, TrainingMetricsResponse
from pydantic import BaseModel as PydanticBaseModel
@ -1244,11 +1245,17 @@ async def start_diffusion_training(
@router.post("/diffusion/stop")
async def stop_diffusion_training(current_subject: str = Depends(get_current_subject)):
"""Request a clean stop of the running diffusion training job (partial adapter saved)."""
async def stop_diffusion_training(
body: Optional[DiffusionTrainingStopRequest] = None,
current_subject: str = Depends(get_current_subject),
):
"""Request a clean stop of the running diffusion training job. The optional body's
``save`` mirrors the LLM /stop: true (default, also for an empty POST) exports the
partial adapter, false cancels without saving one."""
from core.training.diffusion_training_service import get_diffusion_training_service
stopped = get_diffusion_training_service().stop()
save = body.save if body is not None else True
stopped = get_diffusion_training_service().stop(save = save)
return {"status": "stopping" if stopped else "idle"}

334
studio/backend/tests/test_diffusion_train_perf.py Normal file
View file

@ -0,0 +1,334 @@
# 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 diffusion training performance work.
Covers the new pure helpers and small policy functions that the perf PR adds:
the seed-deterministic latent-cache crop/flip plan, the per-family collate fns, the
index-based sigma gather, the new config validation + request-model fields, the
torch.compile policy, the stop save/cancel flag, and the ``preparing`` / ``warning``
service events. No GPU / model load: the collates and gathers run on CPU tensors, the
scheduler is default-initialised (no ``from_pretrained``), and the route/service tests
inject in-thread fakes exactly like ``test_diffusion_training.py``.
"""
from __future__ import annotations
import itertools
import pytest
import torch
from fastapi import FastAPI
from fastapi.testclient import TestClient
from auth.authentication import get_current_subject
from core.training.diffusion_dit_trainer import (
_flux_collate,
_gather_sigmas,
_qwen_collate,
_sample_timesteps,
_should_compile,
_zimage_collate,
)
from core.training.diffusion_train_common import (
DiffusionLoraConfig,
_apply_perf_flags,
_config_from_dict,
_plan_cache_variants,
_restore_perf_flags,
)
from core.training.diffusion_training_service import DiffusionTrainingService
from models.training import DiffusionTrainingStartRequest, DiffusionTrainingStopRequest
from routes.training import router as training_router
# A trainable SDXL base so DiffusionLoraConfig.normalized() resolves a family without a
# network call (resolve_trainable_family is pure name matching for this repo).
_SDXL = "stabilityai/stable-diffusion-xl-base-1.0"
def _cfg(**kw) -> DiffusionLoraConfig:
return DiffusionLoraConfig(base_model = _SDXL, data_dir = "d", output_dir = "o", **kw)
# ── _plan_cache_variants (pure, seed-deterministic) ───────────────────────────
def test_plan_cache_variants_deterministic_and_deduped():
# Same seed -> byte-identical plan (its own rng stream, so it is fully reproducible).
p1 = _plan_cache_variants(3, 4, center_crop = False, random_flip = True, seed = 123)
p2 = _plan_cache_variants(3, 4, center_crop = False, random_flip = True, seed = 123)
assert p1 == p2
assert len(p1) == 3
# cache_variants=1 -> exactly one variant per image.
p_one = _plan_cache_variants(3, 1, center_crop = False, random_flip = True, seed = 7)
assert [len(v) for v in p_one] == [1, 1, 1]
# A center crop with no flip collapses to a single distinct variant no matter how many
# draws are requested, and that variant is the fixed (0.5, 0.5, False) center.
p_cc = _plan_cache_variants(2, 8, center_crop = True, random_flip = False, seed = 7)
assert [len(v) for v in p_cc] == [1, 1]
assert p_cc[0][0] == (0.5, 0.5, False)
# A center crop WITH flip has at most two distinct variants (flip on/off; crop is fixed).
p_cf = _plan_cache_variants(2, 8, center_crop = True, random_flip = True, seed = 7)
assert all(len(v) <= 2 for v in p_cf)
# Every crop fraction is a valid unit fraction the loader can map onto its crop range.
for u_left, u_top, flip in itertools.chain.from_iterable(p1):
assert 0.0 <= u_left < 1.0
assert 0.0 <= u_top < 1.0
assert isinstance(flip, bool)
# ── 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)
]
pe, pooled, text_ids = _flux_collate(entries, "cpu", torch.float32)
assert pe.shape == (3, 512, 32)
assert pooled.shape == (3, 16)
assert text_ids.shape == (512, 3)
# Position ids stay float32 regardless of the requested weight dtype.
assert pe.dtype == torch.float32
assert pooled.dtype == torch.float32
assert text_ids.dtype == torch.float32
def test_qwen_collate_pads_and_masks():
dim = 8
# A short (mask=None) and a long (mask=ones) entry -> pad to the batch max and build the
# validity mask, with the padded tail of the short sample masked out.
short = (torch.randn(1, 5, dim), None)
long = (torch.randn(1, 9, dim), torch.ones(1, 9, dtype = torch.int64))
pe, mask = _qwen_collate([short, long], "cpu", torch.float32)
assert pe.shape == (2, 9, dim)
assert mask.shape == (2, 9)
assert torch.equal(mask[0, 5:], torch.zeros(4, dtype = mask.dtype))
# A single unpadded sample with a None mask keeps the legacy None mask (no behaviour delta).
pe1, mask1 = _qwen_collate([(torch.randn(1, 5, dim), None)], "cpu", torch.float32)
assert pe1.shape == (1, 5, dim)
assert mask1 is None
# A single sample pinned to a compile pad bucket must pad AND expose a mask so the padded
# positions are attended to as invalid.
pe2, mask2 = _qwen_collate([(torch.randn(1, 5, dim), None)], "cpu", torch.float32, pad_to = 16)
assert pe2.shape == (1, 16, dim)
assert mask2 is not None
assert torch.equal(mask2[0, 5:], torch.zeros(11, dtype = mask2.dtype))
def test_zimage_collate_list():
# Z-Image uses list I/O: the batch is one tuple carrying a list of per-sample tensors, each
# cast to the requested dtype.
entries = [(torch.randn(7, 2560),), (torch.randn(9, 2560),)]
out = _zimage_collate(entries, "cpu", torch.float32)
assert isinstance(out, tuple) and len(out) == 1
(caps,) = out
assert isinstance(caps, list) and len(caps) == 2
assert all(t.dtype == torch.float32 for t in caps)
# ── index-based sigma gather ──────────────────────────────────────────────────
def test_gather_sigmas_matches_search_based_gather():
from diffusers import FlowMatchEulerDiscreteScheduler
torch.manual_seed(0)
sched = FlowMatchEulerDiscreteScheduler() # default init, no from_pretrained / no network
timesteps, indices = _sample_timesteps(sched, 16, "cpu")
# The index path must return exactly what the old per-item timestep-matching search did.
schedule_timesteps = sched.timesteps.to("cpu")
step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
assert step_indices == indices.tolist()
sigma = _gather_sigmas(sched, indices, "cpu", torch.float32, 4)
assert sigma.ndim == 4
expected = sched.sigmas[step_indices].flatten()
while expected.ndim < 4:
expected = expected.unsqueeze(-1)
assert torch.equal(sigma, expected)
# ── config validation of the new perf fields ──────────────────────────────────
def test_config_validates_new_fields():
# Defaults normalize cleanly and carry the new perf fields through.
norm = _cfg().normalized()
assert norm.cache_variants == 4
assert norm.compile_transformer == "auto"
assert norm.enable_tf32 is True
assert norm.cache_latents is True
# cache_variants is bounded to 1..16 inclusive.
for bad in (0, 17):
with pytest.raises(ValueError):
_cfg(cache_variants = bad).normalized()
# An unknown compile mode is rejected.
with pytest.raises(ValueError):
_cfg(compile_transformer = "banana").normalized()
# compile_transformer is case/space-insensitive and stored lowered.
assert _cfg(compile_transformer = " ON ").normalized().compile_transformer == "on"
# The generic Studio dict path preserves the flags without inventing defaults.
cfg = _config_from_dict(
{
"base_model": _SDXL,
"data_dir": "d",
"output_dir": "o",
"enable_tf32": False,
"cache_latents": False,
}
)
assert cfg.enable_tf32 is False
assert cfg.cache_latents is False
# ── torch.compile policy ──────────────────────────────────────────────────────
def test_should_compile_policy():
# off never compiles, even on cuda.
assert _should_compile(_cfg(compile_transformer = "off"), False, "cuda") is False
# on always compiles on cuda.
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
# 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
# ── service stop save/cancel flag ─────────────────────────────────────────────
class _StopQueue:
"""Records what stop() puts on the wire (put-only for these tests)."""
def __init__(self) -> None:
self.items: list = []
def put(self, x) -> None:
self.items.append(x)
class _AliveProc:
def is_alive(self) -> bool:
return True
def test_service_stop_save_flag():
svc = DiffusionTrainingService()
# Nothing running -> stop is a no-op and returns False.
assert svc.stop() is False
# Attach a fake live proc + stop queue so stop() has a target.
svc._proc = _AliveProc()
q = _StopQueue()
svc._stop_queue = q
# save=False is the cancel path: the dict form {"save": False} goes on the queue.
assert svc.stop(save = False) is True
assert q.items[-1] == {"save": False}
# The default (save) path keeps the bare-True wire format.
assert svc.stop() is True
assert q.items[-1] is True
# ── 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}
)
st = svc.status()
assert st["status"] == "running"
assert st["in_model_load"] is True
assert "4/8" in st["message"]
svc._apply_event({"type": "warning", "message": "compile disabled"})
assert svc.status()["message"] == "compile disabled"
# A stop with no saved adapter reports the no-adapter message and the stopped status.
svc_no = DiffusionTrainingService()
svc_no._apply_event({"type": "complete", "stopped": True, "lora_path": None})
st_no = svc_no.status()
assert st_no["status"] == "stopped"
assert st_no["message"] == "Stopped (no adapter saved)."
# A stop that DID save a partial adapter reports the partial-adapter message.
svc_partial = DiffusionTrainingService()
svc_partial._apply_event(
{"type": "complete", "stopped": True, "lora_path": "/o/pytorch_lora_weights.safetensors"}
)
assert svc_partial.status()["message"] == "Stopped (partial adapter saved)."
# ── route: stop body forwards the save flag ───────────────────────────────────
class _FakeService:
"""Records the save flag the /diffusion/stop route forwards. A local copy of the
test_diffusion_training.py pattern so the two suites stay decoupled."""
def __init__(self) -> None:
self._running = True
self.stopped_with_save = None
def stop(self, save = True):
self.stopped_with_save = save
was = self._running
self._running = False
return was
@pytest.fixture
def client(monkeypatch):
fake = _FakeService()
monkeypatch.setattr(
"core.training.diffusion_training_service.get_diffusion_training_service", lambda: fake
)
app = FastAPI()
app.include_router(training_router, prefix = "/api/train")
app.dependency_overrides[get_current_subject] = lambda: "test-user"
c = TestClient(app)
c._fake = fake # type: ignore[attr-defined]
return c
def test_route_stop_save_body(client):
# An explicit {"save": false} body forwards save=False to the service.
r = client.post("/api/train/diffusion/stop", json = {"save": False})
assert r.status_code == 200, r.text
assert client._fake.stopped_with_save is False
# A body-less POST defaults to save=True.
r2 = client.post("/api/train/diffusion/stop")
assert r2.status_code == 200, r2.text
assert client._fake.stopped_with_save is True
# ── request models: new perf fields + stop schema ─────────────────────────────
def test_request_models_new_fields():
req = DiffusionTrainingStartRequest(base_model = "b", data_dir = "d", output_dir = "o")
assert req.cache_latents is True
assert req.cache_variants == 4
assert req.compile_transformer == "auto"
assert req.enable_tf32 is True
# cache_variants is validated against its 1..16 bound by pydantic.
with pytest.raises(Exception):
DiffusionTrainingStartRequest(
base_model = "b", data_dir = "d", output_dir = "o", cache_variants = 32
)
# The stop request defaults to saving a partial adapter.
assert DiffusionTrainingStopRequest().save is True
# ── perf flags round-trip on cpu ──────────────────────────────────────────────
def test_perf_flags_cpu_roundtrip():
# On a cpu device (or a torch build without cuda), applying the perf flags is a no-op
# snapshot path and restoring it must not raise.
snap = _apply_perf_flags(_cfg(), "cpu")
assert isinstance(snap, dict)
_restore_perf_flags(snap) # no exception

4
studio/backend/tests/test_diffusion_training.py
View file

@ -210,6 +210,7 @@ class _FakeService:
def __init__(self):
self._running = False
self.started_with = None
self.stopped_with_save = None
# Extra keys merged into status() so a test can inject metric history / perf fields.
self.status_extra: dict = {}
@ -218,7 +219,8 @@ class _FakeService:
self._running = True
return "job-123"
def stop(self):
def stop(self, save = True):
self.stopped_with_save = save
was = self._running
self._running = False
return was