[BugFix][Pooling] Skip weight-prefix probe when model has WeightsMapper (#50890)
Signed-off-by: Juqi Li <2223621784@qq.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Li Juqi committed
Aug 6, 2026 at 12:12 UTC
821717118fc26667dd474b9b0ab81d29259dfc5c
2 files changed
+108
-2
tests/models/test_adapters.py
+95
-1
@@ -10,7 +10,11 @@ from vllm.model_executor.models.adapters import (
10
_create_pooling_model_cls,
11
_resolve_num_labels,
12
)
13
-from vllm.model_executor.models.utils import AutoWeightsLoader, StageMissingLayer
13
+from vllm.model_executor.models.utils import (
14
+ AutoWeightsLoader,
15
+ StageMissingLayer,
16
+ WeightsMapper,
17
+)
18
19
pytestmark = pytest.mark.cpu_test
20
@@ -152,6 +156,96 @@ def test_pooling_load_weights_clones_probed_weights():
156
_load_and_compare(_make_pooling_model(PackedWeightModel), ref, expected)
157
158
159
+class _LanguageModelInner(torch.nn.Module):
160
+ def __init__(self):
161
+ super().__init__()
162
+ self.embed = torch.nn.Linear(4, 8, bias=False)
163
+ self.layer0 = torch.nn.Linear(8, 8, bias=False)
164
+
165
+
166
+class _LanguageModelWrapper(torch.nn.Module):
167
+ def __init__(self):
168
+ super().__init__()
169
+ self.model = _LanguageModelInner()
170
+
171
+
172
+class ModelWithWeightsMapper(torch.nn.Module):
173
+ """Stand-in for models whose keys only align after hf_to_vllm_mapper.
174
+
175
+ Checkpoint keys like ``model.language_model.*`` never match
176
+ ``""`` / ``model.`` + raw name against ``language_model.model.*`` params.
177
+ The pooling prefix probe must consult the mapper so it can early-exit.
178
+ """
179
+
180
+ hf_to_vllm_mapper = WeightsMapper(
181
+ orig_to_new_prefix={
182
+ "model.language_model.": "language_model.model.",
183
+ }
184
+ )
185
+
186
+ def __init__(self):
187
+ super().__init__()
188
+ self.language_model = _LanguageModelWrapper()
189
+ self.lm_head = torch.nn.Linear(8, 16, bias=False)
190
+
191
+ def load_weights(self, weights):
192
+ params = dict(self.named_parameters())
193
+ loaded = set()
194
+ self.seen_names = []
195
+ for name, tensor in weights:
196
+ self.seen_names.append(name)
197
+ mapped = self.hf_to_vllm_mapper._map_name(name)
198
+ if mapped is not None and mapped in params:
199
+ params[mapped].data.copy_(tensor)
200
+ loaded.add(mapped)
201
+ return loaded
202
+
203
+
204
+def test_pooling_prefix_probe_uses_hf_to_vllm_mapper(monkeypatch):
205
+ """Mapper lets the prefix probe early-exit without cloning the full ckpt.
206
+
207
+ Checkpoint keys only align after ``hf_to_vllm_mapper``. The probe must use
208
+ the mapped name for membership, keep forwarding original names to the
209
+ parent loader, and stop after the first hit instead of cloning everything.
210
+ """
211
+ ref = {
212
+ "model.language_model.embed.weight": torch.randn(8, 4),
213
+ "model.language_model.layer0.weight": torch.randn(8, 8),
214
+ "model.language_model.extra0.weight": torch.randn(8, 8),
215
+ "model.language_model.extra1.weight": torch.randn(8, 8),
216
+ "model.language_model.extra2.weight": torch.randn(8, 8),
217
+ }
218
+
219
+ # Sanity: none of these keys match the generic "" / "model." probe.
220
+ model = _make_pooling_model(ModelWithWeightsMapper)
221
+ params = dict(model.named_parameters())
222
+ for name in ref:
223
+ assert name not in params
224
+ assert f"model.{name}" not in params
225
+
226
+ clone_count = 0
227
+ original_clone = torch.Tensor.clone
228
+
229
+ def counting_clone(self, *args, **kwargs):
230
+ nonlocal clone_count
231
+ clone_count += 1
232
+ return original_clone(self, *args, **kwargs)
233
+
234
+ monkeypatch.setattr(torch.Tensor, "clone", counting_clone)
235
+
236
+ loaded = model.load_weights(iter(ref.items()))
237
+
238
+ # First key maps onto a param, so the probe clones once and breaks.
239
+ assert clone_count == 1
240
+ assert clone_count < len(ref)
241
+ # Parent still receives original (unmapped, unprefixed) checkpoint names.
242
+ assert model.seen_names == list(ref.keys())
243
+ assert loaded == {
244
+ "language_model.model.embed.weight",
245
+ "language_model.model.layer0.weight",
246
+ }
247
+
248
+
249
def _composite_config(outer_labels=None, inner_labels=None):
250
"""Build a multimodal-style config whose text_config is a separate object."""
251
config = Gemma3Config(text_config={"num_hidden_layers": 1})
vllm/model_executor/models/adapters.py
+13
-1
@@ -175,6 +175,11 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T:
175
176
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
177
params_dict = dict(self.named_parameters())
178
+ # Use model-owned remapping only for the probe membership check
179
+ # (e.g. Gemma4 model.language_model.* -> language_model.model.*)
180
+ # so the "" / "model." scan can early-exit. Forwarded names stay
181
+ # unmapped to avoid double-applying the mapper in the parent loader.
182
+ hf_to_vllm_mapper = getattr(self, "hf_to_vllm_mapper", None)
183
184
# We support loading from both `*ForCausalLM` and `*Model`
185
candidate_prefixes = ["", "model."]
@@ -185,11 +190,18 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T:
190
# Clone because the iterator may reuse the tensor buffer
191
seen_weights.append((name, loaded_weight.clone()))
192
193
+ lookup_name = name
194
+ if hf_to_vllm_mapper is not None:
195
+ mapped = hf_to_vllm_mapper._map_name(name)
196
+ if mapped is None:
197
+ continue
198
+ lookup_name = mapped
199
+
200
try:
201
target_prefix = next(
202
prefix
203
for prefix in candidate_prefixes
192
- if prefix + name in params_dict
204
+ if prefix + lookup_name in params_dict
205
)
206
break
207
except StopIteration: