From 5a8b81ebebd28b5bc23a8a0721d39998496e7015 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Tue, 25 Aug 2026 23:13:41 +0800 Subject: [PATCH] [Model] Optimize MiniMax-H3 DLO component lifecycle (#6526) Signed-off-by: hsliuustc0106 --- .../distributed_layerwise_offload.md | 27 ++ docs/design/module/diffusion/offloader.md | 7 +- .../test_minimax_h3_dlo_lifecycle.py | 101 +++++++ .../offloader/test_module_residency.py | 230 ++++++++++++++ .../diffusion/models/minimax_h3/encoder.py | 56 +++- .../models/minimax_h3/pipeline_minimax_h3.py | 44 ++- vllm_omni/diffusion/models/minimax_h3/vae.py | 28 +- vllm_omni/diffusion/offloader/__init__.py | 3 +- .../diffusion/offloader/module_residency.py | 280 +++++++++++++++--- 9 files changed, 701 insertions(+), 75 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py create mode 100644 tests/diffusion/offloader/test_module_residency.py diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 1184a70ea..bf282d8c5 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -169,6 +169,31 @@ This mode means: copy per rank. - The scheduler does not require a synchronized DP request wave for DLO. +### Component allocator-cache retention + +MiniMax-H3 DLO keeps one bounded PyTorch allocator cache across its staged +encoder, DiT, and VAE component boundaries so allocations released by one +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 is at most 25% of device capacity and at least 5% of device capacity +remains physically free. Crossing either bound releases the allocator cache. +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 +`empty_cache` override, so the executor's unconditional shutdown cleanup is +preserved. + +Retaining allocator blocks trades higher reserved and peak device memory for +lower lifecycle latency. The benefit is device- and workload-dependent, and a +smaller or externally loaded device may cross a bound and return to the +ordinary per-component cache-release behavior. Current performance validation +covers MiniMax-H3 no-AllGather DLO on CUDA. Other MiniMax-H3 DLO transfer +topologies and platforms receive the same bounded policy but are not claimed +performance targets; other DLO models remain unchanged unless they explicitly +adopt it. + ### Final-layout Host Weight Runtime consumer > **Status:** PR2 implementation. The representation-independent identity, @@ -390,6 +415,8 @@ Current source-level validation includes: transforms without parameter-side flags; - resident-layer requests requiring no-AllGather; - DP request-wave validation for denoising-step compatibility; +- bounded component allocator-cache retention, conservative/forced release, + OOM retry, and immutable encoder non-block staging; - sharding, double-buffer, AllGather-size, and heterogeneous-block regression tests. diff --git a/docs/design/module/diffusion/offloader.md b/docs/design/module/diffusion/offloader.md index 919b1f39a..7ac5d2cc1 100644 --- a/docs/design/module/diffusion/offloader.md +++ b/docs/design/module/diffusion/offloader.md @@ -18,7 +18,7 @@ validation_paths: - tests/diffusion/offloader/** upstream_refs: - torch.nn.Module.to -last_reviewed: 2026-07-16 +last_reviewed: 2026-08-25 --- # Diffusion offloader @@ -48,6 +48,11 @@ and correctness required by the selected execution mode. **Rule:** Retained host and device copies MUST respect configured limits and provide deterministic teardown. +Allocator-cache retention MUST remain local to an explicit component owner, +declare both cache and physical-free-memory bounds, release conservatively +when telemetry is unavailable, and force release on failure or memory +pressure. It MUST NOT replace unconditional executor shutdown cleanup. + ## Safe-change guide Test enabled and disabled paths, repeated execution, asynchronous transfer, diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py new file mode 100644 index 000000000..b0ebf1038 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _OffloadAbort(BaseException): + pass + + +def test_encoder_non_block_children_use_one_shared_snapshot_stager(monkeypatch, mocker): + from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module + + class VisionStack(torch.nn.Module): + def __init__(self): + super().__init__() + self.patch_embed = torch.nn.Linear(2, 2) + self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)]) + + class TextStack(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed_tokens = torch.nn.Embedding(4, 2) + self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)]) + + encoder = object.__new__(encoder_module.MiniMaxH3Qwen3VLEncoder) + torch.nn.Module.__init__(encoder) + encoder.device_target = torch.device("cpu") + encoder.vision = VisionStack() + encoder.text_model = TextStack() + hook = mocker.Mock() + hook.pin_memory = False + encoder._omni_layerwise_hooks = [hook] + encoder._omni_layerwise_enabled = True + cache = mocker.Mock() + encoder.set_omni_component_cache(cache) + stager = mocker.Mock() + stager_cls = mocker.Mock(return_value=stager) + monkeypatch.setattr(encoder_module, "PinnedModuleStager", stager_cls) + + encoder.load_to_device() + encoder.offload_to_cpu() + + stager_cls.assert_called_once_with( + [encoder.vision.patch_embed, encoder.text_model.embed_tokens], + torch.device("cpu"), + pin_memory=False, + cache_retention=cache, + ) + stager.load.assert_called_once_with() + stager.offload.assert_called_once_with() + hook.offload_layer.assert_called_once_with() + + +def test_manual_component_failure_forces_retained_cache_release(mocker): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = mocker.Mock() + pipeline.od_config.enable_layerwise_offload = False + pipeline.od_config.enable_distributed_layerwise_offload = True + pipeline._model_cpu_offload_modules = [] + pipeline._dlo_component_cache = mocker.Mock() + component = mocker.Mock() + component.offload_to_cpu.side_effect = _OffloadAbort("offload failed") + + with pytest.raises(RuntimeError, match="component failed"): + with pipeline._component_on_device(component): + raise RuntimeError("component failed") + + component.load_to_device.assert_called_once_with() + component.offload_to_cpu.assert_called_once_with() + pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True) + + +def test_manual_component_offload_failure_forces_retained_cache_release(mocker): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = mocker.Mock() + pipeline.od_config.enable_layerwise_offload = False + pipeline.od_config.enable_distributed_layerwise_offload = True + pipeline._model_cpu_offload_modules = [] + pipeline._dlo_component_cache = mocker.Mock() + component = mocker.Mock() + component.offload_to_cpu.side_effect = [None, _OffloadAbort("offload failed")] + + with pipeline._component_on_device(component): + pass + with pytest.raises(_OffloadAbort, match="offload failed"): + with pipeline._component_on_device(component): + pass + + assert component.load_to_device.call_count == 2 + assert component.offload_to_cpu.call_count == 2 + pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True) diff --git a/tests/diffusion/offloader/test_module_residency.py b/tests/diffusion/offloader/test_module_residency.py new file mode 100644 index 000000000..adf3c3afc --- /dev/null +++ b/tests/diffusion/offloader/test_module_residency.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from contextlib import contextmanager + +import pytest +import torch +from torch import nn + +import vllm_omni.diffusion.offloader.module_residency as residency_module +from vllm_omni.diffusion.offloader.module_residency import ( + BoundedAllocatorCache, + PinnedModuleStager, +) + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +class _DummyStream: + def wait_stream(self, _stream) -> None: + return None + + def wait_event(self, _event) -> None: + return None + + +class _DummyEvent: + def __init__(self) -> None: + self.record_count = 0 + + def record(self, _stream) -> None: + self.record_count += 1 + + +@contextmanager +def _dummy_stream(_stream): + yield + + +@pytest.fixture +def patched_runtime(monkeypatch, mocker): + empty_cache = mocker.Mock() + synchronize = mocker.Mock() + monkeypatch.setattr(residency_module.current_omni_platform, "Stream", _DummyStream) + monkeypatch.setattr(residency_module.current_omni_platform, "Event", _DummyEvent) + monkeypatch.setattr(residency_module.current_omni_platform, "current_stream", _DummyStream) + monkeypatch.setattr(residency_module.current_omni_platform, "stream", _dummy_stream) + monkeypatch.setattr(residency_module.current_omni_platform, "empty_cache", empty_cache) + monkeypatch.setattr(residency_module.current_omni_platform, "synchronize", synchronize) + return empty_cache, synchronize + + +def _aliased_modules() -> tuple[nn.Module, nn.Module]: + parameter_storage = torch.arange(12, dtype=torch.float32) + buffer_storage = torch.arange(8, dtype=torch.int64) + left = nn.Module() + right = nn.Module() + left.weight = nn.Parameter(parameter_storage[:8].view(2, 4)) + right.weight = nn.Parameter(parameter_storage[4:].view(2, 4)) + left.register_buffer("state", buffer_storage[:6]) + right.register_buffer("state", buffer_storage[2:]) + return left, right + + +def test_module_group_preserves_parameter_and_buffer_storage_aliases(patched_runtime, mocker): + left, right = _aliased_modules() + expected_left_weight = left.weight.detach().clone() + expected_right_weight = right.weight.detach().clone() + expected_left_state = left.state.clone() + expected_right_state = right.state.clone() + retention = mocker.Mock(spec=BoundedAllocatorCache) + stager = PinnedModuleStager( + [left, right], + torch.device("cpu"), + pin_memory=False, + cache_retention=retention, + ) + + master_weight_ptr = left.weight.untyped_storage().data_ptr() + master_state_ptr = left.state.untyped_storage().data_ptr() + assert right.weight.untyped_storage().data_ptr() == master_weight_ptr + assert right.state.untyped_storage().data_ptr() == master_state_ptr + assert (left.weight.storage_offset(), right.weight.storage_offset()) == (0, 4) + assert (left.state.storage_offset(), right.state.storage_offset()) == (0, 2) + + stager.load() + + assert left.weight.untyped_storage().data_ptr() != master_weight_ptr + assert right.weight.untyped_storage().data_ptr() == left.weight.untyped_storage().data_ptr() + assert right.state.untyped_storage().data_ptr() == left.state.untyped_storage().data_ptr() + left.weight.data[1, 0] = -1 + left.state[2] = -1 + assert right.weight[0, 0].item() == -1 + assert right.state[0].item() == -1 + + stager.offload() + + assert left.weight.untyped_storage().data_ptr() == master_weight_ptr + assert right.weight.untyped_storage().data_ptr() == master_weight_ptr + assert left.state.untyped_storage().data_ptr() == master_state_ptr + assert right.state.untyped_storage().data_ptr() == master_state_ptr + assert torch.equal(left.weight, expected_left_weight) + assert torch.equal(right.weight, expected_right_weight) + assert torch.equal(left.state, expected_left_state) + assert torch.equal(right.state, expected_right_state) + retention.release_if_needed.assert_called_once_with(force=False) + + +def test_staging_transitions_are_idempotent_and_reuse_one_event(patched_runtime, mocker): + module = nn.Linear(2, 2) + retention = mocker.Mock(spec=BoundedAllocatorCache) + stager = PinnedModuleStager( + module, + torch.device("cpu"), + pin_memory=False, + cache_retention=retention, + ) + ready_event = stager._ready_event + + stager.load() + first_storage = module.weight.untyped_storage().data_ptr() + stager.load() + assert module.weight.untyped_storage().data_ptr() == first_storage + assert stager._ready_event is ready_event + assert ready_event.record_count == 1 + + stager.offload() + stager.offload() + assert not stager.loaded + retention.release_if_needed.assert_called_once_with(force=False) + + +def test_failed_load_restores_master_and_forces_cache_release(monkeypatch, patched_runtime, mocker): + module = nn.Linear(2, 2) + expected = module.weight.detach().clone() + retention = mocker.Mock(spec=BoundedAllocatorCache) + stager = PinnedModuleStager( + module, + torch.device("cpu"), + pin_memory=False, + cache_retention=retention, + ) + master_ptr = module.weight.untyped_storage().data_ptr() + + def fail_after_rebind() -> None: + device_storages = [torch.empty_like(group.master) for group in stager._groups] + stager._bind(device_storages) + module.weight.data.zero_() + raise RuntimeError("copy failed") + + monkeypatch.setattr(stager, "_load_once", fail_after_rebind) + with pytest.raises(RuntimeError, match="copy failed"): + stager.load() + + assert not stager.loaded + assert module.weight.untyped_storage().data_ptr() == master_ptr + assert torch.equal(module.weight, expected) + retention.release_if_needed.assert_called_once_with(force=True) + + +def test_out_of_memory_releases_cache_and_retries_once(monkeypatch, patched_runtime, mocker): + module = nn.Linear(2, 2) + retention = mocker.Mock(spec=BoundedAllocatorCache) + stager = PinnedModuleStager( + module, + torch.device("cpu"), + pin_memory=False, + cache_retention=retention, + ) + load_once = stager._load_once + attempts = 0 + + def fail_once() -> None: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise torch.OutOfMemoryError("memory pressure") + load_once() + + monkeypatch.setattr(stager, "_load_once", fail_once) + stager.load() + + assert stager.loaded + assert attempts == 2 + retention.release_if_needed.assert_called_once_with(force=True) + + +@pytest.mark.parametrize( + ("reserved", "allocated", "free", "expected_release"), + [ + (20, 10, 80, False), + (40, 10, 80, True), + (20, 10, 4, True), + ], +) +def test_allocator_cache_policy_enforces_cache_and_free_memory_bounds( + 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")) + + 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 + + def unavailable(_device): + raise NotImplementedError + + monkeypatch.setattr(torch.accelerator, "memory_reserved", unavailable) + retention = BoundedAllocatorCache(torch.device("cpu")) + + assert retention.release_if_needed() + empty_cache.assert_called_once_with() diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index 88f0ab193..55393209a 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 Qwen3-VL layer-50 text/vision encoder. The encoder is reimplemented on top of vLLM-style tensor-parallel building @@ -42,6 +43,10 @@ from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMetho from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion.layers.norm import RMSNorm as DiffusionRMSNorm +from vllm_omni.diffusion.offloader.module_residency import ( + BoundedAllocatorCache, + PinnedModuleStager, +) MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50 MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120 @@ -1211,12 +1216,18 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): if not self.is_loaded: return if getattr(self, "_omni_layerwise_enabled", False): - for name, child in self.vision.named_children(): - if name != "blocks": - child.to(self.device_target) - for name, child in self.text_model.named_children(): - if name != "layers": - child.to(self.device_target) + stager = getattr(self, "_omni_non_block_stager", None) + if stager is None: + hooks = getattr(self, "_omni_layerwise_hooks", ()) + pin_memory = bool(getattr(hooks[0], "pin_memory", True)) if hooks else True + stager = PinnedModuleStager( + self._omni_non_block_modules(), + self.device_target, + pin_memory=pin_memory, + cache_retention=getattr(self, "_omni_component_cache", None), + ) + self._omni_non_block_stager = stager + stager.load() return self.vision.to(self.device_target) self.text_model.to(self.device_target) @@ -1227,18 +1238,39 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): if getattr(self, "_omni_layerwise_enabled", False): for hook in self._omni_layerwise_hooks: hook.offload_layer() - for name, child in self.vision.named_children(): - if name != "blocks": + stager = getattr(self, "_omni_non_block_stager", None) + if stager is not None: + stager.offload() + else: + # The initial DLO setup reaches this branch before the first + # load. These modules were constructed on CPU, so this is not + # a device-to-host rematerialization. + for child in self._omni_non_block_modules(): child.to("cpu") - for name, child in self.text_model.named_children(): - if name != "layers": - child.to("cpu") - torch.accelerator.empty_cache() + self._release_omni_component_cache() return self.vision.to("cpu") self.text_model.to("cpu") torch.accelerator.empty_cache() + def _omni_non_block_modules(self) -> list[nn.Module]: + modules = [child for name, child in self.vision.named_children() if name != "blocks"] + modules.extend(child for name, child in self.text_model.named_children() if name != "layers") + return modules + + def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None: + self._omni_component_cache = cache + stager = getattr(self, "_omni_non_block_stager", None) + if stager is not None: + stager.set_cache_retention(cache) + + def _release_omni_component_cache(self) -> None: + cache = getattr(self, "_omni_component_cache", None) + if cache is None: + torch.accelerator.empty_cache() + else: + cache.release_if_needed() + def enable_omni_layerwise_offload(self, *, pin_memory: bool = True) -> None: """Stream the TP-local Qwen vision/text blocks for low-HBM serving. diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 115cdc79d..ff3990b32 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -46,6 +46,7 @@ from vllm_omni.diffusion.models.interface import ( ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.offloader import ( + BoundedAllocatorCache, OffloadPlan, apply_sequential_offload, remove_sequential_offload, @@ -952,6 +953,12 @@ class MiniMaxH3Pipeline( # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``. self.vae = self.video_vae + self._dlo_component_cache = None + if getattr(od_config, "enable_distributed_layerwise_offload", False): + self._dlo_component_cache = BoundedAllocatorCache(self.device) + for component in (self.text_encoder, self.video_vae, self.audio_vae): + component.set_omni_component_cache(self._dlo_component_cache) + self._quality_policy = MiniMaxH3QualityPolicy(od_config) self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self) @@ -1371,11 +1378,8 @@ class MiniMaxH3Pipeline( ): # Layerwise DiT offload already provides the low-residency encoder # phase used by the checkpoint reference. - self.text_encoder.load_to_device() - try: + with self._component_on_device(self.text_encoder): return self.text_encoder.encode_ids(input_ids, **vision_kwargs) - finally: - self.text_encoder.offload_to_cpu() # Keep both Qwen and DiT resident across requests. Moving either model # here makes encoder latency include a tens-of-gigabytes PCIe transfer, @@ -1433,13 +1437,35 @@ class MiniMaxH3Pipeline( yield return staged = self._uses_manual_component_offload() - if staged: - component.load_to_device() try: - yield - finally: if staged: - component.offload_to_cpu() + component.load_to_device() + yield + except BaseException: + if staged: + try: + component.offload_to_cpu() + except BaseException: + logger.exception("Failed to release %s after component failure", component.__class__.__name__) + cache = getattr(self, "_dlo_component_cache", None) + if cache is not None: + try: + cache.release_if_needed(force=True) + except BaseException: + logger.exception("Failed to release retained allocator cache after component failure") + raise + else: + if staged: + try: + component.offload_to_cpu() + except BaseException: + cache = getattr(self, "_dlo_component_cache", None) + if cache is not None: + try: + cache.release_if_needed(force=True) + except BaseException: + logger.exception("Failed to release retained allocator cache after offload failure") + raise def _encode_visual_conditions( self, diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index ffde5ee7c..b52417f06 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 remote-code VAE adapters and exact latent contracts.""" from __future__ import annotations @@ -21,7 +22,10 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor DistributedVaeMixin, ) from vllm_omni.diffusion.distributed.parallel_state import get_world_group -from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager +from vllm_omni.diffusion.offloader.module_residency import ( + BoundedAllocatorCache, + PinnedModuleStager, +) from .packed_tokens import minimax_h3_patchify_video_latent @@ -148,12 +152,21 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): else: self.remote.to(self._device_target) + def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None: + self._omni_component_cache = cache + if self._stager is not None: + self._stager.set_cache_retention(cache) + def offload_to_cpu(self) -> None: if self._stager is not None: self._stager.offload() else: self.remote.to("cpu") - torch.accelerator.empty_cache() + cache = getattr(self, "_omni_component_cache", None) + if cache is None: + torch.accelerator.empty_cache() + else: + cache.release_if_needed() def set_parallel_size( self, @@ -411,12 +424,21 @@ class MiniMaxH3AudioVAE(nn.Module): else: self.remote.to(self._device_target) + def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None: + self._omni_component_cache = cache + if self._stager is not None: + self._stager.set_cache_retention(cache) + def offload_to_cpu(self) -> None: if self._stager is not None: self._stager.offload() else: self.remote.to("cpu") - torch.accelerator.empty_cache() + cache = getattr(self, "_omni_component_cache", None) + if cache is None: + torch.accelerator.empty_cache() + else: + cache.release_if_needed() @torch.inference_mode() def encode_waveform( diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index f768cd14c..8145ac624 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -17,7 +17,7 @@ from .distributed_layerwise_backend import ( remove_distributed_block_hook, ) from .layerwise_backend import LayerWiseOffloadBackend -from .module_residency import PinnedModuleStager +from .module_residency import BoundedAllocatorCache, PinnedModuleStager from .offload_plan import OffloadPlan, get_offload_plan from .sequential_backend import ( ModelLevelOffloadBackend, @@ -46,6 +46,7 @@ __all__ = [ "DistributedLayerwiseOffloadBackend", "DistributedLayerwiseOffloadHook", "ModelLevelOffloadBackend", + "BoundedAllocatorCache", "PinnedModuleStager", "apply_sequential_offload", "remove_sequential_offload", diff --git a/vllm_omni/diffusion/offloader/module_residency.py b/vllm_omni/diffusion/offloader/module_residency.py index da32ab520..ab584c790 100644 --- a/vllm_omni/diffusion/offloader/module_residency.py +++ b/vllm_omni/diffusion/offloader/module_residency.py @@ -1,89 +1,271 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""On-demand module staging backed by an immutable pinned CPU snapshot.""" +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""On-demand module staging backed by immutable pinned CPU storage.""" from __future__ import annotations +from collections.abc import Iterable +from dataclasses import dataclass from itertools import chain from typing import Any import torch from torch import nn +from vllm.logger import init_logger from vllm_omni.platforms import current_omni_platform -from .tensor_utils import set_tensor_storage +from .tensor_utils import is_dtensor, set_tensor_storage + +logger = init_logger(__name__) -class PinnedModuleStager: - """Stage an immutable module without copying device weights back to CPU. +class BoundedAllocatorCache: + """Retain reusable allocator blocks without monopolizing device memory. - ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter - after each forward. In inference the weights are immutable, so retain one - pinned CPU master instead. ``load`` materializes fresh device tensors from - that master; ``offload`` only rebinds the Parameters/Buffers to it and - releases the device allocations. + Component offload normally calls ``empty_cache`` after every stage. That + makes the next stage return to the device allocator even though PyTorch's + cached blocks are immediately reusable. This policy keeps the cache while + both of these bounds hold: + + * cached-but-unallocated memory is at most 25% of device capacity; and + * at least 5% of device capacity is physically free. + + 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__( self, - module: nn.Module, + device: torch.device, + *, + max_cached_fraction: float = 0.25, + min_free_fraction: float = 0.05, + ) -> None: + if not 0.0 <= max_cached_fraction <= 1.0: + raise ValueError(f"max_cached_fraction must be in [0, 1], got {max_cached_fraction}") + if not 0.0 <= min_free_fraction <= 1.0: + raise ValueError(f"min_free_fraction must be in [0, 1], got {min_free_fraction}") + self.device = device + self.max_cached_fraction = max_cached_fraction + self.min_free_fraction = min_free_fraction + + def _should_release(self) -> bool: + reserved = int(torch.accelerator.memory_reserved(self.device)) + 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 > int(total * self.max_cached_fraction) or free < int(total * self.min_free_fraction) + + def release_if_needed(self, *, force: bool = False) -> bool: + """Release cached blocks when a bound is crossed or release is forced.""" + if not force: + try: + if not self._should_release(): + return False + except Exception as exc: + # Preserve the pre-retention behavior on platforms that do not + # expose allocator telemetry through torch.accelerator. + logger.debug("Allocator cache telemetry unavailable; releasing cache: %s", exc) + current_omni_platform.empty_cache() + return True + + +@dataclass(frozen=True) +class _TensorBinding: + target: torch.Tensor + dtype: torch.dtype + shape: tuple[int, ...] + stride: tuple[int, ...] + storage_offset: int + + +@dataclass +class _StorageGroup: + master: torch.Tensor + bindings: list[_TensorBinding] + + +class PinnedModuleStager: + """Stage immutable module groups without copying device weights to CPU. + + ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter + after each forward. In inference the weights are immutable, so retain one + pinned CPU master instead. ``load`` materializes device storage from that + master; ``offload`` only rebinds Parameters and buffers to the master. + + A module iterable is treated as one staging group. It uses one copy stream + and one reusable completion event. Tensors sharing storage keep their + shapes, strides, offsets, dtypes, and aliases across every transition. + """ + + def __init__( + self, + module: nn.Module | Iterable[nn.Module], device: torch.device, *, pin_memory: bool = True, copy_stream: Any | None = None, + cache_retention: BoundedAllocatorCache | None = None, ) -> None: - self.device = device - self.copy_stream = copy_stream or current_omni_platform.Stream() - self.loaded = False - self._entries: list[tuple[torch.Tensor, torch.Tensor]] = [] - self._device_tensors: list[torch.Tensor] = [] + modules = (module,) if isinstance(module, nn.Module) else tuple(module) + if not modules or not all(isinstance(item, nn.Module) for item in modules): + raise ValueError("PinnedModuleStager requires at least one nn.Module") - # parameters()/buffers() remove duplicate objects by default, which - # preserves tied weights and buffer aliases when their storage moves. - for target in chain(module.parameters(), module.buffers()): - master = target.detach() - if master.device.type != "cpu": - master = master.to("cpu") + self.device = device + self.copy_stream = copy_stream if copy_stream is not None else current_omni_platform.Stream() + self._ready_event = current_omni_platform.Event() + self.cache_retention = cache_retention + self.loaded = False + self._groups = self._snapshot_groups(modules, pin_memory=pin_memory) + self._device_storages: list[torch.Tensor] = [] + self._restore_masters() + + @staticmethod + def _local_tensor(target: torch.Tensor) -> torch.Tensor: + return target.to_local() if is_dtensor(target) else target + + @classmethod + def _snapshot_groups( + cls, + modules: tuple[nn.Module, ...], + *, + pin_memory: bool, + ) -> list[_StorageGroup]: + targets: list[torch.Tensor] = [] + seen_targets: set[int] = set() + for target in chain.from_iterable(chain(item.parameters(), item.buffers()) for item in modules): + if id(target) not in seen_targets: + seen_targets.add(id(target)) + targets.append(target) + + grouped: dict[tuple[Any, ...], tuple[torch.Tensor, list[_TensorBinding]]] = {} + for target in targets: + local = cls._local_tensor(target) + if local.is_meta: + raise ValueError("PinnedModuleStager cannot snapshot a meta tensor") + storage = local.untyped_storage() + if storage.nbytes() == 0: + storage_key: tuple[Any, ...] = ("empty", id(target)) + else: + storage_key = ( + local.device.type, + local.device.index, + storage.data_ptr(), + storage.nbytes(), + ) + binding = _TensorBinding( + target=target, + dtype=local.dtype, + shape=tuple(local.shape), + stride=tuple(local.stride()), + storage_offset=local.storage_offset(), + ) + if storage_key not in grouped: + grouped[storage_key] = (local, []) + grouped[storage_key][1].append(binding) + + groups: list[_StorageGroup] = [] + for source, bindings in grouped.values(): + storage = source.untyped_storage() + storage_view = torch.empty(0, dtype=torch.uint8, device=source.device).set_( + storage, + 0, + (storage.nbytes(),), + (1,), + ) + master = storage_view.detach() if storage_view.device.type == "cpu" else storage_view.to("cpu") if pin_memory and not master.is_pinned(): master = master.pin_memory() - set_tensor_storage(target, master) - self._entries.append((target, master)) + groups.append(_StorageGroup(master=master, bindings=bindings)) + return groups + + @staticmethod + def _view(backing: torch.Tensor, binding: _TensorBinding) -> torch.Tensor: + return torch.empty(0, dtype=binding.dtype, device=backing.device).set_( + backing.untyped_storage(), + binding.storage_offset, + binding.shape, + binding.stride, + ) + + def _bind(self, storages: list[torch.Tensor]) -> None: + for storage, group in zip(storages, self._groups): + for binding in group.bindings: + set_tensor_storage(binding.target, self._view(storage, binding)) + + def _restore_masters(self) -> None: + self._bind([group.master for group in self._groups]) + + def set_cache_retention(self, cache_retention: BoundedAllocatorCache | None) -> None: + self.cache_retention = cache_retention + + def _release_cache(self, *, force: bool = False) -> None: + if self.cache_retention is None: + current_omni_platform.empty_cache() + else: + self.cache_retention.release_if_needed(force=force) + + def _load_once(self) -> None: + device_storages = [torch.empty_like(group.master, device=self.device) for group in self._groups] + compute_stream = current_omni_platform.current_stream() + self.copy_stream.wait_stream(compute_stream) + with current_omni_platform.stream(self.copy_stream): + for device_storage, group in zip(device_storages, self._groups): + device_storage.copy_(group.master, non_blocking=group.master.is_pinned()) + self._ready_event.record(self.copy_stream) + + self._bind(device_storages) + compute_stream.wait_event(self._ready_event) + self._device_storages = device_storages + self.loaded = True + + def _cleanup_failed_load(self) -> None: + try: + current_omni_platform.synchronize() + except Exception: + logger.debug("Device synchronization failed while cleaning up module staging", exc_info=True) + self._restore_masters() + self._device_storages.clear() + self.loaded = False + self._release_cache(force=True) def load(self) -> None: if self.loaded: return - - # Allocate on the compute stream so the caching allocator can reuse - # memory released by the preceding encoder/DiT stage. - device_tensors = [torch.empty_like(master, device=self.device) for _, master in self._entries] - compute_stream = current_omni_platform.current_stream() - self.copy_stream.wait_stream(compute_stream) - ready = current_omni_platform.Event() - with current_omni_platform.stream(self.copy_stream): - for device_tensor, (_, master) in zip(device_tensors, self._entries): - device_tensor.copy_(master, non_blocking=master.is_pinned()) - ready.record(self.copy_stream) - - for device_tensor, (target, _) in zip(device_tensors, self._entries): - set_tensor_storage(target, device_tensor) - compute_stream.wait_event(ready) - self._device_tensors = device_tensors - self.loaded = True + try: + self._load_once() + except torch.OutOfMemoryError: + # A retained cache is normally reusable by this process, but an + # explicit flush gives external memory pressure one bounded retry. + self._cleanup_failed_load() + try: + self._load_once() + except BaseException: + self._cleanup_failed_load() + raise + except BaseException: + self._cleanup_failed_load() + raise def offload(self) -> None: if not self.loaded: return # The module has completed on the compute stream. Synchronize once at - # the stage boundary, then release HBM by rebinding to the unchanged - # host masters. No device-to-host transfer is necessary. - current_omni_platform.synchronize() - for target, master in self._entries: - set_tensor_storage(target, master) - self._device_tensors.clear() + # the stage boundary, then discard device storage without any D2H copy. + try: + current_omni_platform.synchronize() + self._restore_masters() + except BaseException: + self._device_storages.clear() + self.loaded = False + self._release_cache(force=True) + raise + self._device_storages.clear() self.loaded = False - current_omni_platform.empty_cache() + self._release_cache() -__all__ = ["PinnedModuleStager"] +__all__ = ["BoundedAllocatorCache", "PinnedModuleStager"]