From 092fc91e081a2035de11dab8cb60f9211b6a6453 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 08:48:13 +0000 Subject: [PATCH] diffusion: don't double-count the prefetched transformer in the dense re-plan The dense-quant re-plan passes transformer_resident_override_mib (the bf16 build peak) AND computes companions via _companion_cache_bytes(base), which sums every flat blob in the HF cache. Because the dense path prefetches the base transformer/ shards into that same cache before load_pipeline runs, the transformer is counted twice, inflating the footprint (~44 GB instead of ~20 GB in the reproduction) and wrongly forcing offload for models that fit resident -- the case this path exists to enable. Add companion_override_mib and pass the auto-policy's own text-encoder plus VAE estimate on the re-plan so the cache (with its prefetched transformer) is not read for this artifact. --- studio/backend/core/inference/diffusion.py | 24 ++++++++-- .../backend/tests/test_diffusion_backend.py | 47 +++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 9f71789265..44030925a8 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1059,6 +1059,10 @@ class DiffusionBackend: kind = kind, repo_id = repo_id, transformer_resident_override_mib = (candidate.transient_transformer_mib), + # The dense path prefetches the base transformer/ shards into the + # cache _companion_cache_bytes reads; pass the auto-policy's own + # companion estimate so the re-plan does not double-count them. + companion_override_mib = candidate.companions_mib, ) if replanned.offload_policy == OFFLOAD_NONE: quant_plan = replanned @@ -1555,6 +1559,7 @@ class DiffusionBackend: kind: str = "gguf", repo_id: Optional[str] = None, transformer_resident_override_mib: Optional[int] = None, + companion_override_mib: Optional[int] = None, ): """Build the memory plan for this load: snapshot free device memory and estimate the model's resident footprint, then let the planner pick an @@ -1567,7 +1572,11 @@ class DiffusionBackend: pipeline is one cached download (transformer + companions), already compressed. ``transformer_resident_override_mib`` replaces the file-size transformer estimate when the loader is planning for a DIFFERENT artifact than the file on disk (the - dense transformer-quant candidate, whose footprint the auto-policy estimates).""" + dense transformer-quant candidate, whose footprint the auto-policy estimates); + ``companion_override_mib`` likewise replaces the cached companion total on that + re-plan, so the base repo's PREFETCHED transformer/ shards -- which land in the + same blob cache _companion_cache_bytes sums -- are not counted as companions on + top of transformer_resident_override_mib (a double-count of the transformer).""" device_memory = snapshot_device_memory(target) if kind == "pipeline": # The whole repo (transformer + companions) is one cached download; the @@ -1600,8 +1609,17 @@ class DiffusionBackend: # LOCAL diffusers base -- the on-disk component weights (the blob cache is # empty for a local path, which would otherwise fold multi-GB companions to 0 # and let auto planning pick a resident placement that OOMs). - companion = self._companion_cache_bytes(base) - companion_mib = int(companion // (1024 * 1024)) if companion else None + if companion_override_mib is not None: + # Re-planning the dense transformer-quant candidate: the dense path + # prefetches the base repo's transformer/ shards into the SAME blob cache + # _companion_cache_bytes sums, so reading it here would count the + # transformer AGAIN on top of transformer_resident_override_mib and make + # the resident quant plan look far too large. Use the auto-policy's own + # companion (text-encoder + VAE) estimate for this artifact instead. + companion_mib = companion_override_mib + else: + companion = self._companion_cache_bytes(base) + companion_mib = int(companion // (1024 * 1024)) if companion else None model_dense_mib = None if transformer_resident is not None: model_dense_mib = transformer_resident + (companion_mib or 0) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 1d18745d2a..f134521d71 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2120,6 +2120,53 @@ def test_companion_cache_bytes_local_dir_excludes_transformer(tmp_path): assert total == 150 # vae + text_encoder only; transformer/ and json excluded +def test_plan_memory_dense_replan_does_not_double_count_prefetched_transformer(monkeypatch): + # Re-planning the dense transformer-quant candidate: the dense path prefetches the + # base repo's transformer/ shards into the SAME blob cache _companion_cache_bytes + # sums. If the re-plan read that cache it would count the transformer TWICE (once as + # transformer_resident_override_mib, once as a "companion") and force offload even + # when the quantised artifact fits resident. The re-plan must use the auto-policy's + # companion estimate instead. Here the cache is stubbed to the transformer-inflated + # value; the plan must still stay resident. + from core.inference import diffusion as dmod + from core.inference.diffusion_memory import OFFLOAD_NONE, DeviceMemory + + backend = DiffusionBackend() + target = types.SimpleNamespace( + device = "cuda", backend = "cuda", supports_model_cpu_offload = True + ) + # 40 GiB discrete card, 40000 MiB free: comfortably fits transformer + real + # companions + headroom, but NOT a second copy of the bf16 transformer. + monkeypatch.setattr( + dmod, + "snapshot_device_memory", + lambda t: DeviceMemory("cuda", "cuda", "discrete_vram", 40000, 40960), + ) + monkeypatch.setattr(dmod, "estimate_image_runtime_mib", lambda **kw: 4000) + # The cache is inflated by the prefetched bf16 transformer (~24000) on top of the + # ~8000 real companions; if the re-plan consulted it the plan would offload. + monkeypatch.setattr( + DiffusionBackend, + "_companion_cache_bytes", + staticmethod(lambda base: (8000 + 24000) * 1024 * 1024), + ) + fam = types.SimpleNamespace(name = "z-image") + plan = backend._plan_memory( + target, + None, + "org/base", + fam, + None, + False, + kind = "gguf", + transformer_resident_override_mib = 12000, # int8 candidate transient (~half bf16) + companion_override_mib = 8000, # auto-policy text-encoder + VAE estimate + ) + # 12000 + 8000 + 4000 + 2048 overhead = 26048 MiB, fits the ~36 GiB budget. + # A double-count (12000 + [8000+24000] + ...) would have exceeded it and offloaded. + assert plan.offload_policy == OFFLOAD_NONE + + def test_reset_step_cache_helper_is_best_effort(): # Calls the transformer's reset hook when present. calls = []