From da3a79468ea2271fdeaa7332b62c18d4e5dcb9a4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 07:49:30 +0000 Subject: [PATCH 1/2] Use permutation-cycle index sampling in diffusion trainers and guard non-object run records Replace the with-replacement per-batch index draw in the SDXL and DiT LoRA trainers with a shared PermutationBatchSampler that visits every image once per cycle before repeating, so short runs cover the whole dataset. The sampler reshuffles from the run's rng so the index stream stays seed-deterministic. Guard the diffusion run detail route against a valid-JSON non-object record, which previously raised TypeError and returned a 500; it now 404s like the list path's shape check. Add regression tests for both. --- .../core/training/diffusion_dit_trainer.py | 7 ++- .../core/training/diffusion_lora_trainer.py | 9 +++- .../core/training/diffusion_train_common.py | 41 +++++++++++++++ studio/backend/routes/training.py | 6 ++- .../backend/tests/test_diffusion_training.py | 52 +++++++++++++++++++ 5 files changed, 112 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index cbaf2f70f2..48e632ae5c 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -49,6 +49,7 @@ from core.training.diffusion_train_common import ( _restore_perf_flags, discover_image_caption_pairs, has_functional_torchao, + PermutationBatchSampler, repo_is_prequantized, resolve_train_steps, ) @@ -1158,6 +1159,10 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto transformer.train() n_images = len(image_paths) batch_size = cfg.train_batch_size + # Permutation-cycle index sampler (shared with the SDXL trainer): visits every image once + # per cycle before repeating, so a short run covers the whole dataset instead of the old + # with-replacement draw. Uses the loop's own rng to stay seed-deterministic. + index_sampler = PermutationBatchSampler(n_images, rng) stopped = False running_loss = 0.0 peak_gb = 0.0 @@ -1177,7 +1182,7 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto optimizer.zero_grad(set_to_none = True) step_loss = 0.0 for _ in range(cfg.gradient_accumulation_steps): - idxs = [rng.randrange(n_images) for _ in range(batch_size)] + idxs = index_sampler.next_batch(batch_size) if latent_cache is not None: latents = _sample_cached_latents( latent_cache, idxs, variant_rng, device, weight_dtype diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index b89ad350a7..dcf2812d82 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -59,6 +59,7 @@ from core.training.diffusion_train_common import ( # noqa: F401 _restore_perf_flags, discover_image_caption_pairs, get_trainer, + PermutationBatchSampler, resolve_train_steps, ) @@ -436,8 +437,14 @@ def run_diffusion_lora_training( _emit(on_event, "model_load_completed") + # Permutation-cycle index sampler (shared with the DiT trainer): each dataset image is + # visited once per cycle before any repeat, so a short run does not leave part of a + # small dataset unseen. Draws from the loop's own rng so the sequence stays + # seed-deterministic. + index_sampler = PermutationBatchSampler(len(pairs), rng) + def _next_batch() -> tuple[list[int], list[str], list[str]]: - idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs))) + idx = index_sampler.next_batch(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] diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index db3819569c..5599d31962 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -431,6 +431,47 @@ def resolve_train_steps(cfg: "DiffusionLoraConfig", n_images: int) -> int: return cfg.train_steps +class PermutationBatchSampler: + """Yields batch indices as consecutive slices of a reshuffled permutation of + ``range(n)``, so every index is visited exactly once per cycle before any repeats -- + an epoch-style full pass instead of the with-replacement draw that leaves part of a + small dataset unseen at low step counts (num_epochs converts to a step budget, but the + per-batch index draw is what decides coverage). When a cycle is exhausted the order is + reshuffled from the run's own ``rng`` so the index stream stays seed-deterministic and + each cycle differs. + + Both trainers share this so the SDXL ``_next_batch`` path and the DiT per-sample draw + select indices the same way. Only the index selection changes (with-replacement -> + permutation cycles); step count and batch shapes are unchanged. + """ + + def __init__(self, n: int, rng: random.Random) -> None: + if n <= 0: + raise ValueError("PermutationBatchSampler needs at least one item") + self._n = n + self._rng = rng + self._order: list[int] = [] + self._pos = 0 + + def _reshuffle(self) -> None: + self._order = list(range(self._n)) + self._rng.shuffle(self._order) + self._pos = 0 + + def next_batch(self, k: int) -> list[int]: + # k may exceed n (batch larger than the dataset): the permutation is refilled across + # as many cycles as needed so the caller always gets exactly k indices and the batch + # never shrinks, matching the old sampler's fixed batch shape. + out: list[int] = [] + while len(out) < k: + if self._pos >= len(self._order): + self._reshuffle() + take = min(k - len(out), len(self._order) - self._pos) + out.extend(self._order[self._pos : self._pos + take]) + self._pos += take + return out + + def discover_image_caption_pairs( data_dir: str | os.PathLike[str], *, diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 35d6921605..bc131ad413 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1337,7 +1337,11 @@ async def get_diffusion_training_run( from core.training.diffusion_training_service import get_diffusion_run rec = get_diffusion_run(job_id) - if rec is None: + # A valid-JSON file that is not an object (a truncated / hand-edited [] record) would make + # DiffusionTrainingRunDetail(**rec) raise TypeError -- not the ValidationError caught below + # -- and 500 the endpoint. Treat any non-dict record as absent, matching the list route's + # shape check. + if not isinstance(rec, dict): raise HTTPException(status_code = 404, detail = "No such training run.") try: return DiffusionTrainingRunDetail(**rec) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 40db06e945..7af3ddd015 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -427,6 +427,45 @@ def test_config_from_dict_epoch_mode_drops_max_steps_sentinel(): assert cfg_explicit.train_steps == 25 +def test_permutation_sampler_covers_dataset_once_per_cycle(): + # Every index must appear exactly once per cycle before any repeat (epoch-style pass), + # so a short run over a small dataset never leaves images unseen the way the old + # with-replacement draw did. Consecutive cycles must be reshuffled (differ). + import random + + from core.training.diffusion_train_common import PermutationBatchSampler + + n = 100 + sampler = PermutationBatchSampler(n, random.Random(0)) + + # Draw exactly one cycle in batches of 3 (n not divisible by the batch, so a batch spans + # the cycle boundary); the first n indices must be a permutation of range(n). + drawn: list[int] = [] + while len(drawn) < n: + drawn.extend(sampler.next_batch(3)) + first_cycle = drawn[:n] + assert sorted(first_cycle) == list(range(n)) # each index once, none missing + + # The next full cycle is also a permutation, and it is reshuffled (order differs). + fresh = PermutationBatchSampler(n, random.Random(0)) + cycle_a = fresh.next_batch(n) + cycle_b = fresh.next_batch(n) + assert sorted(cycle_a) == list(range(n)) + assert sorted(cycle_b) == list(range(n)) + assert cycle_a != cycle_b # cycles are reshuffled, not repeated in the same order + + # A seed replays the exact index stream (determinism for reproducible runs). + replay = PermutationBatchSampler(n, random.Random(0)) + assert replay.next_batch(n) == cycle_a + + # A batch larger than the dataset refills across cycles so it never shrinks (batch shape + # preserved), even though it must then repeat indices within the batch. + big = PermutationBatchSampler(4, random.Random(1)) + batch = big.next_batch(10) + assert len(batch) == 10 + assert set(batch) == {0, 1, 2, 3} + + def test_route_start_accepts_zero_max_grad_norm(client): # 0 is the documented "disable clipping" value (the trainer skips clip_grad_norm_); # the request model must not reject it. @@ -937,3 +976,16 @@ def test_runs_route_tolerates_bad_field_record(client, _isolated_runs_dir): assert r.status_code == 200, r.text adapters = [x["adapter"] for x in r.json()["runs"]] assert adapters == ["good"] # the bad-field record was skipped, the good one remained + + +def test_run_detail_route_non_object_record_is_404(client, _isolated_runs_dir): + # A valid-JSON but non-object record (a truncated / hand-edited [] file named with a real + # job id) makes DiffusionTrainingRunDetail(**rec) raise TypeError, not ValidationError; the + # detail route must shape-check like the list path and 404 instead of 500. + import json + + job_id = "a" * 32 + (_isolated_runs_dir / f"{job_id}.json").write_text(json.dumps([])) + + r = client.get(f"/api/train/diffusion/runs/{job_id}") + assert r.status_code == 404, r.text From 8c00f81a5d5e690929983c30ada09c94cf5fdc0a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 07:53:12 +0000 Subject: [PATCH 2/2] Size-gate the automatic diffusion latent cache The latent cache holds two fp32 posterior tensors per crop/flip variant per image, pinned on CUDA hosts, so datasets with thousands of images can exhaust host or pinned memory with no fallback. Estimate the cache size from the first real encoded latent and fall back to per-step VAE encoding when it exceeds a 4 GiB budget. UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE bypasses the gate; the existing UNSLOTH_DIFFUSION_NO_LATENT_CACHE opt-out is unchanged. --- .../core/training/diffusion_dit_trainer.py | 56 ++++++++--- .../core/training/diffusion_lora_trainer.py | 50 ++++++++-- .../core/training/diffusion_train_common.py | 36 ++++++++ .../tests/test_diffusion_train_perf.py | 92 +++++++++++++++++++ 4 files changed, 213 insertions(+), 21 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index c27f1f6fef..14a4a3a71c 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -40,10 +40,13 @@ from core.training.diffusion_train_common import ( DEFAULT_LORA_TARGETS, DiffusionLoraConfig, EventCb, + LATENT_CACHE_OVER_BUDGET, StopCb, _apply_perf_flags, _assert_trusted_base_model, _emit, + _latent_cache_forced, + _latent_cache_over_budget, _plan_cache_variants, _publish_to_lora_catalog, _restore_perf_flags, @@ -680,6 +683,9 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev cache: list[list[tuple]] = [] total = len(image_paths) + total_variants = sum(len(v) for v in plan) + forced = _latent_cache_forced() + gated = False for i, path in enumerate(image_paths): variants = [] for u_left, u_top, flip in plan[i]: @@ -691,7 +697,30 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev .to(device) ) a, b = spec.encode_latent_stats(vae, px) - variants.append((_hold(a), _hold(b))) + a, b = _hold(a), _hold(b) + if not forced and not gated: + # Size-gate the automatic cache off the first REAL encoded variant, before + # building the rest: packed 16-channel DiT latents x variants x images of two + # fp32 tensors can exhaust host/pinned RAM. Over budget we bail with the VAE + # still resident so the loop encodes latents per step instead. ``b`` is None + # for a deterministic-latent family, so only ``a`` contributes bytes there. + per_variant = a.numel() * a.element_size() + if b is not None: + per_variant += b.numel() * b.element_size() + if _latent_cache_over_budget(per_variant, total_variants): + _emit( + on_event, + "warning", + message = ( + "Latent cache disabled: estimated " + f"{per_variant * total_variants / 1024 ** 3:.1f} GiB over the " + "budget; encoding latents per step instead. Set " + "UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE=1 to keep it." + ), + ) + return LATENT_CACHE_OVER_BUDGET + gated = True + variants.append((a, 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) @@ -893,7 +922,11 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto 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 + if latent_cache is LATENT_CACHE_OVER_BUDGET: + # The estimated cache exceeded the host-memory budget; keep the VAE resident and + # fall through to the in-loop encode path (latent_cache stays None). + latent_cache = None + elif latent_cache is None: # stopped during the cache build; nothing trained yet _emit( on_event, "complete", @@ -903,15 +936,16 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto 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() + else: + 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) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index b7c65a0ac4..8c3a79afaf 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -47,12 +47,15 @@ from core.training.diffusion_train_common import ( # noqa: F401 EventCb, StopCb, DiffusionLoraConfig, + LATENT_CACHE_OVER_BUDGET, _apply_perf_flags, _assert_trusted_base_model, _coerce_gradient_checkpointing, _config_from_dict, _CONFIG_ALIASES, _emit, + _latent_cache_forced, + _latent_cache_over_budget, _plan_cache_variants, _publish_to_lora_catalog, _restore_perf_flags, @@ -202,6 +205,9 @@ def _build_sdxl_latent_cache( cache: list[list[tuple]] = [] total = len(image_paths) + total_variants = sum(len(v) for v in plan) + forced = _latent_cache_forced() + gated = False for i, path in enumerate(image_paths): variants = [] for u_left, u_top, flip in plan[i]: @@ -213,6 +219,25 @@ def _build_sdxl_latent_cache( dist = vae.encode(pixel_values).latent_dist a = _hold(dist.mean * vae_scale) b = _hold(dist.std * vae_scale) + if not forced and not gated: + # Size-gate the automatic cache off the first REAL encoded variant, before + # building the rest: thousands of images x variants of two fp32 tensors can + # exhaust host/pinned RAM with no fallback. Over budget we bail with the VAE + # still resident so the loop encodes latents per step instead. + per_variant = a.numel() * a.element_size() + b.numel() * b.element_size() + if _latent_cache_over_budget(per_variant, total_variants): + _emit( + on_event, + "warning", + message = ( + "Latent cache disabled: estimated " + f"{per_variant * total_variants / 1024 ** 3:.1f} GiB over the " + "budget; encoding latents per step instead. Set " + "UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE=1 to keep it." + ), + ) + return LATENT_CACHE_OVER_BUDGET + gated = True variants.append((a, b, tuple(time_ids))) cache.append(variants) if (i + 1) % 4 == 0 or i + 1 == total: @@ -404,7 +429,11 @@ def run_diffusion_lora_training( on_event, _check_stop, ) - if latent_cache is None: # stopped during the cache build; nothing trained yet + if latent_cache is LATENT_CACHE_OVER_BUDGET: + # The estimated cache exceeded the host-memory budget; keep the VAE resident + # and fall through to the in-loop encode path (latent_cache stays None). + latent_cache = None + elif latent_cache is None: # stopped during the cache build; nothing trained yet out_dir = Path(cfg.output_dir).expanduser() _emit( on_event, @@ -415,15 +444,16 @@ def run_diffusion_lora_training( 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() + else: + 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) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index fba15669fd..01bce2e6b9 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -411,6 +411,42 @@ def _plan_cache_variants( return plan +# Host-memory budget for the AUTOMATIC latent cache. The cache holds two fp32 posterior +# tensors (mean/std, VAE scale folded in) per crop/flip variant per image, pinned on a CUDA +# host. At 1024px an SDXL variant is ~0.5 MiB and a 16-channel DiT variant several times +# that, so a few thousand images x cache_variants can exhaust host or pinned RAM with no +# fallback. Over this budget the default falls back to per-step VAE encoding. A fixed +# constant (rather than a psutil RAM fraction) keeps the gate dependency-free and identical +# across hosts; it is deliberately conservative, well under a typical training host's RAM. +_LATENT_CACHE_BUDGET_BYTES = 4 * 1024 ** 3 # 4 GiB + +# Returned by the cache builders when the estimated cache exceeds the budget: the caller +# keeps the VAE resident and encodes each step's latents in-loop. A distinct sentinel from +# ``None`` (which means a stop was requested mid-build) so the two are not conflated. +LATENT_CACHE_OVER_BUDGET: Any = object() + + +def _latent_cache_forced() -> bool: + """The user explicitly forced the latent cache on, bypassing the size gate. This is the + explicit opt-in counterpart to ``UNSLOTH_DIFFUSION_NO_LATENT_CACHE`` (the explicit + opt-out); only the automatic default is size-gated, so an explicit choice is honoured + verbatim in either direction.""" + return os.environ.get("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", "") in ("1", "true") + + +def _latent_cache_over_budget( + per_variant_bytes: int, total_variants: int, budget_bytes: Optional[int] = None +) -> bool: + """True when a cache of ``total_variants`` entries, each two fp32 tensors totalling + ``per_variant_bytes``, is estimated to exceed ``budget_bytes``. ``per_variant_bytes`` is + measured from a real encoded latent, so the estimate tracks the actual per-family tensor + shape (SDXL 4-channel vs. a packed 16-channel DiT latent) rather than a guess. The budget + is read from the module constant at call time when not given, so tests can override it.""" + if budget_bytes is None: + budget_bytes = _LATENT_CACHE_BUDGET_BYTES + return per_variant_bytes * max(0, total_variants) > budget_bytes + + def _apply_perf_flags( cfg: "DiffusionLoraConfig", device: str, diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py index 4fd8c85730..8fa1cc2e23 100644 --- a/studio/backend/tests/test_diffusion_train_perf.py +++ b/studio/backend/tests/test_diffusion_train_perf.py @@ -32,11 +32,16 @@ from core.training.diffusion_dit_trainer import ( ) from core.training.diffusion_train_common import ( DiffusionLoraConfig, + LATENT_CACHE_OVER_BUDGET, _apply_perf_flags, _config_from_dict, + _latent_cache_forced, + _latent_cache_over_budget, _plan_cache_variants, _restore_perf_flags, ) +import core.training.diffusion_lora_trainer as sdxl_trainer +import core.training.diffusion_train_common as train_common from core.training.diffusion_training_service import DiffusionTrainingService from models.training import DiffusionTrainingStartRequest, DiffusionTrainingStopRequest from routes.training import router as training_router @@ -372,3 +377,90 @@ def test_perf_flags_tf32_off_clears_flags(): torch.get_float32_matmul_precision(), ) assert after == before + + +# ── latent cache size gate ──────────────────────────────────────────────────── +class _FakeLatentDist: + def __init__(self, shape): + self.mean = torch.zeros(shape, dtype = torch.float32) + self.std = torch.ones(shape, dtype = torch.float32) + + +class _FakeEncoded: + def __init__(self, shape): + self.latent_dist = _FakeLatentDist(shape) + + +class _FakeVae: + # Minimal VAE stand-in: encode() returns a posterior of the requested latent shape so the + # builder measures a real per-variant byte size without a model load or image files. + def __init__(self, shape): + self._shape = shape + + def encode(self, pixel_values): + return _FakeEncoded(self._shape) + + +def _fake_planned_loader(path, resolution, center_crop, u_left, u_top, flip): + # The fake VAE ignores pixels; return a valid tensor + square SDXL time_ids. + tensor = torch.zeros(3, resolution, resolution, dtype = torch.float32) + return tensor, (resolution, resolution, 0, 0, resolution, resolution) + + +def _build_fake_sdxl_cache(monkeypatch, num_images, latent_shape): + # center_crop + no flip collapses to one variant per image, so total_variants == num_images. + monkeypatch.setattr(sdxl_trainer, "_load_image_tensor_planned", _fake_planned_loader) + cfg = _cfg(cache_variants = 1, center_crop = True, random_flip = False).normalized() + return sdxl_trainer._build_sdxl_latent_cache( + _FakeVae(latent_shape), + 1.0, + [f"img{i}.png" for i in range(num_images)], + cfg, + "cpu", + torch.float32, + None, + lambda: False, + ) + + +def test_latent_cache_over_budget_boundary(): + # 32 bytes per variant x 4 variants = 128 bytes; exactly at budget is not "over". + assert _latent_cache_over_budget(32, 4, budget_bytes = 200) is False + assert _latent_cache_over_budget(32, 4, budget_bytes = 128) is False + assert _latent_cache_over_budget(32, 4, budget_bytes = 127) is True + # An empty plan can never overflow. + assert _latent_cache_over_budget(1_000_000, 0, budget_bytes = 1) is False + + +def test_latent_cache_forced_env(monkeypatch): + monkeypatch.delenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", raising = False) + assert _latent_cache_forced() is False + monkeypatch.setenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", "1") + assert _latent_cache_forced() is True + + +def test_sdxl_cache_built_under_budget(monkeypatch): + # Default (4 GiB) budget: a handful of tiny latents fits, so the full cache is returned. + monkeypatch.delenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", raising = False) + cache = _build_fake_sdxl_cache(monkeypatch, num_images = 3, latent_shape = (1, 4, 8, 8)) + assert cache is not LATENT_CACHE_OVER_BUDGET and cache is not None + assert len(cache) == 3 + assert all(len(variants) == 1 for variants in cache) + + +def test_sdxl_cache_gated_over_budget(monkeypatch): + # A budget below one variant forces the gate on the first encode: the sentinel is returned + # so the caller keeps the VAE resident and encodes per step. + monkeypatch.delenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", raising = False) + monkeypatch.setattr(train_common, "_LATENT_CACHE_BUDGET_BYTES", 8) + cache = _build_fake_sdxl_cache(monkeypatch, num_images = 3, latent_shape = (1, 4, 8, 8)) + assert cache is LATENT_CACHE_OVER_BUDGET + + +def test_sdxl_cache_force_bypasses_gate(monkeypatch): + # An explicit force-on must be honoured verbatim even when the estimate is over budget. + monkeypatch.setenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", "1") + monkeypatch.setattr(train_common, "_LATENT_CACHE_BUDGET_BYTES", 8) + cache = _build_fake_sdxl_cache(monkeypatch, num_images = 3, latent_shape = (1, 4, 8, 8)) + assert cache is not LATENT_CACHE_OVER_BUDGET and cache is not None + assert len(cache) == 3