A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 06:00 UTC
[Bugfix][Pooling] Fix silent weight corruption with buffer-reusing iterators (#39650)
Signed-off-by: Pedram Razavi <pedram.razavi@gmail.com>
Pedram Razavi committed
Apr 13, 2026 at 12:37 UTC
3693f922ff0e15b1e395cdc862f7c6e492a11bf4
2 files changed
+154
-2
tests/models/test_adapters.py
new
+148
@@ -0,0 +1,148 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for model adapter weight loading (adapters.py)."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.models.adapters import _create_pooling_model_cls
9
+from vllm.model_executor.models.utils import AutoWeightsLoader, StageMissingLayer
10
+
11
+pytestmark = pytest.mark.cpu_test
12
+
13
+
14
+class SimpleInnerModel(torch.nn.Module):
15
+ def __init__(self):
16
+ super().__init__()
17
+ self.embed = torch.nn.Linear(4, 8, bias=False)
18
+ self.layer0 = torch.nn.Linear(8, 8, bias=False)
19
+ self.layer1 = torch.nn.Linear(8, 8, bias=False)
20
+ self.norm = torch.nn.Linear(8, 4, bias=False)
21
+
22
+ def load_weights(self, weights):
23
+ params = dict(self.named_parameters())
24
+ loaded = set()
25
+ for name, tensor in weights:
26
+ if name in params:
27
+ params[name].data.copy_(tensor)
28
+ loaded.add(name)
29
+ return loaded
30
+
31
+
32
+class SimpleModel(torch.nn.Module):
33
+ def __init__(self):
34
+ super().__init__()
35
+ self.model = SimpleInnerModel()
36
+ self.lm_head = torch.nn.Linear(8, 16, bias=False)
37
+
38
+ def load_weights(self, weights):
39
+ loader = AutoWeightsLoader(self)
40
+ return loader.load_weights(weights)
41
+
42
+
43
+class PackedWeightInnerModel(torch.nn.Module):
44
+ """Remaps q_proj/k_proj into a fused qkv_proj (Qwen2/Llama pattern)."""
45
+
46
+ def __init__(self):
47
+ super().__init__()
48
+ self.qkv_proj = torch.nn.Linear(4, 16, bias=False)
49
+ self.out = torch.nn.Linear(8, 4, bias=False)
50
+
51
+ def load_weights(self, weights):
52
+ params = dict(self.named_parameters())
53
+ loaded = set()
54
+ for name, tensor in weights:
55
+ if name == "q_proj.weight":
56
+ params["qkv_proj.weight"].data[:8].copy_(tensor)
57
+ loaded.add("qkv_proj.weight")
58
+ elif name == "k_proj.weight":
59
+ params["qkv_proj.weight"].data[8:].copy_(tensor)
60
+ loaded.add("qkv_proj.weight")
61
+ elif name in params:
62
+ params[name].data.copy_(tensor)
63
+ loaded.add(name)
64
+ return loaded
65
+
66
+
67
+class PackedWeightModel(torch.nn.Module):
68
+ def __init__(self):
69
+ super().__init__()
70
+ self.model = PackedWeightInnerModel()
71
+ self.lm_head = torch.nn.Linear(4, 8, bias=False)
72
+
73
+ def load_weights(self, weights):
74
+ loader = AutoWeightsLoader(self)
75
+ return loader.load_weights(weights)
76
+
77
+
78
+def _buffer_reusing_iterator(weight_dict):
79
+ """Yield weights through a shared buffer overwritten each step.
80
+
81
+ Mimics ``runai_model_streamer`` with ``RUNAI_STREAMER_MEMORY_LIMIT=0``.
82
+ """
83
+ buf = None
84
+ for name, tensor in weight_dict.items():
85
+ if buf is None or buf.numel() < tensor.numel():
86
+ buf = torch.empty(tensor.numel(), dtype=tensor.dtype)
87
+ view = buf[: tensor.numel()].view(tensor.shape)
88
+ view.copy_(tensor)
89
+ yield name, view
90
+
91
+
92
+def _make_pooling_model(base_cls=SimpleModel):
93
+ PoolingModel = _create_pooling_model_cls(base_cls)
94
+ model = base_cls()
95
+ model.__class__ = PoolingModel
96
+ model.lm_head = StageMissingLayer("output", model.lm_head)
97
+ return model
98
+
99
+
100
+def _make_reference_weights():
101
+ torch.manual_seed(42)
102
+ return {
103
+ "model.embed.weight": torch.randn(8, 4),
104
+ "model.layer0.weight": torch.randn(8, 8),
105
+ "model.layer1.weight": torch.randn(8, 8),
106
+ "model.norm.weight": torch.randn(4, 8),
107
+ "lm_head.weight": torch.randn(16, 8),
108
+ }
109
+
110
+
111
+def _make_packed_reference_weights():
112
+ torch.manual_seed(42)
113
+ return {
114
+ "model.q_proj.weight": torch.randn(8, 4),
115
+ "model.k_proj.weight": torch.randn(8, 4),
116
+ "model.out.weight": torch.randn(4, 8),
117
+ "lm_head.weight": torch.randn(8, 4),
118
+ }
119
+
120
+
121
+def _load_and_compare(model, ref, expected):
122
+ for p in model.parameters():
123
+ p.data.zero_()
124
+ model.load_weights(_buffer_reusing_iterator(ref))
125
+ for name, param in model.named_parameters():
126
+ assert torch.equal(param.data, expected[name]), name
127
+
128
+
129
+def test_pooling_load_weights_with_buffer_reuse():
130
+ """Ensure ModelForPooling.load_weights works with buffer-reusing iterators."""
131
+ ref = _make_reference_weights()
132
+
133
+ ground_truth = SimpleModel()
134
+ ground_truth.load_weights(ref.items())
135
+ expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()}
136
+
137
+ _load_and_compare(_make_pooling_model(), ref, expected)
138
+
139
+
140
+def test_pooling_load_weights_clones_probed_weights():
141
+ """Ensure probed weights survive buffer reuse during packed remapping."""
142
+ ref = _make_packed_reference_weights()
143
+
144
+ ground_truth = PackedWeightModel()
145
+ ground_truth.load_weights(ref.items())
146
+ expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()}
147
+
148
+ _load_and_compare(_make_pooling_model(PackedWeightModel), ref, expected)
vllm/model_executor/models/adapters.py
+6
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import itertools
5
from collections.abc import Iterable
6
from contextlib import contextmanager
7
from typing import TYPE_CHECKING, Any, TypeVar, cast
@@ -181,7 +182,8 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T:
182
183
seen_weights = list[tuple[str, torch.Tensor]]()
184
for name, loaded_weight in weights:
184
- seen_weights.append((name, loaded_weight))
185
+ # Clone because the iterator may reuse the tensor buffer
186
+ seen_weights.append((name, loaded_weight.clone()))
187
188
try:
189
target_prefix = next(
@@ -208,9 +210,11 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T:
210
self._get_name(),
211
)
212
213
+ # Lazy chain so buffer-reusing weight iterators (e.g.
214
+ # runai_streamer) are consumed one tensor at a time.
215
mapped_weights = (
216
(target_prefix + name, weight)
213
- for name, weight in (*seen_weights, *weights)
217
+ for name, weight in itertools.chain(seen_weights, weights)
218
)
219
220
def default_load_weights(weights):