diff --git a/tests/basic_correctness/test_cpu_offload.py b/tests/basic_correctness/test_cpu_offload.py index f3ea25e58a..a63f1189aa 100644 --- a/tests/basic_correctness/test_cpu_offload.py +++ b/tests/basic_correctness/test_cpu_offload.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from operator import attrgetter + import pytest +import torch.nn as nn import vllm.envs as envs from vllm.model_executor.offloader import ( @@ -82,3 +85,59 @@ def test_mrv2_weight_offloading( finally: set_offloader(original_offloader) envs.disable_envs_cache() + + +def _is_offloaded(p: nn.Parameter) -> bool: + return p.device.type == "cpu" or getattr(p, "_vllm_is_uva_offloaded", False) + + +@pytest.mark.parametrize("disable_uva", [False, True]) +def test_tower_weight_offloading(vllm_runner, monkeypatch, disable_uva): + """`cpu_offload_params` segments must reach towers built outside make_layers. + + Regression test: `wrap_modules` was only called from `make_layers`, so a + directly-constructed vision tower never reached the offloader and segments + targeting it silently matched nothing. + """ + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + if disable_uva: + monkeypatch.setenv("VLLM_WEIGHT_OFFLOADING_DISABLE_UVA", "1") + envs.disable_envs_cache() + original_offloader = get_offloader() + + try: + with vllm_runner( + "Qwen/Qwen3.5-0.8B", + enforce_eager=True, + # allocate more vram as Qwen 3.5 has 1.6 GiB of weights + gpu_memory_utilization=0.3, + max_model_len=128, + max_num_seqs=1, + enable_prefix_caching=False, + cpu_offload_gb=1, + cpu_offload_params={"visual"}, + ) as vllm_model: + engine_core = vllm_model.llm.llm_engine.engine_core.engine_core + model_runner = engine_core.model_executor.driver_worker.worker.model_runner + + offloader = get_offloader() + assert isinstance(offloader, UVAOffloader) + assert offloader.cpu_offload_bytes > 0 + + model = model_runner.get_model() + assert model._tower_model_names + for name in model._tower_model_names: + tower = attrgetter(name)(model) + params = list(tower.parameters()) + assert params + assert all(_is_offloaded(p) for p in params) + if disable_uva: + # non-UVA fallback: weights live on CPU and are moved + # back on first forward + assert "forward" in vars(tower) + + # The language model must stay resident. + assert not any(_is_offloaded(p) for p in model.language_model.parameters()) + finally: + set_offloader(original_offloader) + envs.disable_envs_cache() diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index d9f84be396..b3102f567a 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -11,6 +11,7 @@ from collections.abc import ( ) from contextlib import ExitStack, contextmanager, nullcontext from dataclasses import dataclass +from operator import attrgetter from typing import ( TYPE_CHECKING, Any, @@ -321,7 +322,13 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol): If `targets` is set, instead include descendants that are an instance of `targets`, even if they aren't direct children. + + Marked components are also routed through the active offloader (when + it supports tower offloading), since `make_layers` only ever sees the + decoder layer stack. """ + from vllm.model_executor.offloader import get_offloader + from .utils import StageMissingLayer, collect_children, no_init_weights if isinstance(modalities, str): @@ -348,6 +355,14 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol): self._tower_model_names = children_names + # Towers are constructed directly, so `make_layers` never routes them + # through the offloader. Do it here, at the same construction stage, + # so offloaded tower weights are never allocated on the device. + offloader = get_offloader() + if offloader.supports_tower_offload: + for name in children_names: + offloader.wrap_modules(iter([attrgetter(name)(self)]), prefix=name) + @contextmanager def _mark_composite_model( self, diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index ceff60cd4c..cb9c35b79c 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -51,15 +51,29 @@ class BaseOffloader(ABC): inference. Different strategies trade memory for compute/transfer time. """ + supports_tower_offload: bool = False + """Whether `wrap_modules` also accepts modules routed by + `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`. + """ + @abstractmethod def wrap_modules( self, modules_generator: Generator[nn.Module, None, None], + prefix: str = "", ) -> list[nn.Module]: """Wrap modules with offloading logic. Args: modules_generator: Generator yielding modules to potentially offload. + prefix: Name prefix prepended to parameter names before matching + them against the offloading parameter set. Used when the + modules are not the full model, so that name segments stay + fully qualified (e.g. `visual` for a tower module). Returns: List of modules, potentially with offloading hooks installed. @@ -99,6 +113,7 @@ class NoopOffloader(BaseOffloader): def wrap_modules( self, modules_generator: Generator[nn.Module, None, None], + prefix: str = "", ) -> list[nn.Module]: """Return modules unchanged.""" return list(modules_generator) diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index 466d8c13ce..810acff6bc 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -163,12 +163,16 @@ class PrefetchOffloader(BaseOffloader): def wrap_modules( self, modules_generator: Generator[nn.Module, None, None], + prefix: str = "", ) -> list[nn.Module]: """Wrap modules with prefetch offloading logic.""" assert len(self.module_offloaders) == 0, ( "wrap_modules should only be called once" ) + if prefix: + prefix = f"{prefix}." + all_modules = [] offload_modules = [] @@ -182,7 +186,9 @@ class PrefetchOffloader(BaseOffloader): whitelist = [ name for name, _ in module.named_parameters() - if any(f".{p}." in f".{name}." for p in self.offload_params) + if any( + f".{p}." in f".{prefix}{name}." for p in self.offload_params + ) ] else: whitelist = [name for name, _ in module.named_parameters()] diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index 51eb1a14fc..5107870ffa 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -34,6 +34,8 @@ class UVAOffloader(BaseOffloader): offload. If empty, all parameters are eligible up to the byte limit. """ + supports_tower_offload = True + def __init__( self, cpu_offload_max_bytes: int, @@ -51,9 +53,14 @@ class UVAOffloader(BaseOffloader): def wrap_modules( self, modules_generator: Generator[nn.Module, None, None], + prefix: str = "", ) -> list[nn.Module]: """Wrap modules with UVA offloading.""" - modules = [self._maybe_offload_to_cpu(module) for module in modules_generator] + if prefix: + prefix = f"{prefix}." + modules = [ + self._maybe_offload_to_cpu(module, prefix) for module in modules_generator + ] if self.cpu_offload_bytes > 0: logger.info( "Total CPU offloaded parameters: %s", @@ -61,7 +68,7 @@ class UVAOffloader(BaseOffloader): ) return modules - def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: + def _maybe_offload_to_cpu(self, module: nn.Module, prefix: str = "") -> nn.Module: """Offload module parameters to CPU using UVA if budget allows.""" if (params := next(module.parameters(), None)) is None: return module @@ -83,13 +90,20 @@ class UVAOffloader(BaseOffloader): # one module might have some parameters offloaded and some not break + # Skip parameters an earlier wrap_modules call already offloaded. + # The UVA path leaves p.device as the accelerator (a view of CPU + # memory), so the marker is the only way to recognize those. + if p.device.type == "cpu" or getattr(p, "_vllm_is_uva_offloaded", False): + continue + if self.cpu_offload_params: # Check if parameter belongs to the offloading set # Add dots here to ensure we match full segments only # e.g., "experts.w2_weight" matches "mlp.experts.w2_weight" # but not "mlp.experts.w2_weight_scale" should_offload = any( - f".{param}." in f".{name}." for param in self.cpu_offload_params + f".{param}." in f".{prefix}{name}." + for param in self.cpu_offload_params ) if not should_offload: continue