diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 82b65ee5ad..4da16d764a 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -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: diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index deef4d03ce..ceb3802ea7 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -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) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 9458c4e5c9..f1989f6305 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -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 diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 2ad30e8347..fa44cceabe 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -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.", ) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 2479571357..8dafb7dfa4 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -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): diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index a5b4c3b06e..271abebd78 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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"} diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py new file mode 100644 index 0000000000..b8d805ead2 --- /dev/null +++ b/studio/backend/tests/test_diffusion_train_perf.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 886a3e5d4f..f3b7d2b58a 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -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