From c5295f3ffe94a5ec150bda6271a9d3b827d9b771 Mon Sep 17 00:00:00 2001 From: hsliuustc0106 Date: Wed, 26 Aug 2026 11:05:17 +0800 Subject: [PATCH] [Model] Cap derived cache budget and account for stager detach Address review on #6629: - Cap the derived retention budget at a quarter of device capacity so a large staged total cannot monopolize a small device; the multiplier itself stays pending an instrumented target run. - Log every release decision (reason, cached bytes, effective budget and cap, staged total, free/total memory) and every staging OOM retry with its staged byte total, so a single target run validates the headroom multiplier directly. - Make stager detach symmetric: set_cache_retention now shrinks the previous cache's budget before growing the new one, so cache -> None -> same cache counts the staged bytes exactly once (regression test added). Signed-off-by: hsliuustc0106 --- .../distributed_layerwise_offload.md | 11 ++-- .../offloader/test_module_residency.py | 41 ++++++++++++++ .../diffusion/offloader/module_residency.py | 54 ++++++++++++++++--- 3 files changed, 97 insertions(+), 9 deletions(-) diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index f0ffd946a..a2c215ffd 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -180,9 +180,14 @@ After a component is offloaded, cached-but-unallocated memory is retained only while it stays within a budget derived from the model topology — the staged component byte total reported by the attached `PinnedModuleStager` instances times a fixed headroom multiplier — and at least 5% of device capacity -remains physically free. Crossing either bound releases the allocator cache, -so the budget scales with the staged components rather than with device -capacity and needs no user-chosen fraction. +remains physically free. The derived budget is capped at a quarter of device +capacity so a large staged total cannot monopolize a small device, and a +detached stager returns its bytes to the budget. Crossing either bound +releases the allocator cache, so the budget scales with the staged components +rather than with device capacity and needs no user-chosen fraction. Each +release decision logs the staged byte total, the effective budget and cap, +cached-but-unallocated bytes, and free memory, so a target run can validate +the headroom multiplier directly. Missing allocator telemetry also releases it conservatively. Component or staging failure forces a release, and an out-of-memory allocation gets one retry after release. This is a component-local policy rather than a global diff --git a/tests/diffusion/offloader/test_module_residency.py b/tests/diffusion/offloader/test_module_residency.py index 897cc17bc..7775edb66 100644 --- a/tests/diffusion/offloader/test_module_residency.py +++ b/tests/diffusion/offloader/test_module_residency.py @@ -237,12 +237,53 @@ def test_allocator_cache_budget_is_derived_from_attached_stager_bytes(patched_ru stager.set_cache_retention(retention) assert retention.budget_bytes == 2 * expected_bytes + # Detaching and re-attaching the same cache also counts the bytes once. + stager.set_cache_retention(None) + assert retention.budget_bytes == 0 + stager.set_cache_retention(retention) + assert retention.budget_bytes == 2 * expected_bytes + # A second stager grows the budget by its own staged bytes. extra = PinnedModuleStager(torch.nn.Linear(2, 2), torch.device("cpu"), pin_memory=False) extra.set_cache_retention(retention) assert retention.budget_bytes == 2 * (expected_bytes + extra.staged_bytes) +@pytest.mark.parametrize( + ("reserved", "allocated", "free", "expected_release"), + [ + # Derived budget 2 * 1000 = 2000 uncapped, but capacity cap = 25 and + # cached = 30 exceeds it, so the cap must trigger the release. + (40, 10, 80, True), + # Within the cap (cached = 20 <= 25) and the free floor holds. + (30, 10, 80, False), + ], +) +def test_allocator_cache_budget_is_capped_by_device_capacity( + monkeypatch, + patched_runtime, + reserved, + allocated, + free, + expected_release, +): + empty_cache, _ = patched_runtime + monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: reserved) + monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: allocated) + monkeypatch.setattr( + residency_module.current_omni_platform, + "get_device_memory", + lambda _device: (free, 100), + ) + retention = BoundedAllocatorCache(torch.device("cpu")) + retention.grow_retention_budget(1000) + + released = retention.release_if_needed() + + assert released is expected_release + assert empty_cache.call_count == int(expected_release) + + def test_allocator_cache_releases_when_telemetry_is_unavailable(monkeypatch, patched_runtime): empty_cache, _ = patched_runtime diff --git a/vllm_omni/diffusion/offloader/module_residency.py b/vllm_omni/diffusion/offloader/module_residency.py index 0f07ac2e7..708181769 100644 --- a/vllm_omni/diffusion/offloader/module_residency.py +++ b/vllm_omni/diffusion/offloader/module_residency.py @@ -25,6 +25,10 @@ logger = init_logger(__name__) # weights at component boundaries, so the reusable payload alone under-bounds # the working set the policy must keep to be effective. _RETENTION_HEADROOM_MULTIPLIER = 2.0 +# Hard ceiling for the derived budget so a large staged total can never +# monopolize a small device; the #6526 L20X retention plateau (~23.8 GB on a +# ~140 GB card) stayed well below this fraction. +_MAX_BUDGET_CAPACITY_FRACTION = 0.25 class BoundedAllocatorCache: @@ -37,15 +41,16 @@ class BoundedAllocatorCache: * cached-but-unallocated memory stays within the derived budget: the staged component byte total reported by attached ``PinnedModuleStager`` - instances times a fixed headroom multiplier; and + instances times a fixed headroom multiplier, capped at a quarter of + device capacity; and * at least 5% of device capacity is physically free. The budget derives from the model topology instead of a user-chosen fraction, so it scales with the staged components rather than with device - capacity. Missing memory telemetry is handled conservatively by releasing - the cache. Failure paths can force release; normal executor shutdown keeps - its own unconditional device-cache cleanup because this policy is not - global. + capacity. Detaching a stager returns its bytes to the budget. Missing + memory telemetry is handled conservatively by releasing the cache. + Failure paths can force release; normal executor shutdown keeps its own + unconditional device-cache cleanup because this policy is not global. """ def __init__( @@ -66,6 +71,12 @@ class BoundedAllocatorCache: raise ValueError(f"staged_bytes must be >= 0, got {staged_bytes}") self._staged_bytes += int(staged_bytes) + def shrink_retention_budget(self, staged_bytes: int) -> None: + """Return bytes previously grown, e.g. when a stager detaches.""" + if staged_bytes < 0: + raise ValueError(f"staged_bytes must be >= 0, got {staged_bytes}") + self._staged_bytes = max(0, self._staged_bytes - int(staged_bytes)) + @property def budget_bytes(self) -> int: """Reusable-cache ceiling: staged bytes times the headroom multiplier.""" @@ -76,7 +87,32 @@ class BoundedAllocatorCache: allocated = int(torch.accelerator.memory_allocated(self.device)) free, total = current_omni_platform.get_device_memory(self.device) cached = max(0, reserved - allocated) - return cached > self.budget_bytes or free < int(total * self.min_free_fraction) + capacity_cap = int(total * _MAX_BUDGET_CAPACITY_FRACTION) + budget = min(self.budget_bytes, capacity_cap) + if cached > budget or free < int(total * self.min_free_fraction): + reason = "cached-over-budget" if cached > budget else "low-free-memory" + logger.debug( + "Releasing retained allocator cache (%s): cached=%d budget=%d staged=%d " + "capacity_cap=%d free=%d total=%d", + reason, + cached, + budget, + self._staged_bytes, + capacity_cap, + free, + total, + ) + return True + logger.debug( + "Retaining allocator cache: cached=%d budget=%d staged=%d capacity_cap=%d free=%d total=%d", + cached, + budget, + self._staged_bytes, + capacity_cap, + free, + total, + ) + return False def release_if_needed(self, *, force: bool = False) -> bool: """Release cached blocks when a bound is crossed or release is forced.""" @@ -227,6 +263,8 @@ class PinnedModuleStager: def set_cache_retention(self, cache_retention: BoundedAllocatorCache | None) -> None: if cache_retention is self.cache_retention: return + if self.cache_retention is not None: + self.cache_retention.shrink_retention_budget(self.staged_bytes) self.cache_retention = cache_retention if cache_retention is not None: cache_retention.grow_retention_budget(self.staged_bytes) @@ -269,6 +307,10 @@ class PinnedModuleStager: except torch.OutOfMemoryError: # A retained cache is normally reusable by this process, but an # explicit flush gives external memory pressure one bounded retry. + logger.info( + "Staging out of memory (%d staged bytes); releasing retained cache and retrying once", + self.staged_bytes, + ) self._cleanup_failed_load() try: self._load_once()