| 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 | from transformers import Gemma3Config, PreTrainedConfig, Qwen2Config |
| 8 | |
| 9 | 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 ( |
| 14 | AutoWeightsLoader, |
| 15 | StageMissingLayer, |
| 16 | WeightsMapper, |
| 17 | ) |
| 18 | |
| 19 | pytestmark = pytest.mark.cpu_test |
| 20 | |
| 21 | |
| 22 | class SimpleInnerModel(torch.nn.Module): |
| 23 | def __init__(self): |
| 24 | super().__init__() |
| 25 | self.embed = torch.nn.Linear(4, 8, bias=False) |
| 26 | self.layer0 = torch.nn.Linear(8, 8, bias=False) |
| 27 | self.layer1 = torch.nn.Linear(8, 8, bias=False) |
| 28 | self.norm = torch.nn.Linear(8, 4, bias=False) |
| 29 | |
| 30 | def load_weights(self, weights): |
| 31 | params = dict(self.named_parameters()) |
| 32 | loaded = set() |
| 33 | for name, tensor in weights: |
| 34 | if name in params: |
| 35 | params[name].data.copy_(tensor) |
| 36 | loaded.add(name) |
| 37 | return loaded |
| 38 | |
| 39 | |
| 40 | class SimpleModel(torch.nn.Module): |
| 41 | def __init__(self): |
| 42 | super().__init__() |
| 43 | self.model = SimpleInnerModel() |
| 44 | self.lm_head = torch.nn.Linear(8, 16, bias=False) |
| 45 | |
| 46 | def load_weights(self, weights): |
| 47 | loader = AutoWeightsLoader(self) |
| 48 | return loader.load_weights(weights) |
| 49 | |
| 50 | |
| 51 | class PackedWeightInnerModel(torch.nn.Module): |
| 52 | """Remaps q_proj/k_proj into a fused qkv_proj (Qwen2/Llama pattern).""" |
| 53 | |
| 54 | def __init__(self): |
| 55 | super().__init__() |
| 56 | self.qkv_proj = torch.nn.Linear(4, 16, bias=False) |
| 57 | self.out = torch.nn.Linear(8, 4, bias=False) |
| 58 | |
| 59 | def load_weights(self, weights): |
| 60 | params = dict(self.named_parameters()) |
| 61 | loaded = set() |
| 62 | for name, tensor in weights: |
| 63 | if name == "q_proj.weight": |
| 64 | params["qkv_proj.weight"].data[:8].copy_(tensor) |
| 65 | loaded.add("qkv_proj.weight") |
| 66 | elif name == "k_proj.weight": |
| 67 | params["qkv_proj.weight"].data[8:].copy_(tensor) |
| 68 | loaded.add("qkv_proj.weight") |
| 69 | elif name in params: |
| 70 | params[name].data.copy_(tensor) |
| 71 | loaded.add(name) |
| 72 | return loaded |
| 73 | |
| 74 | |
| 75 | class PackedWeightModel(torch.nn.Module): |
| 76 | def __init__(self): |
| 77 | super().__init__() |
| 78 | self.model = PackedWeightInnerModel() |
| 79 | self.lm_head = torch.nn.Linear(4, 8, bias=False) |
| 80 | |
| 81 | def load_weights(self, weights): |
| 82 | loader = AutoWeightsLoader(self) |
| 83 | return loader.load_weights(weights) |
| 84 | |
| 85 | |
| 86 | def _buffer_reusing_iterator(weight_dict): |
| 87 | """Yield weights through a shared buffer overwritten each step. |
| 88 | |
| 89 | Mimics ``runai_model_streamer`` with ``RUNAI_STREAMER_MEMORY_LIMIT=0``. |
| 90 | """ |
| 91 | buf = None |
| 92 | for name, tensor in weight_dict.items(): |
| 93 | if buf is None or buf.numel() < tensor.numel(): |
| 94 | buf = torch.empty(tensor.numel(), dtype=tensor.dtype) |
| 95 | view = buf[: tensor.numel()].view(tensor.shape) |
| 96 | view.copy_(tensor) |
| 97 | yield name, view |
| 98 | |
| 99 | |
| 100 | def _make_pooling_model(base_cls=SimpleModel): |
| 101 | PoolingModel = _create_pooling_model_cls(base_cls) |
| 102 | model = base_cls() |
| 103 | model.__class__ = PoolingModel |
| 104 | model.lm_head = StageMissingLayer("output", model.lm_head) |
| 105 | return model |
| 106 | |
| 107 | |
| 108 | def _make_reference_weights(): |
| 109 | torch.manual_seed(42) |
| 110 | return { |
| 111 | "model.embed.weight": torch.randn(8, 4), |
| 112 | "model.layer0.weight": torch.randn(8, 8), |
| 113 | "model.layer1.weight": torch.randn(8, 8), |
| 114 | "model.norm.weight": torch.randn(4, 8), |
| 115 | "lm_head.weight": torch.randn(16, 8), |
| 116 | } |
| 117 | |
| 118 | |
| 119 | def _make_packed_reference_weights(): |
| 120 | torch.manual_seed(42) |
| 121 | return { |
| 122 | "model.q_proj.weight": torch.randn(8, 4), |
| 123 | "model.k_proj.weight": torch.randn(8, 4), |
| 124 | "model.out.weight": torch.randn(4, 8), |
| 125 | "lm_head.weight": torch.randn(8, 4), |
| 126 | } |
| 127 | |
| 128 | |
| 129 | def _load_and_compare(model, ref, expected): |
| 130 | for p in model.parameters(): |
| 131 | p.data.zero_() |
| 132 | model.load_weights(_buffer_reusing_iterator(ref)) |
| 133 | for name, param in model.named_parameters(): |
| 134 | assert torch.equal(param.data, expected[name]), name |
| 135 | |
| 136 | |
| 137 | def test_pooling_load_weights_with_buffer_reuse(): |
| 138 | """Ensure ModelForPooling.load_weights works with buffer-reusing iterators.""" |
| 139 | ref = _make_reference_weights() |
| 140 | |
| 141 | ground_truth = SimpleModel() |
| 142 | ground_truth.load_weights(ref.items()) |
| 143 | expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()} |
| 144 | |
| 145 | _load_and_compare(_make_pooling_model(), ref, expected) |
| 146 | |
| 147 | |
| 148 | def test_pooling_load_weights_clones_probed_weights(): |
| 149 | """Ensure probed weights survive buffer reuse during packed remapping.""" |
| 150 | ref = _make_packed_reference_weights() |
| 151 | |
| 152 | ground_truth = PackedWeightModel() |
| 153 | ground_truth.load_weights(ref.items()) |
| 154 | expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()} |
| 155 | |
| 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}) |
| 252 | if outer_labels is not None: |
| 253 | config.num_labels = outer_labels |
| 254 | if inner_labels is not None: |
| 255 | config.get_text_config().num_labels = inner_labels |
| 256 | return config |
| 257 | |
| 258 | |
| 259 | def test_resolve_num_labels_text_only_config(): |
| 260 | config = Qwen2Config(num_labels=7) |
| 261 | assert config.get_text_config() is config |
| 262 | assert _resolve_num_labels(config, config.get_text_config()) == 7 |
| 263 | |
| 264 | |
| 265 | def test_resolve_num_labels_defaults_when_undeclared(): |
| 266 | config = _composite_config() |
| 267 | assert ( |
| 268 | _resolve_num_labels(config, config.get_text_config()) |
| 269 | == PreTrainedConfig().num_labels |
| 270 | ) |
| 271 | |
| 272 | |
| 273 | def test_resolve_num_labels_declared_on_outer_config(): |
| 274 | """Multimodal checkpoints keep id2label/problem_type on the top-level config.""" |
| 275 | config = _composite_config(outer_labels=20) |
| 276 | assert config.get_text_config().num_labels == PreTrainedConfig().num_labels |
| 277 | assert _resolve_num_labels(config, config.get_text_config()) == 20 |
| 278 | |
| 279 | |
| 280 | def test_resolve_num_labels_declared_on_text_config(): |
| 281 | """Overrides written into text_config keep working.""" |
| 282 | config = _composite_config(inner_labels=5) |
| 283 | assert _resolve_num_labels(config, config.get_text_config()) == 5 |
| 284 | |
| 285 | |
| 286 | def test_resolve_num_labels_outer_wins_when_both_declared(): |
| 287 | config = _composite_config(outer_labels=20, inner_labels=5) |
| 288 | assert _resolve_num_labels(config, config.get_text_config()) == 20 |