From c9bb65902e4e575abde837fd7941bc9415ebb32b Mon Sep 17 00:00:00 2001 From: hsliuustc0106 Date: Wed, 26 Aug 2026 13:43:00 +0800 Subject: [PATCH] [Model] Learn the retention budget directly from observed peaks Drop the staged-bytes seed per review direction: the budget is now min(observed_peak * 1.25, 25% of device capacity), nothing else. A cold start pays one release per observation step while the peak climbs (two on the validated MiniMax workload), then boundaries retain steadily; the 5% free floor and capacity cap are unchanged. The staged-byte ledger remains only to detect full detachment, which now resets the observed peak so one workload cannot permanently influence another. Trade-off accepted per review: lower reserved memory and a simpler policy over a flush-free first request. Signed-off-by: hsliuustc0106 --- .../distributed_layerwise_offload.md | 25 ++--- .../offloader/test_module_residency.py | 93 ++++++++++--------- .../diffusion/offloader/module_residency.py | 57 ++++++------ 3 files changed, 89 insertions(+), 86 deletions(-) diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 19429231a..34abbfec8 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -177,18 +177,19 @@ component can be reused by the next. The DiT prefetch path and its two shared device buffers are unchanged. After a component is offloaded, cached-but-unallocated memory is retained only -while it stays within a budget derived from the workload itself and at least -5% of device capacity remains physically free. The budget starts from the -model topology — the staged component byte total reported by the attached -`PinnedModuleStager` instances times a seed multiplier — and then tracks the -largest cached value observed at any component boundary: if a workload's -transient footprint outgrows the seed, the first over-budget boundary -releases once, the observation raises the budget to the observed peak times -a margin, and later boundaries retain. The budget is capped at a quarter of -device capacity so neither a large staged total nor a spurious observation -can monopolize a small device, and a detached stager returns its bytes to the -seed. No user-chosen fraction and no warmup flag are needed. Each decision -logs the staged byte total, observed peak, effective budget and cap, +while it stays within a budget learned from the running workload and at least +5% of device capacity remains physically free. The budget is the largest +cached value observed at any component boundary times a margin, so it tracks +the weight-reuse payload plus the transient footprint whatever the +resolution. It starts at zero: a cold start pays one release per observation +step while the peak climbs, after which boundaries retain steadily, and a +workload change whose cached footprint grows by more than the margin pays one +further release while the new peak lands. The budget is capped at a quarter +of device capacity so a spurious observation cannot monopolize a small +device, and detaching the last attached stager resets the learned peak so one +workload never permanently influences another. No user-chosen fraction, no +multiplier constant, and no warmup flag are needed. Each decision logs the +staged byte total, observed peak, effective budget and cap, cached-but-unallocated bytes, and free memory. Missing allocator telemetry also releases it conservatively. Component or staging failure forces a release, and an out-of-memory allocation gets one diff --git a/tests/diffusion/offloader/test_module_residency.py b/tests/diffusion/offloader/test_module_residency.py index f1cd59c96..16b9202c6 100644 --- a/tests/diffusion/offloader/test_module_residency.py +++ b/tests/diffusion/offloader/test_module_residency.py @@ -202,15 +202,23 @@ def test_allocator_cache_policy_enforces_cache_and_free_memory_bounds( expected_release, ): empty_cache, _ = patched_runtime - monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: reserved) - monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: allocated) + telemetry = {"reserved": 20, "allocated": 10, "free": 80} + monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: telemetry["reserved"]) + monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: telemetry["allocated"]) monkeypatch.setattr( residency_module.current_omni_platform, "get_device_memory", - lambda _device: (free, 100), + lambda _device: (telemetry["free"], 400), ) retention = BoundedAllocatorCache(torch.device("cpu")) - retention.grow_retention_budget(10) + + # Prime the learned budget: the cold-start check releases (budget 0) and + # records cached = 10, so later checks retain at budget 12.5. + assert retention.release_if_needed() + assert retention.observed_cached_peak == 10 + empty_cache.reset_mock() + + telemetry.update(reserved=reserved, allocated=allocated, free=free) released = retention.release_if_needed() @@ -218,7 +226,7 @@ def test_allocator_cache_policy_enforces_cache_and_free_memory_bounds( assert empty_cache.call_count == int(expected_release) -def test_allocator_cache_budget_is_derived_from_attached_stager_bytes(patched_runtime): +def test_attached_stagers_form_a_detach_ledger_that_resets_the_observed_peak(patched_runtime): left, right = _aliased_modules() retention = BoundedAllocatorCache(torch.device("cpu")) assert retention.budget_bytes == 0 @@ -231,57 +239,57 @@ def test_allocator_cache_budget_is_derived_from_attached_stager_bytes(patched_ru ) expected_bytes = 12 * 4 + 8 * 8 # float32 parameter + int64 buffer storages assert stager.staged_bytes == expected_bytes - assert retention.budget_bytes == 2 * expected_bytes + assert retention._staged_bytes == expected_bytes # Re-attaching the same cache does not double count. stager.set_cache_retention(retention) - assert retention.budget_bytes == 2 * expected_bytes + assert retention._staged_bytes == 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. + # A second stager adds its own staged bytes to the ledger. 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) + assert retention._staged_bytes == expected_bytes + extra.staged_bytes + + # A learned peak survives a partial detach... + retention._observed_cached_peak = 30 + extra.set_cache_retention(None) + assert retention.observed_cached_peak == 30 + assert retention.budget_bytes == 37 + + # ...but the last detach resets it so one workload cannot influence the next. + stager.set_cache_retention(None) + assert retention.observed_cached_peak == 0 + assert retention.budget_bytes == 0 -@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, -): +def test_allocator_cache_budget_is_capped_by_device_capacity(monkeypatch, patched_runtime): empty_cache, _ = patched_runtime - monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: reserved) - monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: allocated) + telemetry = {"reserved": 40, "allocated": 10, "free": 80} + monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: telemetry["reserved"]) + monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: telemetry["allocated"]) monkeypatch.setattr( residency_module.current_omni_platform, + # Capacity 100 keeps the 25% cap at 25, below the learned budget of + # 30 * 1.25 = 37.5, so the cap must govern. "get_device_memory", - lambda _device: (free, 100), + lambda _device: (telemetry["free"], 100), ) retention = BoundedAllocatorCache(torch.device("cpu")) - retention.grow_retention_budget(1000) - released = retention.release_if_needed() + # Cold start: budget 0 releases and records the peak of 30. + assert retention.release_if_needed() + assert retention.observed_cached_peak == 30 + empty_cache.reset_mock() - assert released is expected_release - assert empty_cache.call_count == int(expected_release) + # The learned budget 37.5 is capped to 25, so cached 30 keeps releasing. + assert retention.release_if_needed() + assert retention.release_if_needed() + empty_cache.reset_mock() + + # Within the cap (cached 20 <= 25) the policy retains. + telemetry.update(reserved=30, allocated=10) + assert not retention.release_if_needed() + assert empty_cache.call_count == 0 def test_allocator_cache_releases_when_telemetry_is_unavailable(monkeypatch, patched_runtime): @@ -310,10 +318,9 @@ def test_allocator_cache_budget_adapts_to_observed_peak(monkeypatch, patched_run lambda _device: (telemetry["free"], 400), ) retention = BoundedAllocatorCache(torch.device("cpu")) - retention.grow_retention_budget(10) - # First boundary: cached 30 exceeds the seed budget 20, so the policy - # releases once and records the observed peak. + # Cold-start boundary: budget 0, so the policy releases and records the + # observed peak. assert retention.release_if_needed() assert retention.observed_cached_peak == 30 diff --git a/vllm_omni/diffusion/offloader/module_residency.py b/vllm_omni/diffusion/offloader/module_residency.py index cc523707b..d97a9e45c 100644 --- a/vllm_omni/diffusion/offloader/module_residency.py +++ b/vllm_omni/diffusion/offloader/module_residency.py @@ -20,19 +20,14 @@ from .tensor_utils import is_dtensor, set_tensor_storage logger = init_logger(__name__) -# Seed headroom over the staged component bytes. Retained blocks include -# allocator fragmentation and transient storages freed alongside the weights -# at component boundaries, so the reusable payload alone under-bounds the -# working set the policy must keep to be effective. Observations take over -# from this seed at the first component boundary (see below). -_RETENTION_HEADROOM_MULTIPLIER = 2.0 # Margin over the largest cached-but-unallocated value observed at a component -# boundary. The observed peak already includes the transient footprint of the -# running workload, so this only covers check-to-check variation. +# boundary. The observed peak already includes the weight-reuse payload and +# the transient footprint of the running workload, so this only covers +# check-to-check variation. _OBSERVED_PEAK_MARGIN = 1.25 -# Hard ceiling for the derived budget so a large staged total or a spurious -# observation can never monopolize a small device; the #6526 L20X retention -# plateau (~23.8 GB on a ~140 GB card) stayed well below this fraction. +# Hard ceiling for the learned budget so a spurious observation 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 @@ -44,22 +39,22 @@ class BoundedAllocatorCache: cached blocks are immediately reusable. This policy keeps the cache while both of these bounds hold: - * cached-but-unallocated memory stays within the derived budget, capped - at a quarter of device capacity; and + * cached-but-unallocated memory stays within the learned budget — the + largest cached value observed at any component boundary times a margin, + capped at a quarter of device capacity; and * at least 5% of device capacity is physically free. - The budget needs no user-chosen fraction and no warmup flag: it starts - from the staged component byte total reported by attached - ``PinnedModuleStager`` instances (times the seed multiplier) and then - tracks the largest cached value observed at any boundary. If a workload's - transient footprint outgrows the seed, the first over-budget boundary - releases once, the observation raises the budget to the observed peak - times its margin, and later boundaries retain. The budget therefore never - exceeds real observed usage by more than the margin, whatever the - resolution or workload. Detaching a stager returns its bytes to the seed. - 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. + The budget needs no user-chosen fraction, no multiplier constant, and no + warmup flag. It starts at zero, so the first boundaries of a cold start + release once per observation step while the peak climbs; afterwards the + budget tracks the observed peak and boundaries retain steadily. A + workload change whose cached footprint grows by more than the margin pays + one further release while the new peak lands. Detaching the last attached + ``PinnedModuleStager`` resets the observed peak, so one workload never + permanently influences another. 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__( @@ -76,16 +71,18 @@ class BoundedAllocatorCache: self._observed_cached_peak = 0 def grow_retention_budget(self, staged_bytes: int) -> None: - """Grow the seed cache budget by a component's staged byte total.""" + """Record an attached stager's staged byte total.""" if staged_bytes < 0: 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.""" + """Record a stager detach; the last detach resets the observed peak.""" 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)) + if self._staged_bytes == 0: + self._observed_cached_peak = 0 @property def observed_cached_peak(self) -> int: @@ -94,10 +91,8 @@ class BoundedAllocatorCache: @property def budget_bytes(self) -> int: - """Reusable-cache ceiling: seed or observed peak, whichever is larger.""" - seed = int(self._staged_bytes * _RETENTION_HEADROOM_MULTIPLIER) - learned = int(self._observed_cached_peak * _OBSERVED_PEAK_MARGIN) - return max(seed, learned) + """Reusable-cache ceiling: the observed peak times its margin.""" + return int(self._observed_cached_peak * _OBSERVED_PEAK_MARGIN) def _should_release(self) -> bool: reserved = int(torch.accelerator.memory_reserved(self.device))