From f5497ecb0ab526ed95511ca781e698226aa0b808 Mon Sep 17 00:00:00 2001 From: MOVZX Date: Fri, 16 Oct 2026 14:32:14 +0700 Subject: [PATCH 104/105] [Offloader] Stream multimodal towers from host RAM with bulk-copy Add --offload-towers: place multimodal tower weights in pinned host memory and bulk-copy each unit to the GPU only during its forward, keeping towers out of VRAM between encodes. Unlike UVA zero-copy, whole-tensor H2D transfers avoid re-reading operand tiles over PCIe inside the GEMM inner loop. Tower streaming is independent of layer group scheduling, so group_size may be 0 for tower-only offload. Co-authored-by: Hinata Signed-off-by: MOVZX --- .../model_executor/offloader/test_prefetch.py | 65 +++++- vllm/config/offload.py | 30 ++- vllm/engine/arg_utils.py | 5 + vllm/model_executor/offloader/base.py | 12 +- vllm/model_executor/offloader/prefetch.py | 209 +++++++++++++++++- 5 files changed, 307 insertions(+), 14 deletions(-) diff --git a/tests/model_executor/offloader/test_prefetch.py b/tests/model_executor/offloader/test_prefetch.py index 560ac8d..41fc6c2 100644 --- a/tests/model_executor/offloader/test_prefetch.py +++ b/tests/model_executor/offloader/test_prefetch.py @@ -2,8 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch +import torch.nn as nn -from vllm.model_executor.offloader.prefetch import _get_next_prefetch_index +from vllm.model_executor.offloader.prefetch import ( + _get_next_prefetch_index, + _TowerUnitStreamer, +) pytestmark = pytest.mark.skip_global_cleanup @@ -35,3 +40,61 @@ def test_next_prefetch_index_preserves_slot_ownership( target % prefetch_step == index % prefetch_step for index, target in enumerate(targets) ) + + +def test_tower_unit_streamer_roundtrip_cpu(monkeypatch: pytest.MonkeyPatch) -> None: + """Host-resident weights must compute identically to resident weights, + and be back on host storage after each forward.""" + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + lin = nn.Linear(4, 4, bias=False) + weight = lin.weight.detach().clone() + + streamer = _TowerUnitStreamer(lin, ["weight"], torch.device("cpu")) + streamer.offload_to_host_now() + assert lin.weight.device.type == "cpu" + + streamer.sync_storage() + assert streamer.offloaded_bytes == weight.numel() * weight.element_size() + + streamer.hook_forward() + x = torch.randn(2, 4) + out = lin(x) + assert torch.allclose(out, x @ weight.T) + # Weights released back to host after the forward. + assert lin.weight.device.type == "cpu" + + # Re-running post_init-style hooking must not double-wrap the forward. + streamer.hook_forward() + assert torch.allclose(lin(x), out) + + +def test_offload_config_tower_only_prefetch_is_valid() -> None: + from vllm.config.offload import OffloadConfig, PrefetchOffloadConfig + + cfg = OffloadConfig( + offload_backend="prefetch", + prefetch=PrefetchOffloadConfig(offload_towers=True), + ) + assert cfg.prefetch.offload_towers + assert cfg.prefetch.offload_group_size == 0 + + +def test_offload_config_prefetch_requires_layers_or_towers() -> None: + from vllm.config.offload import OffloadConfig + + with pytest.raises(ValueError): + OffloadConfig(offload_backend="prefetch") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_create_offloader_selects_prefetch_for_towers() -> None: + from vllm.config.offload import OffloadConfig, PrefetchOffloadConfig + from vllm.model_executor.offloader.base import create_offloader + from vllm.model_executor.offloader.prefetch import PrefetchOffloader + + offloader = create_offloader( + OffloadConfig(prefetch=PrefetchOffloadConfig(offload_towers=True)) + ) + assert isinstance(offloader, PrefetchOffloader) + assert offloader.supports_tower_offload diff --git a/vllm/config/offload.py b/vllm/config/offload.py index ad65e8a..bb465e3 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -73,6 +73,19 @@ class PrefetchOffloadConfig: parameters of each offloaded layer are offloaded. Uses segment matching: "w13_weight" matches "mlp.experts.w13_weight" but not "mlp.experts.w13_weight_scale". + Applies to decoder layers only; `offload_towers` offloads the full + tower regardless of this set. + """ + + offload_towers: bool = False + """Offload multimodal tower weights (vision/audio encoders) to pinned + host memory, and bulk-copy them to the GPU during each encoder forward. + Towers hold no VRAM between image encodes. Unlike UVA zero-copy, weights + are streamed with whole-tensor H2D copies instead of being re-read tile + by tile inside the GEMM inner loops, so encode cost stays near the + resident-tower baseline. + Works together with, or instead of, layer prefetch: + `offload_group_size` may be 0 when only towers are offloaded. """ @@ -83,7 +96,8 @@ class OffloadConfig: offload_backend: OffloadBackend = "auto" """The backend for weight offloading. Options: - "auto": Selects based on which sub-config has non-default values - (prefetch if offload_group_size > 0, uva if cpu_offload_gb > 0). + (prefetch if offload_group_size > 0 or offload_towers, uva if + cpu_offload_gb > 0). - "uva": UVA (Unified Virtual Addressing) zero-copy offloading. - "prefetch": Async prefetch with group-based layer offloading. """ @@ -97,7 +111,15 @@ class OffloadConfig: @model_validator(mode="after") def validate_offload_config(self) -> "OffloadConfig": """Validate offload configuration constraints.""" - if self.offload_backend == "prefetch" or self.prefetch.offload_group_size > 0: + layer_offload = self.prefetch.offload_group_size > 0 + if self.offload_backend == "prefetch" and not ( + layer_offload or self.prefetch.offload_towers + ): + raise ValueError( + "offload_backend='prefetch' requires offload_group_size > 0" + " or offload_towers=True" + ) + if layer_offload: if self.prefetch.offload_num_in_group > self.prefetch.offload_group_size: raise ValueError( f"offload_num_in_group ({self.prefetch.offload_num_in_group})" @@ -114,7 +136,9 @@ class OffloadConfig: # Warn if both backends have non-default values uva_active = self.uva.cpu_offload_gb > 0 - prefetch_active = self.prefetch.offload_group_size > 0 + prefetch_active = ( + self.prefetch.offload_group_size > 0 or self.prefetch.offload_towers + ) if self.offload_backend == "uva" and prefetch_active: warnings.warn( "Prefetch offload fields are set but offload_backend='uva'. " diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index b420e84..ae1e2df 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -551,6 +551,7 @@ class EngineArgs: offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params") + offload_towers: bool = PrefetchOffloadConfig.offload_towers gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes max_num_batched_tokens: int | None = None @@ -1356,6 +1357,9 @@ class EngineArgs: offload_group.add_argument( "--offload-params", **prefetch_kwargs["offload_params"] ) + offload_group.add_argument( + "--offload-towers", **prefetch_kwargs["offload_towers"] + ) # Multimodal related configs multimodal_kwargs = get_kwargs(MultiModalConfig) @@ -2623,6 +2627,7 @@ class EngineArgs: offload_num_in_group=self.offload_num_in_group, offload_prefetch_step=self.offload_prefetch_step, offload_params=self.offload_params, + offload_towers=self.offload_towers, ), ) diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index cb9c35b..64eba91 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -56,8 +56,9 @@ class BaseOffloader(ABC): `SupportsMultiModal._mark_tower_model`, outside the `make_layers` call. Offloaders whose `wrap_modules` may only be called on the decoder layer - stack (e.g. `PrefetchOffloader`, which schedules prefetches over a - circular layer stack) must keep this `False`. + stack must keep this `False`. `PrefetchOffloader` enables it per-instance + when `offload_towers` is configured, via a tower path that does not use + the circular layer stack or the static buffer pool. """ @abstractmethod @@ -142,8 +143,8 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: """Create an offloader based on the offload configuration. Uses the explicit ``offload_backend`` selector. When set to ``"auto"``, - selects prefetch if ``offload_group_size > 0``, UVA if - ``cpu_offload_gb > 0``, otherwise noop. + selects prefetch if ``offload_group_size > 0`` or ``offload_towers`` is + set, UVA if ``cpu_offload_gb > 0``, otherwise noop. """ from vllm.model_executor.offloader.prefetch import PrefetchOffloader from vllm.model_executor.offloader.uva import UVAOffloader @@ -153,7 +154,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: prefetch = offload_config.prefetch if backend == "auto": - if prefetch.offload_group_size > 0: + if prefetch.offload_group_size > 0 or prefetch.offload_towers: backend = "prefetch" elif uva.cpu_offload_gb > 0: backend = "uva" @@ -166,6 +167,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: num_in_group=prefetch.offload_num_in_group, prefetch_step=prefetch.offload_prefetch_step, offload_params=prefetch.offload_params, + offload_towers=prefetch.offload_towers, mode="cpu", ) elif backend == "uva": diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index e682fd0..f276673 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -150,6 +150,10 @@ class PrefetchOffloader(BaseOffloader): group_size: Group every N layers together. num_in_group: Offload this many layers per group (last N of each group). prefetch_step: Number of layers to prefetch ahead. + offload_towers: Also offload multimodal tower weights to pinned host + memory, streaming them to the GPU during each encoder forward + (bulk-copy). Group scheduling still applies to decoder layers + only; `group_size` may be 0 when only towers are offloaded. mode: Offload mode ("cpu" is currently supported). """ @@ -159,19 +163,24 @@ class PrefetchOffloader(BaseOffloader): num_in_group: int, prefetch_step: int, offload_params: set[str] | None = None, + offload_towers: bool = False, mode: str = "cpu", ): self.group_size = group_size self.num_in_group = num_in_group self.prefetch_step = prefetch_step self.offload_params = offload_params or set() + self.offload_towers = offload_towers self.mode = mode + if offload_towers: + self.supports_tower_offload = True # Copy stream for async H2D transfers self.copy_stream = torch.cuda.Stream() # Module offloaders and buffer pool (populated in wrap_modules/post_init) self.module_offloaders: list[_ModuleOffloader] = [] + self.tower_streamers: list[_TowerUnitStreamer] = [] self.buffer_pool: StaticBufferPool | None = None self.total_offloaded_bytes = 0 @@ -181,17 +190,29 @@ class PrefetchOffloader(BaseOffloader): prefix: str = "", ) -> list[nn.Module]: """Wrap modules with prefetch offloading logic.""" - assert len(self.module_offloaders) == 0, ( - "wrap_modules should only be called once" - ) + modules = list(modules_generator) if prefix: - prefix = f"{prefix}." + # Routed by _mark_tower_model: a whole tower module, handled by + # the streaming path, independent of the decoder group schedule. + if self.offload_towers: + for module in modules: + self._wrap_tower(module, prefix) + return modules + + if self.group_size <= 0: + # Prefetch backend selected for towers only; decoder layers + # stay resident on the device. + return modules + + assert len(self.module_offloaders) == 0, ( + "wrap_modules should only be called once for the decoder stack" + ) all_modules = [] offload_modules = [] - for module_index, module in enumerate(modules_generator): + for module_index, module in enumerate(modules): all_modules.append(module) # Select layers to offload based on group pattern @@ -227,6 +248,39 @@ class PrefetchOffloader(BaseOffloader): return all_modules + def _wrap_tower(self, tower: nn.Module, prefix: str): + """Split a tower into parameter-bearing units and place each one + in pinned host RAM until its forward executes (bulk-copy). + """ + first_param = next(tower.parameters(), None) + if first_param is None: + # Tower was staged out (e.g. --limit-mm-per-prompt zero for all + # its modalities) or carries no parameters. + return + device = first_param.device + if device.type == "cpu": + return + + num_units = 0 + for submodule in tower.modules(): + param_names = [ + name for name, _ in submodule.named_parameters(recurse=False) + ] + if not param_names: + continue + streamer = _TowerUnitStreamer(submodule, param_names, device) + streamer.offload_to_host_now() + self.tower_streamers.append(streamer) + num_units += 1 + + logger.info( + "[PrefetchOffloader] Queued tower '%s' for host streaming: " + "%d parameter units (weights move to the GPU only during " + "encoder forwards).", + prefix, + num_units, + ) + def _hook_module_forward(self, index: int, module: nn.Module): """Hook module's forward with torch.compile-compatible sync.""" original_forward = module.forward @@ -339,6 +393,21 @@ class PrefetchOffloader(BaseOffloader): (in _CpuParamOffloader.__init__), so GPU memory is available for the static buffer pool. """ + # Tower streamers: re-snapshot host storage from the final processed + # weights, then install the streaming forwards. + for streamer in self.tower_streamers: + streamer.sync_storage() + for streamer in self.tower_streamers: + streamer.hook_forward() + if self.tower_streamers: + tower_bytes = sum(s.offloaded_bytes for s in self.tower_streamers) + self.total_offloaded_bytes += tower_bytes + logger.info_once( + f"[PrefetchOffloader] Tower offload active: " + f"{len(self.tower_streamers)} units, " + f"{tower_bytes / 1e9:.4f} GB in pinned host RAM." + ) + # Sync CPU storage with current param.data BEFORE collecting param info. # This is needed because process_weights_after_loading may have: # 1. Transformed weights (quantization, transpose, etc.) @@ -390,6 +459,136 @@ class PrefetchOffloader(BaseOffloader): self.module_offloaders[i].start_onload_to_static() +class _TowerUnitStreamer: + """Bulk-copy placement for one parameter-bearing submodule of a tower. + + The unit's parameters live in pinned host memory between forwards. + On each forward, they are bulk-copied to freshly allocated GPU tensors + on the compute stream (so copies are stream-ordered ahead of the unit's + kernels with no explicit sync) and released right after. At steady + state only the unit currently executing holds VRAM, and freed blocks + are recycled by the caching allocator for the next unit's weights. + + This is the bulk-copy alternative to UVA zero-copy for multimodal + towers: whole-tensor H2D transfers instead of operand tiles re-read + over PCIe inside the GEMM inner loop. + + Interaction notes: + - If an encoder CUDA graph captures a streamed unit, correctness is + preserved (the memcpy is captured and replays from host memory), but + the captured allocations persist in the graph pool, so the VRAM + saving is lost. Prefer eager encoder execution (`cudagraph_mm_encoder` + off) together with tower offloading. + - Sleep/wake weight reload is not intercepted; storage is re-synced + from `param.data` at `post_init` only. + """ + + def __init__(self, unit: nn.Module, param_names: list[str], device: Any): + self.unit = unit + self.param_names = param_names + self.device = device + self.offloaded_bytes = 0 + # name -> pinned host tensor holding the weights between forwards. + # Populated/refreshed in sync_storage() (after weight loading and + # process_weights_after_loading), which is the authoritative point. + self._storages: dict[str, torch.Tensor] = {} + + def _get_param(self, name: str) -> nn.Parameter: + return getattr(self.unit, name) + + @staticmethod + def _empty_cpu_like(data: torch.Tensor, pin_memory: bool) -> torch.Tensor: + return torch.empty_strided( + size=data.size(), + stride=data.stride(), + dtype=data.dtype, + layout=data.layout, + device="cpu", + pin_memory=pin_memory, + ) + + def offload_to_host_now(self) -> None: + """Move current parameter data to pinned host storage during + construction, so weight loading writes to host memory and the + device never accumulates tower weights. + """ + pin_memory = should_pin_memory() + for name in self.param_names: + p = self._get_param(name) + if p.device.type == "cpu": + continue + cpu_data = self._empty_cpu_like(p.data, pin_memory) + cpu_data.copy_(p.data) + p.data = cpu_data + + def sync_storage(self) -> None: + """Re-snapshot host storage from param.data. + + Called from post_init() after weight loading and + process_weights_after_loading. device_loading_context may have + replaced param.data with fresh (already pinned) CPU tensors, and + quantization processing may have left data on the device; handle + both so _storages always references the final processed weights. + """ + self.offloaded_bytes = 0 + for name in list(self.param_names): + try: + p = self._get_param(name) + except AttributeError: + # Parameter deleted by process_weights_after_loading. + self.param_names.remove(name) + continue + data = p.data + if data.device.type != "cpu": + cpu_data = self._empty_cpu_like(data, pin_memory=should_pin_memory()) + cpu_data.copy_(data) + data = cpu_data + p.data = data + elif should_pin_memory() and not data.is_pinned(): + pinned = self._empty_cpu_like(data, pin_memory=True) + pinned.copy_(data) + data = pinned + p.data = data + self._storages[name] = data + self.offloaded_bytes += data.numel() * data.element_size() + + def hook_forward(self) -> None: + if getattr(self.unit, "_vllm_tower_streamed", False): + return # post_init ran more than once + original_forward = self.unit.forward + streamer = self + + def forward(*args, **kwargs): + if torch.cuda.is_current_stream_capturing(): + logger.warning_once( + "[_TowerUnitStreamer] Streaming a tower unit during " + "CUDA graph capture: results stay correct but the " + "weights will be captured into the graph pool, so no " + "VRAM is saved. Disable encoder CUDA graphs " + "(cudagraph_mm_encoder) to keep the saving." + ) + for name in streamer.param_names: + storage = streamer._storages.get(name) + if storage is None: + continue + # Async H2D on the compute stream: ordered ahead of this + # unit's kernels, no explicit synchronization needed. + streamer._get_param(name).data = storage.to( + streamer.device, non_blocking=True + ) + try: + return original_forward(*args, **kwargs) + finally: + for name in streamer.param_names: + storage = streamer._storages.get(name) + if storage is None: + continue + streamer._get_param(name).data = storage + + self.unit._vllm_tower_streamed = True + self.unit.forward = forward + + class _ModuleOffloader: """Manages offloading for a single module. -- 2.47.3