[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>
ray24777 committed
Aug 26, 2026 at 18:37 UTC
61d4f56635e0f2faae4bd60e6da88a53c9d3affb
5 files changed
+113
-4
tests/basic_correctness/test_cpu_offload.py
+59
@@ -1,7 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from operator import attrgetter
5
+
6
import pytest
7
+import torch.nn as nn
8
9
import vllm.envs as envs
10
from vllm.model_executor.offloader import (
@@ -82,3 +85,59 @@ def test_mrv2_weight_offloading(
85
finally:
86
set_offloader(original_offloader)
87
envs.disable_envs_cache()
88
+
89
+
90
+def _is_offloaded(p: nn.Parameter) -> bool:
91
+ return p.device.type == "cpu" or getattr(p, "_vllm_is_uva_offloaded", False)
92
+
93
+
94
+@pytest.mark.parametrize("disable_uva", [False, True])
95
+def test_tower_weight_offloading(vllm_runner, monkeypatch, disable_uva):
96
+ """`cpu_offload_params` segments must reach towers built outside make_layers.
97
+
98
+ Regression test: `wrap_modules` was only called from `make_layers`, so a
99
+ directly-constructed vision tower never reached the offloader and segments
100
+ targeting it silently matched nothing.
101
+ """
102
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
103
+ if disable_uva:
104
+ monkeypatch.setenv("VLLM_WEIGHT_OFFLOADING_DISABLE_UVA", "1")
105
+ envs.disable_envs_cache()
106
+ original_offloader = get_offloader()
107
+
108
+ try:
109
+ with vllm_runner(
110
+ "Qwen/Qwen3.5-0.8B",
111
+ enforce_eager=True,
112
+ # allocate more vram as Qwen 3.5 has 1.6 GiB of weights
113
+ gpu_memory_utilization=0.3,
114
+ max_model_len=128,
115
+ max_num_seqs=1,
116
+ enable_prefix_caching=False,
117
+ cpu_offload_gb=1,
118
+ cpu_offload_params={"visual"},
119
+ ) as vllm_model:
120
+ engine_core = vllm_model.llm.llm_engine.engine_core.engine_core
121
+ model_runner = engine_core.model_executor.driver_worker.worker.model_runner
122
+
123
+ offloader = get_offloader()
124
+ assert isinstance(offloader, UVAOffloader)
125
+ assert offloader.cpu_offload_bytes > 0
126
+
127
+ model = model_runner.get_model()
128
+ assert model._tower_model_names
129
+ for name in model._tower_model_names:
130
+ tower = attrgetter(name)(model)
131
+ params = list(tower.parameters())
132
+ assert params
133
+ assert all(_is_offloaded(p) for p in params)
134
+ if disable_uva:
135
+ # non-UVA fallback: weights live on CPU and are moved
136
+ # back on first forward
137
+ assert "forward" in vars(tower)
138
+
139
+ # The language model must stay resident.
140
+ assert not any(_is_offloaded(p) for p in model.language_model.parameters())
141
+ finally:
142
+ set_offloader(original_offloader)
143
+ envs.disable_envs_cache()
vllm/model_executor/models/interfaces.py
+15
@@ -11,6 +11,7 @@ from collections.abc import (
11
)
12
from contextlib import ExitStack, contextmanager, nullcontext
13
from dataclasses import dataclass
14
+from operator import attrgetter
15
from typing import (
16
TYPE_CHECKING,
17
Any,
@@ -321,7 +322,13 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):
322
323
If `targets` is set, instead include descendants that are an instance
324
of `targets`, even if they aren't direct children.
325
+
326
+ Marked components are also routed through the active offloader (when
327
+ it supports tower offloading), since `make_layers` only ever sees the
328
+ decoder layer stack.
329
"""
330
+ from vllm.model_executor.offloader import get_offloader
331
+
332
from .utils import StageMissingLayer, collect_children, no_init_weights
333
334
if isinstance(modalities, str):
@@ -348,6 +355,14 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):
355
356
self._tower_model_names = children_names
357
358
+ # Towers are constructed directly, so `make_layers` never routes them
359
+ # through the offloader. Do it here, at the same construction stage,
360
+ # so offloaded tower weights are never allocated on the device.
361
+ offloader = get_offloader()
362
+ if offloader.supports_tower_offload:
363
+ for name in children_names:
364
+ offloader.wrap_modules(iter([attrgetter(name)(self)]), prefix=name)
365
+
366
@contextmanager
367
def _mark_composite_model(
368
self,
vllm/model_executor/offloader/base.py
+15
@@ -51,15 +51,29 @@ class BaseOffloader(ABC):
51
inference. Different strategies trade memory for compute/transfer time.
52
"""
53
54
+ supports_tower_offload: bool = False
55
+ """Whether `wrap_modules` also accepts modules routed by
56
+ `SupportsMultiModal._mark_tower_model`, outside the `make_layers` call.
57
+
58
+ Offloaders whose `wrap_modules` may only be called on the decoder layer
59
+ stack (e.g. `PrefetchOffloader`, which schedules prefetches over a
60
+ circular layer stack) must keep this `False`.
61
+ """
62
+
63
@abstractmethod
64
def wrap_modules(
65
self,
66
modules_generator: Generator[nn.Module, None, None],
67
+ prefix: str = "",
68
) -> list[nn.Module]:
69
"""Wrap modules with offloading logic.
70
71
Args:
72
modules_generator: Generator yielding modules to potentially offload.
73
+ prefix: Name prefix prepended to parameter names before matching
74
+ them against the offloading parameter set. Used when the
75
+ modules are not the full model, so that name segments stay
76
+ fully qualified (e.g. `visual` for a tower module).
77
78
Returns:
79
List of modules, potentially with offloading hooks installed.
@@ -99,6 +113,7 @@ class NoopOffloader(BaseOffloader):
113
def wrap_modules(
114
self,
115
modules_generator: Generator[nn.Module, None, None],
116
+ prefix: str = "",
117
) -> list[nn.Module]:
118
"""Return modules unchanged."""
119
return list(modules_generator)
vllm/model_executor/offloader/prefetch.py
+7
-1
@@ -163,12 +163,16 @@ class PrefetchOffloader(BaseOffloader):
163
def wrap_modules(
164
self,
165
modules_generator: Generator[nn.Module, None, None],
166
+ prefix: str = "",
167
) -> list[nn.Module]:
168
"""Wrap modules with prefetch offloading logic."""
169
assert len(self.module_offloaders) == 0, (
170
"wrap_modules should only be called once"
171
)
172
173
+ if prefix:
174
+ prefix = f"{prefix}."
175
+
176
all_modules = []
177
offload_modules = []
178
@@ -182,7 +186,9 @@ class PrefetchOffloader(BaseOffloader):
186
whitelist = [
187
name
188
for name, _ in module.named_parameters()
185
- if any(f".{p}." in f".{name}." for p in self.offload_params)
189
+ if any(
190
+ f".{p}." in f".{prefix}{name}." for p in self.offload_params
191
+ )
192
]
193
else:
194
whitelist = [name for name, _ in module.named_parameters()]
vllm/model_executor/offloader/uva.py
+17
-3
@@ -34,6 +34,8 @@ class UVAOffloader(BaseOffloader):
34
offload. If empty, all parameters are eligible up to the byte limit.
35
"""
36
37
+ supports_tower_offload = True
38
+
39
def __init__(
40
self,
41
cpu_offload_max_bytes: int,
@@ -51,9 +53,14 @@ class UVAOffloader(BaseOffloader):
53
def wrap_modules(
54
self,
55
modules_generator: Generator[nn.Module, None, None],
56
+ prefix: str = "",
57
) -> list[nn.Module]:
58
"""Wrap modules with UVA offloading."""
56
- modules = [self._maybe_offload_to_cpu(module) for module in modules_generator]
59
+ if prefix:
60
+ prefix = f"{prefix}."
61
+ modules = [
62
+ self._maybe_offload_to_cpu(module, prefix) for module in modules_generator
63
+ ]
64
if self.cpu_offload_bytes > 0:
65
logger.info(
66
"Total CPU offloaded parameters: %s",
@@ -61,7 +68,7 @@ class UVAOffloader(BaseOffloader):
68
)
69
return modules
70
64
- def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module:
71
+ def _maybe_offload_to_cpu(self, module: nn.Module, prefix: str = "") -> nn.Module:
72
"""Offload module parameters to CPU using UVA if budget allows."""
73
if (params := next(module.parameters(), None)) is None:
74
return module
@@ -83,13 +90,20 @@ class UVAOffloader(BaseOffloader):
90
# one module might have some parameters offloaded and some not
91
break
92
93
+ # Skip parameters an earlier wrap_modules call already offloaded.
94
+ # The UVA path leaves p.device as the accelerator (a view of CPU
95
+ # memory), so the marker is the only way to recognize those.
96
+ if p.device.type == "cpu" or getattr(p, "_vllm_is_uva_offloaded", False):
97
+ continue
98
+
99
if self.cpu_offload_params:
100
# Check if parameter belongs to the offloading set
101
# Add dots here to ensure we match full segments only
102
# e.g., "experts.w2_weight" matches "mlp.experts.w2_weight"
103
# but not "mlp.experts.w2_weight_scale"
104
should_offload = any(
92
- f".{param}." in f".{name}." for param in self.cpu_offload_params
105
+ f".{param}." in f".{prefix}{name}."
106
+ for param in self.cpu_offload_params
107
)
108
if not should_offload:
109
continue