From cf1910c92de9ed480b4e9afc7e6f532430709247 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 17 Jul 2026 08:46:12 +0000 Subject: [PATCH] Harden the diffusion memory plan against transient free-VRAM undercounts A cold FLUX.2-dev int8 load on an idle 183 GB B200 planned offload=model (companions exceed budget) and silently served the GGUF as-is; the identical retry went resident and engaged the hosted prequant. The plan arithmetic was byte-identical across both loads (required 90,228 MiB, resident needs free of about 124 GB); the only divergent input was torch.cuda.mem_get_info, which is device-wide and instantaneous: a transient foreign CUDA context briefly held about 100 GB at the first snapshot, and the planner trusted that single read. Three changes: - settled_snapshot_device_memory: on cuda, synchronize + empty_cache (best-effort) and take the MAX free over up to 3 spaced reads. A transient can only shrink free, so the max rejects transient undercounts while a persistent tenant still caps every read. _plan_memory now uses it. - plan_fits_total_capacity + one replan retry: when the dense/prequant candidate fits TOTAL device capacity under the standard reserve and the 0.85 resident margin, an offload verdict can only stem from the free reading, so the loader re-snapshots and replans once before declining the fast path. Explicit balanced/low_vram modes skip the retry (they offload by mode). - diffusion.transformer_quant_declined log line with required/budget/free and the plan reasons, so the next decline is diagnosable from the server log (previously silent). Verified: cold FLUX.2-dev int8 first load in a fresh server now engages the hosted prequant resident (offload=none). --- studio/backend/core/inference/diffusion.py | 68 +++++++--- .../core/inference/diffusion_memory.py | 75 ++++++++++- .../backend/tests/test_diffusion_backend.py | 121 +++++++++++++++++- studio/backend/tests/test_diffusion_memory.py | 74 +++++++++++ 4 files changed, 312 insertions(+), 26 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 29f3d1fe11..3cd5f1199e 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -59,7 +59,8 @@ from .diffusion_memory import ( file_size_mib, normalize_memory_mode, plan_diffusion_memory, - snapshot_device_memory, + plan_fits_total_capacity, + settled_snapshot_device_memory, ) from .diffusion_speed import ( SPEED_DEFAULT, @@ -1232,22 +1233,50 @@ class DiffusionBackend: logger = logger, ) if candidate is not None: - replanned = self._plan_memory( - target, - single_file_path, - base, - fam, - memory_mode, - cpu_offload, - kind = kind, - repo_id = repo_id, - transformer_resident_override_mib = ( - candidate.transient_transformer_mib - ), - # Pass the auto-policy's companion estimate so the prefetched base - # transformer/ shards in the cache aren't double-counted. - companion_override_mib = candidate.companions_mib, - ) + def _replan_candidate(): + return self._plan_memory( + target, + single_file_path, + base, + fam, + memory_mode, + cpu_offload, + kind = kind, + repo_id = repo_id, + transformer_resident_override_mib = ( + candidate.transient_transformer_mib + ), + # Pass the auto-policy's companion estimate so the prefetched base + # transformer/ shards in the cache aren't double-counted. + companion_override_mib = candidate.companions_mib, + ) + + replanned = _replan_candidate() + if ( + replanned.offload_policy != OFFLOAD_NONE + # Explicit balanced/low_vram picks offload BY MODE; a fresh + # snapshot cannot change that, so don't waste a retry. + and normalize_memory_mode(memory_mode) + not in (MEMORY_MODE_BALANCED, MEMORY_MODE_LOW_VRAM) + and plan_fits_total_capacity(replanned) + ): + # The candidate fits TOTAL device capacity with the standard + # reserve + resident margin, yet the instantaneous free reading + # said no: a transient foreign allocation (measured on B200: + # ~100 GB held for under a minute on an idle card) must not + # force the GGUF fallback. Re-snapshot (settled) and replan + # once before declining. + replanned = _replan_candidate() + if replanned.offload_policy != OFFLOAD_NONE: + logger.info( + "diffusion.transformer_quant_declined: required=%s MiB " + "budget=%s MiB free=%s MiB policy=%s (%s)", + replanned.estimates.get("resident_required_mib"), + replanned.estimates.get("safe_device_budget_mib"), + getattr(replanned.device_memory, "free_mib", None), + replanned.offload_policy, + "; ".join(replanned.reasons), + ) if replanned.offload_policy == OFFLOAD_NONE: quant_plan = replanned # The GGUF plan already declined resident; a prequant-sized @@ -1862,7 +1891,10 @@ class DiffusionBackend: 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) + # Settled (max-over-reads) on cuda: a transient foreign allocation at the wrong instant + # otherwise makes an empty card look full and silently declines the resident/quant fast + # path (see settled_snapshot_device_memory). + device_memory = settled_snapshot_device_memory(target) if kind == "pipeline": # The whole repo is one cached download; cached bytes are the resident estimate # (bnb-4bit/fp8 stay compressed). A LOCAL path isn't cached, so sum its on-disk weights. diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 29b16ac579..db2505555e 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -148,6 +148,46 @@ def snapshot_device_memory(target: Any) -> DeviceMemory: return DeviceMemory(backend, device, "system_memory", free, total) +def settled_snapshot_device_memory( + target: Any, attempts: int = 3, delay_s: float = 1.0 +) -> DeviceMemory: + """``snapshot_device_memory`` hardened against TRANSIENT free-VRAM undercounts on cuda. + + ``torch.cuda.mem_get_info`` is device-wide and instantaneous: a neighbouring process (or a + just-spawned subprocess context) briefly holding tens of GB at the wrong moment makes an + empty card look full, and the planner then silently declines the resident/quant fast path + (measured on B200: a cold FLUX.2-dev int8 load saw free < 74 GB on an idle 183 GB card and + fell back to offloaded GGUF; the identical retry saw >= 124 GB and went resident). Settle + the allocator (synchronize + empty_cache, best-effort) and take the MAX free over a few + spaced reads: a transient can only SHRINK free, so the max rejects transient undercounts + while a persistent tenant still caps every read. Non-cuda targets keep the single read.""" + if getattr(target, "device", "cpu") != "cuda": + return snapshot_device_memory(target) + try: + import torch + + torch.cuda.synchronize() + torch.cuda.empty_cache() + except Exception: # noqa: BLE001 — settle is best-effort; the snapshot below still runs + pass + best = snapshot_device_memory(target) + for _ in range(max(0, attempts - 1)): + if best.free_mib is not None and best.total_mib is not None: + # Free already within the reserve of total: nothing transient to wait out. + if best.free_mib >= best.total_mib - max(2048, int(best.total_mib * 0.10)): + break + try: + import time + + time.sleep(delay_s) + except Exception: # noqa: BLE001 + break + nxt = snapshot_device_memory(target) + if nxt.free_mib is not None and (best.free_mib is None or nxt.free_mib > best.free_mib): + best = nxt + return best + + def _cuda_memory(backend: str) -> tuple[Optional[int], Optional[int], str]: try: import torch @@ -274,19 +314,40 @@ def estimate_video_runtime_mib( return max(3072, int(4096 + 3.0 * decoded_mib)) +def _reserve_mib(memory_kind: str, base: int) -> int: + if memory_kind == "unified_memory": + return max(2048, int(base * 0.20)) # OS + CPU share this pool + if memory_kind == "system_memory": + return max(1024, int(base * 0.10)) + return max(2048, int(base * 0.10)) + + def _safe_device_budget_mib(memory: DeviceMemory) -> Optional[int]: """Free memory minus a headroom reserve (room for fragmentation + other tenants). None when free memory is unknown.""" if memory.free_mib is None: return None base = memory.total_mib or memory.free_mib - if memory.memory_kind == "unified_memory": - reserve = max(2048, int(base * 0.20)) # OS + CPU share this pool - elif memory.memory_kind == "system_memory": - reserve = max(1024, int(base * 0.10)) - else: - reserve = max(2048, int(base * 0.10)) - return max(0, int(memory.free_mib) - reserve) + return max(0, int(memory.free_mib) - _reserve_mib(memory.memory_kind, base)) + + +def plan_fits_total_capacity(plan: Any) -> bool: + """Whether ``plan``'s resident requirement fits TOTAL device capacity under the standard + reserve + the 0.85 resident margin -- i.e. an offload decision can only stem from the + instantaneous FREE reading (something else held VRAM at snapshot time), never from the + device being too small. Used to retry a declined resident/quant plan once with a fresh + settled snapshot instead of trusting a single transient undercount. False on any missing + input (unknown sizes keep today's behaviour).""" + try: + required = plan.estimates.get("resident_required_mib") + memory = plan.device_memory + total = memory.total_mib + kind = memory.memory_kind + except Exception: # noqa: BLE001 — malformed plan: no retry + return False + if required is None or total is None: + return False + return int(required) <= int((int(total) - _reserve_mib(kind, int(total))) * 0.85) def _sum_required(*values: Optional[int]) -> Optional[int]: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index b6e2af4473..5c6c5e3ca5 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2712,6 +2712,125 @@ def test_dense_quant_prequant_proceeds_but_forbids_dense_fallback(fake_runtime, assert attempted == [False] # ...fast path still attempted, dense fallback forbidden +def test_dense_quant_replan_retries_once_on_transient_free_undercount( + fake_runtime, tmp_path, monkeypatch +): + # A transient foreign allocation at snapshot time makes an empty card look full and the + # candidate replan declines resident -- but the candidate FITS total capacity, so the + # loader must retry the replan once with a fresh settled snapshot instead of silently + # falling back to GGUF-as-is (measured: FLUX.2-dev int8 cold load on an idle B200). + import dataclasses + + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "int8" + ) + monkeypatch.setattr( + dmod, + "resolve_dense_quant_candidate", + lambda **kw: types.SimpleNamespace( + transient_transformer_mib = 33_831, companions_mib = 46_157, prequant = True + ), + ) + replan_calls = [] + orig_plan = DiffusionBackend._plan_memory + + def spy_plan(self, *a, transformer_resident_override_mib = None, **k): + real = orig_plan( + self, *a, transformer_resident_override_mib = transformer_resident_override_mib, **k + ) + if transformer_resident_override_mib is None: + # Initial GGUF plan: force offload so the candidate replan branch is entered. + return dataclasses.replace(real, offload_policy = "model") + replan_calls.append(True) + if len(replan_calls) == 1: + # First replan: the transient undercount. Required fits total capacity + # (90,228 <= 0.85 * (183,359 - 18,335)), so a retry must follow. + return types.SimpleNamespace( + offload_policy = "model", + estimates = {"resident_required_mib": 90_228, "safe_device_budget_mib": 40_000}, + device_memory = types.SimpleNamespace( + total_mib = 183_359, memory_kind = "discrete_vram", free_mib = 60_000 + ), + reasons = ("companions exceed budget",), + ) + # Retry: the transient cleared; resident. + return dataclasses.replace(real, offload_policy = "none") + + monkeypatch.setattr(DiffusionBackend, "_plan_memory", spy_plan) + attempted = [] + + def fake_dense_load(self, *a, **k): + attempted.append(k.get("allow_dense_fallback")) + raise RuntimeError("test: stop after reaching the fast path") + + monkeypatch.setattr(DiffusionBackend, "_load_dense_quant_pipeline", fake_dense_load) + (tmp_path / "m.gguf").write_bytes(b"x") + backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "int8", + ) + assert replan_calls == [True, True] # declined once, retried once + assert attempted == [False] # fast path attempted; prequant-sized plan forbids dense fallback + + +def test_dense_quant_replan_no_retry_when_capacity_truly_short( + fake_runtime, tmp_path, monkeypatch +): + # When the candidate does NOT fit total capacity, the decline is real: no retry. + import dataclasses + + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "int8" + ) + monkeypatch.setattr( + dmod, + "resolve_dense_quant_candidate", + lambda **kw: types.SimpleNamespace( + transient_transformer_mib = 33_831, companions_mib = 46_157, prequant = True + ), + ) + replan_calls = [] + orig_plan = DiffusionBackend._plan_memory + + def spy_plan(self, *a, transformer_resident_override_mib = None, **k): + real = orig_plan( + self, *a, transformer_resident_override_mib = transformer_resident_override_mib, **k + ) + if transformer_resident_override_mib is None: + return dataclasses.replace(real, offload_policy = "model") + replan_calls.append(True) + return types.SimpleNamespace( + offload_policy = "model", + estimates = {"resident_required_mib": 150_000, "safe_device_budget_mib": 40_000}, + device_memory = types.SimpleNamespace( + total_mib = 183_359, memory_kind = "discrete_vram", free_mib = 60_000 + ), + reasons = ("companions exceed budget",), + ) + + monkeypatch.setattr(DiffusionBackend, "_plan_memory", spy_plan) + (tmp_path / "m.gguf").write_bytes(b"x") + backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "int8", + ) + assert replan_calls == [True] # genuine capacity shortfall: declined without a retry + + def test_assemble_pipe_routes_krea2_per_component(monkeypatch): # krea's repo ships transformers-5.x configs and no top-level tokenizer files, so # Pipeline.from_pretrained dies in the tokenizer (vocab_file = None). The quant fast @@ -2994,7 +3113,7 @@ def test_plan_memory_dense_replan_does_not_double_count_prefetched_transformer(m # companions + headroom, but NOT a second copy of the bf16 transformer. monkeypatch.setattr( dmod, - "snapshot_device_memory", + "settled_snapshot_device_memory", lambda t: DeviceMemory("cuda", "cuda", "discrete_vram", 40000, 40960), ) monkeypatch.setattr(dmod, "estimate_image_runtime_mib", lambda **kw: 4000) diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index 41475584ef..e92feaf2d1 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -559,3 +559,77 @@ def test_apply_tolerates_pipe_without_vae_savers(): bare = _Bare() _, tiled = apply_memory_plan(bare, _plan(OFFLOAD_NONE, tiling = False), device = "cpu") assert bare.moved == "cpu" and tiled is False + + +# ── settled snapshot + capacity-fit retry helpers ──────────────────────────── + + +def test_settled_snapshot_takes_max_free_over_reads(monkeypatch): + # A transient foreign allocation can only SHRINK free, so the settled snapshot must + # reject a transient undercount (60 GB free on an idle 183 GB card) by keeping the max + # free across the retry reads. Measured incident: FLUX.2-dev int8 cold load. + from core.inference import diffusion_memory as dm + + reads = [ + DeviceMemory("cuda", "cuda", "discrete_vram", free_mib = 60_000, total_mib = 183_359), + DeviceMemory("cuda", "cuda", "discrete_vram", free_mib = 170_000, total_mib = 183_359), + DeviceMemory("cuda", "cuda", "discrete_vram", free_mib = 170_000, total_mib = 183_359), + ] + monkeypatch.setattr(dm, "snapshot_device_memory", lambda target: reads.pop(0)) + snap = dm.settled_snapshot_device_memory(_target(device = "cuda"), attempts = 3, delay_s = 0) + assert snap.free_mib == 170_000 + + +def test_settled_snapshot_stops_early_when_device_already_idle(monkeypatch): + # First read already within the reserve of total: no transient to wait out, one read only. + from core.inference import diffusion_memory as dm + + calls = [] + + def fake_snapshot(target): + calls.append(1) + return DeviceMemory( + "cuda", "cuda", "discrete_vram", free_mib = 170_000, total_mib = 183_359 + ) + + monkeypatch.setattr(dm, "snapshot_device_memory", fake_snapshot) + snap = dm.settled_snapshot_device_memory(_target(device = "cuda"), attempts = 3, delay_s = 0) + assert snap.free_mib == 170_000 + assert calls == [1] + + +def test_settled_snapshot_passthrough_off_cuda(monkeypatch): + # Non-cuda targets keep the single-read behaviour (no settle loop). + from core.inference import diffusion_memory as dm + + calls = [] + + def fake_snapshot(target): + calls.append(1) + return DeviceMemory("mps", "mps", "unified_memory", free_mib = 8_000, total_mib = 16_000) + + monkeypatch.setattr(dm, "snapshot_device_memory", fake_snapshot) + snap = dm.settled_snapshot_device_memory(_target(device = "mps"), attempts = 3, delay_s = 0) + assert snap.memory_kind == "unified_memory" + assert calls == [1] + + +def test_plan_fits_total_capacity(): + # True exactly when required fits (total - reserve) * 0.85: the decline can then only + # stem from the instantaneous free reading, so a settled retry is worthwhile. + from core.inference.diffusion_memory import plan_fits_total_capacity + + def plan(required, total, kind = "discrete_vram"): + return types.SimpleNamespace( + estimates = {"resident_required_mib": required}, + device_memory = DeviceMemory("cuda", "cuda", kind, free_mib = 1, total_mib = total), + ) + + # FLUX.2-dev int8 incident numbers: 90,228 required on a 183,359 MiB card -> fits. + assert plan_fits_total_capacity(plan(90_228, 183_359)) is True + # Larger than the capacity margin (0.85 * (183,359 - 18,335) = 140,270) -> no retry. + assert plan_fits_total_capacity(plan(150_000, 183_359)) is False + # Unknown sizes keep today's behaviour (no retry). + assert plan_fits_total_capacity(plan(None, 183_359)) is False + assert plan_fits_total_capacity(plan(90_228, None)) is False + assert plan_fits_total_capacity(types.SimpleNamespace()) is False