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).
This commit is contained in:
Daniel Han
parent
6d33582b88
commit
cf1910c92d
4 changed files with 312 additions and 26 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue