[Model] Optimize MiniMax-H3 DLO component lifecycle

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

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",

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"]