From a4277a01e49bb63dabbbe4e67fb55f45c098f5cc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 03:51:53 +0000 Subject: [PATCH] Honor memory_mode over legacy cpu_offload and prefetch dense-quant transformer shards plan_diffusion_memory only applies the legacy cpu_offload override when no memory_mode was supplied, matching the documented API contract that memory_mode overrides cpu_offload when set; an explicit fast request now stays resident even if the old flag is also enabled. The transformer-quant dense path fetches the base repo's transformer/ shards inside the locked finalize phase, where unload and cancellation cannot preempt the multi-GB download. The load worker now widens the preemptible prefetch to include those shards when that path can actually run: quant requested and supported for the device, scheme resolvable, and no pre-quantized checkpoint shortcutting the dense build. --- studio/backend/core/inference/diffusion.py | 56 +++++++++++++++++-- .../core/inference/diffusion_memory.py | 12 +++- .../backend/tests/test_diffusion_backend.py | 46 +++++++++++++++ studio/backend/tests/test_diffusion_memory.py | 15 +++++ 4 files changed, 122 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 6323f3dd52..63914103bd 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -219,6 +219,34 @@ class DiffusionBackend: return hf_hub_download(repo_id, gguf_filename, token = hf_token) + def _dense_quant_prefetch_needed(self, fam: DiffusionFamily, kwargs: dict) -> bool: + """True when ``load_pipeline`` may take the dense transformer-quant path, so + the prefetch should also pull the base repo's ``transformer/`` shards. + + Those shards are excluded from the prefetch by default (the GGUF supplies + the transformer), but ``_load_dense_quant_pipeline`` fetches them with + ``from_pretrained(subfolder = "transformer")`` under the load lock during + "finalizing", after the previous pipeline was already evicted, where + unload/cancellation cannot preempt the download. Mirrors the dense-path + gates in ``load_pipeline``: quant requested and supported for this device, + and no pre-quantized checkpoint that would shortcut the dense build.""" + mode = normalize_transformer_quant(kwargs.get("transformer_quant")) + if mode is None: + return False + try: + target = self._resolve_device_target(fam) + if not dense_transformer_supported(target): + return False + scheme = select_transformer_quant_scheme(target, mode) + if scheme is None: + return False + source = resolve_prequant_source( + fam, scheme, path_override = kwargs.get("transformer_prequant_path") + ) + return source is None + except Exception: # noqa: BLE001 — widening the prefetch is best-effort only + return False + def _prefetch_files( self, repo_id: str, @@ -364,7 +392,16 @@ class DiffusionBackend: ) kwargs["base_repo"] = base expected, base_files = self._estimate_download_bytes( - kwargs["repo_id"], kwargs.get("gguf_filename"), base, kwargs.get("hf_token") + kwargs["repo_id"], + kwargs.get("gguf_filename"), + base, + kwargs.get("hf_token"), + # The dense transformer-quant path downloads the base repo's + # transformer/ shards via from_pretrained(subfolder="transformer") + # INSIDE the locked finalize phase, where unload/cancellation cannot + # preempt the multi-GB pull. When that path can actually run, pull the + # shards here in the preemptible prefetch instead. + include_transformer = self._dense_quant_prefetch_needed(fam, kwargs), ) with self._lock: # Stamp progress only if this load is still current; a superseding @@ -433,7 +470,12 @@ class DiffusionBackend: @staticmethod def _estimate_download_bytes( - repo_id: str, gguf_filename: Optional[str], base_repo: str, hf_token: Optional[str] + repo_id: str, + gguf_filename: Optional[str], + base_repo: str, + hf_token: Optional[str], + *, + include_transformer: bool = False, ) -> tuple[int, list[str]]: """Total download size for the progress bar, plus the base-repo files to fetch (the prefetch reuses this list, so the base is listed only once).""" @@ -452,7 +494,7 @@ class DiffusionBackend: total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename) base_info = api.model_info(base_repo, files_metadata = True, token = hf_token) for s in base_info.siblings: - if _base_file_downloaded(s.rfilename): + if _base_file_downloaded(s.rfilename, include_transformer = include_transformer): base_files.append(s.rfilename) total += s.size or 0 except Exception as exc: # noqa: BLE001 — estimate is best-effort @@ -1156,16 +1198,18 @@ def _hf_base_model(repo_id: str, hf_token: Optional[str]) -> Optional[str]: return base if isinstance(base, str) and base.strip() else None -def _base_file_downloaded(rfilename: str) -> bool: +def _base_file_downloaded(rfilename: str, *, include_transformer: bool = False) -> bool: """True for base-repo files ``from_pretrained`` actually fetches. The transformer is supplied by the GGUF, and repo docs (``assets/``, the top-level README/PDF/images) are never downloaded — counting them would peg the progress estimate above what lands on disk, so the bar would sit short of 100% for the whole pipeline-load phase instead of advancing to "finalizing". - """ + ``include_transformer`` admits the ``transformer/`` shards for loads where the + dense transformer-quant path will fetch them anyway (see + ``_dense_quant_prefetch_needed``).""" if rfilename.startswith("transformer/"): - return False + return include_transformer if "/" not in rfilename: # top-level: only the pipeline manifest is fetched return rfilename == "model_index.json" return not rfilename.startswith("assets/") diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index a39565f1fc..656e6abfcf 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -367,7 +367,17 @@ def plan_diffusion_memory( policy = OFFLOAD_MODEL reasons.append("companions exceed budget; whole-module offload of every component") - if explicit_offload and policy == OFFLOAD_NONE and can_offload and not device_memory.is_unified: + # The legacy cpu_offload flag only applies when NO memory_mode was supplied: + # the API documents memory_mode as overriding cpu_offload when set, so an + # explicit `fast` request must stay resident even if the caller also left the + # old flag enabled. + if ( + explicit_offload + and normalize_memory_mode(requested_mode) is None + and policy == OFFLOAD_NONE + and can_offload + and not device_memory.is_unified + ): policy = OFFLOAD_MODEL reasons.append("explicit cpu_offload overrides resident placement") diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 7ea7bb1cc5..1ea3d19e98 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1193,6 +1193,52 @@ def test_transformer_quant_unsupported_scheme_skips_dense_download( assert _FakeTransformer.last["path"] # GGUF from_single_file used +def test_base_file_downloaded_include_transformer_flag(): + # Default: transformer/ shards are the GGUF's job, so they are excluded from + # the prefetch list; the dense transformer-quant path opts them back in. + from core.inference.diffusion import _base_file_downloaded + + assert _base_file_downloaded("transformer/diffusion_pytorch_model-00001.safetensors") is False + assert ( + _base_file_downloaded( + "transformer/diffusion_pytorch_model-00001.safetensors", include_transformer = True + ) + is True + ) + # The flag must not admit anything else that is normally excluded. + assert _base_file_downloaded("assets/teaser.png", include_transformer = True) is False + assert _base_file_downloaded("README.md", include_transformer = True) is False + + +def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): + # The transformer/ prefetch only widens when the dense quant path can really + # run: quant requested + device supported + scheme resolvable + no prequant + # checkpoint shortcutting the dense build. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + fam = detect_family("unsloth/Z-Image-Turbo-GGUF") + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8") + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) + + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is True + # No quant requested -> never widen. + assert backend._dense_quant_prefetch_needed(fam, {}) is False + # A resolvable pre-quantized checkpoint shortcuts the dense download. + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object()) + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False + # Unsupported scheme bails before the dense path (and so must the prefetch). + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) + monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: None) + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False + # Device without dense support (e.g. non-CUDA) never widens. + monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8") + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: False) + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False + + def test_companion_cache_bytes_local_dir_excludes_transformer(tmp_path): # A LOCAL diffusers base: sum the on-disk VAE / text-encoder weights so auto memory # planning sees the resident companions, but exclude transformer/ (the GGUF supplies diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index bc9ba592ce..e6cc054648 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -270,6 +270,21 @@ def test_explicit_cpu_offload_overrides_resident_auto_choice(): assert any("explicit cpu_offload" in r for r in plan.reasons) +def test_explicit_memory_mode_wins_over_legacy_cpu_offload(): + # The API documents memory_mode as overriding cpu_offload when set: fast + + # the legacy flag must stay resident, not silently downgrade to offload. + plan = plan_diffusion_memory( + target = _target(), + device_memory = _discrete(80000), + model_dense_mib = 4000, + runtime_headroom_mib = 2000, + requested_mode = MEMORY_MODE_FAST, + explicit_offload = True, + ) + assert plan.offload_policy == OFFLOAD_NONE + assert not any("explicit cpu_offload" in r for r in plan.reasons) + + def test_explicit_cpu_offload_ignored_on_cpu_target(): plan = plan_diffusion_memory( target = _target(device = "cpu", backend = "cpu", supports_offload = False),