[Offloader] Offload submodules that make_layers never reaches (#53120)

Signed-off-by: ray24777 <103923677+ray24777@users.noreply.github.com>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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()

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,

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)

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()]

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