From 908cb86b8d53f194b82e876ea75a7a2b4274f0f9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 14:04:47 +0000 Subject: [PATCH] Do not widen the dense prefetch under a forced-offload memory policy _dense_quant_prefetch_needed widened the prefetch to pull the base repo's transformer/ shards whenever a dense-quant candidate resolved, but balanced/low_vram (and the legacy cpu_offload flag) force load_pipeline onto offload unconditionally in plan_diffusion_memory, so its re-plan never flips to OFFLOAD_NONE and the dense build never runs. The offloaded GGUF path then never touches those shards, so the widened prefetch only wastes a multi-GB download, and a disk-full on that begin_load pull has no GGUF fallback (unlike the in-load_pipeline dense failure). Mirror plan_diffusion_memory's definite-offload gates so the prefetch stays scoped. --- studio/backend/core/inference/diffusion.py | 16 +++++++++ .../backend/tests/test_diffusion_backend.py | 33 +++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index ac61efccf8..e02d90b87c 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -44,12 +44,15 @@ from .diffusion_device import ( ) from .diffusion_krea2 import KREA2_FAMILY_NAME, load_krea2_pipeline from .diffusion_memory import ( + MEMORY_MODE_BALANCED, + MEMORY_MODE_LOW_VRAM, OFFLOAD_NONE, apply_memory_plan, estimate_gguf_resident_mib, estimate_image_runtime_mib, estimate_safetensors_dense_mib, file_size_mib, + normalize_memory_mode, plan_diffusion_memory, snapshot_device_memory, ) @@ -441,6 +444,19 @@ class DiffusionBackend: if speed is not None and str(speed).strip().lower() == SPEED_OFF: return False try: + # A definite-offload memory policy forces load_pipeline onto offload regardless of the + # dense candidate's smaller footprint, so its re-plan never flips to OFFLOAD_NONE and + # the dense build never runs. balanced -> OFFLOAD_GROUP and low_vram -> OFFLOAD_MODEL are + # set unconditionally in plan_diffusion_memory; the legacy cpu_offload flag forces + # OFFLOAD_MODEL when no memory_mode overrides it. In those cases the GGUF path runs + # offloaded and never touches the base transformer/ shards, so widening the prefetch only + # wastes a multi-GB download -- and a disk-full on that begin_load pull has NO GGUF + # fallback (unlike the in-load_pipeline dense failure). Mirror those offload gates here. + mm = normalize_memory_mode(kwargs.get("memory_mode")) + if mm in (MEMORY_MODE_BALANCED, MEMORY_MODE_LOW_VRAM): + return False + if mm is None and kwargs.get("cpu_offload"): + return False target = self._resolve_device_target(fam) # Only widen the prefetch when the loader would actually take the dense path: resolve # the SAME dense-quant candidate load_pipeline re-plans against, which also checks the diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 0fbfe9f32b..119c362ea7 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2127,6 +2127,39 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): # UNSET defaults to the hardware ladder (Dtype default-auto) -> widens, threading auto. assert backend._dense_quant_prefetch_needed(fam, {}) is True assert seen[-1] == "auto" + # A definite-offload memory policy forces load_pipeline onto offload regardless of the dense + # candidate's smaller footprint, so the dense build never runs and the widened prefetch would + # download base transformer/ shards the offloaded GGUF path never uses (and a disk-full there + # has no GGUF fallback). balanced / low_vram (and the legacy cpu_offload flag when no + # memory_mode overrides it) must NOT widen, even though the candidate itself is dense-viable. + before = len(seen) + assert ( + backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8", "memory_mode": "balanced"}) + is False + ) + assert ( + backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8", "memory_mode": "low_vram"}) + is False + ) + assert ( + backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8", "cpu_offload": True}) + is False + ) + # The gate short-circuits BEFORE resolving the candidate (no wasted resolve). + assert len(seen) == before + # An explicit memory_mode still consulting the candidate: fast/auto can flip resident, so they + # widen when the candidate is dense-viable (memory_mode="fast" does not force offload). + assert ( + backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8", "memory_mode": "fast"}) + is True + ) + # A cpu_offload flag is overridden by an explicit resident memory_mode, so it still widens. + assert ( + backend._dense_quant_prefetch_needed( + fam, {"transformer_quant": "fp8", "memory_mode": "fast", "cpu_offload": True} + ) + is True + ) # An explicit off pins running the GGUF as-is -> never widen (mode resolves to None first). assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "none"}) is False # An explicit Speed="off" (bit-exact) load suppresses the dense path -> never widen.