Fully generalise input embedding handling in Transformers modelling backend (#51247)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Aug 6, 2026 at 17:21 UTC
e7b8d5946095a594af2cf7ca3c314b9806cb7c32
6 files changed
+300
-47
tests/models/transformers/test_backend.py
+176
@@ -2,9 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Test the functionality of the Transformers modeling backend."""
4
5
+import contextlib
6
+import os
7
+import tempfile
8
from typing import Any
9
10
import pytest
11
+import torch
12
+import torch.nn as nn
13
14
from ...conftest import HfRunner, VllmRunner
15
from ...utils import multi_gpu_test, prep_prompts
@@ -299,3 +304,174 @@ def test_pooling(hf_runner, vllm_runner, example_prompts, arch):
304
name_0="hf",
305
name_1="vllm",
306
)
307
+
308
+
309
+VOCAB_SIZE = 64
310
+HIDDEN_SIZE = 8
311
+EMBED_SCALE = 3.0
312
+
313
+
314
+class ScaledWordEmbedding(nn.Embedding):
315
+ """Mirrors Transformers' `*ScaledWordEmbedding` classes."""
316
+
317
+ def __init__(
318
+ self, num_embeddings, embedding_dim, padding_idx=None, embed_scale=1.0
319
+ ):
320
+ super().__init__(num_embeddings, embedding_dim, padding_idx)
321
+ self.scalar_embed_scale = embed_scale
322
+ self.register_buffer("embed_scale", torch.tensor(embed_scale), persistent=False)
323
+
324
+ def forward(self, input_ids):
325
+ return super().forward(input_ids) * self.embed_scale.to(self.weight.dtype)
326
+
327
+
328
+class ComposedWordEmbedding(nn.Module):
329
+ """Scales embeddings, but wraps `nn.Embedding` instead of inheriting from it."""
330
+
331
+ def __init__(self, num_embeddings, embedding_dim, embed_scale=1.0):
332
+ super().__init__()
333
+ self.embed = nn.Embedding(num_embeddings, embedding_dim)
334
+ self.embed_scale = embed_scale
335
+
336
+ def forward(self, input_ids):
337
+ return self.embed(input_ids) * self.embed_scale
338
+
339
+
340
+@pytest.fixture
341
+def tp_init():
342
+ """Single rank tensor parallel state, so vLLM layers can be constructed."""
343
+ from vllm.distributed import (
344
+ cleanup_dist_env_and_memory,
345
+ init_distributed_environment,
346
+ initialize_model_parallel,
347
+ )
348
+ from vllm.platforms import current_platform
349
+
350
+ from ...utils import ensure_current_vllm_config
351
+
352
+ fd, temp_file = tempfile.mkstemp()
353
+ os.close(fd)
354
+ try:
355
+ with ensure_current_vllm_config():
356
+ init_distributed_environment(
357
+ world_size=1,
358
+ rank=0,
359
+ distributed_init_method=f"file://{temp_file}",
360
+ local_rank=0,
361
+ backend=current_platform.dist_backend,
362
+ )
363
+ initialize_model_parallel(1, 1)
364
+ yield
365
+ cleanup_dist_env_and_memory()
366
+ finally:
367
+ with contextlib.suppress(OSError):
368
+ os.unlink(temp_file)
369
+
370
+
371
+@pytest.fixture
372
+def vpe(tp_init):
373
+ """`VocabParallelEmbedding`, imported late so collection does not import vLLM."""
374
+ from vllm.model_executor.layers.vocab_parallel_embedding import (
375
+ VocabParallelEmbedding,
376
+ )
377
+
378
+ return VocabParallelEmbedding
379
+
380
+
381
+def replace(embedding):
382
+ """Replace `embedding` and fill the new weights with recognisable values."""
383
+ from vllm.model_executor.models.transformers.utils import replace_embedding_class
384
+
385
+ new_embedding = replace_embedding_class(embedding)
386
+ for _, param in new_embedding.named_parameters():
387
+ param.data = torch.arange(param.numel(), dtype=param.dtype).view(param.shape)
388
+ return new_embedding
389
+
390
+
391
+def assert_scaled(vpe, module, embedding=None):
392
+ """`module`'s output is `embedding`'s unscaled output times `EMBED_SCALE`."""
393
+ input_ids = torch.arange(VOCAB_SIZE)
394
+ unscaled = vpe.forward(embedding if embedding is not None else module, input_ids)
395
+ torch.testing.assert_close(module(input_ids), unscaled * EMBED_SCALE)
396
+
397
+
398
+def test_replace_plain_embedding(vpe):
399
+ """A plain `nn.Embedding` is replaced outright, leaving no subclass behind."""
400
+ assert type(replace(nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE))) is vpe
401
+
402
+
403
+def test_replace_infers_shape_and_dtype(tp_init):
404
+ """Shape and dtype come from the replaced module, not from the config."""
405
+ embedding = nn.Embedding(VOCAB_SIZE * 2, HIDDEN_SIZE + 1, dtype=torch.float16)
406
+ new_embedding = replace(embedding)
407
+
408
+ assert new_embedding.num_embeddings == VOCAB_SIZE * 2
409
+ assert new_embedding.org_vocab_size == VOCAB_SIZE * 2
410
+ assert new_embedding.embedding_dim == HIDDEN_SIZE + 1
411
+ assert new_embedding.weight.dtype == torch.float16
412
+
413
+
414
+def test_replace_inherited_embedding(vpe):
415
+ """Subclasses keep their extra state and their scaled `forward`."""
416
+ new_embedding = replace(
417
+ ScaledWordEmbedding(
418
+ VOCAB_SIZE, HIDDEN_SIZE, padding_idx=0, embed_scale=EMBED_SCALE
419
+ )
420
+ )
421
+
422
+ assert isinstance(new_embedding, vpe)
423
+ assert new_embedding.scalar_embed_scale == EMBED_SCALE
424
+ assert "embed_scale" in new_embedding._non_persistent_buffers_set
425
+ assert_scaled(vpe, new_embedding)
426
+
427
+
428
+def test_replace_composed_embedding(vpe):
429
+ """Wrappers are left alone; only the `nn.Embedding` they hold is replaced."""
430
+ embedding = ComposedWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE)
431
+ new_embedding = replace(embedding)
432
+
433
+ assert new_embedding is embedding
434
+ assert type(embedding.embed) is vpe
435
+ assert_scaled(vpe, new_embedding, embedding.embed)
436
+
437
+
438
+def test_replace_nested_embedding(vpe):
439
+ """The composed `nn.Embedding` is found and set however deeply it is nested."""
440
+ wrapper = nn.Module()
441
+ wrapper.add_module("inner", nn.Module())
442
+ wrapper.inner.add_module(
443
+ "embed", ScaledWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE)
444
+ )
445
+ replace(wrapper)
446
+
447
+ assert isinstance(wrapper.inner.embed, vpe)
448
+ assert_scaled(vpe, wrapper.inner.embed)
449
+
450
+
451
+@pytest.mark.parametrize("num_embeddings", [0, 2])
452
+def test_replace_ambiguous_embedding(tp_init, num_embeddings):
453
+ """Composing anything but one `nn.Embedding` is an error, not a silent guess."""
454
+ wrapper = nn.Module()
455
+ for i in range(num_embeddings):
456
+ wrapper.add_module(f"embed_{i}", nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE))
457
+
458
+ with pytest.raises(ValueError, match=f"found {num_embeddings}"):
459
+ replace(wrapper)
460
+
461
+
462
+def test_replaced_embedding_exposes_one_vpe(vpe):
463
+ """`CausalMixin` ties `lm_head` to the one `VocabParallelEmbedding` it can find.
464
+
465
+ `tie_weights` reads `.weight`, which a composing module does not have, so it must
466
+ be handed the composed embedding instead.
467
+ """
468
+ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
469
+
470
+ inherited = replace(ScaledWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE))
471
+ composed = replace(ComposedWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, EMBED_SCALE))
472
+
473
+ assert [m for m in inherited.modules() if isinstance(m, vpe)] == [inherited]
474
+ assert [m for m in composed.modules() if isinstance(m, vpe)] == [composed.embed]
475
+
476
+ lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE)
477
+ assert lm_head.tie_weights(composed.embed).weight is composed.embed.weight
vllm/lora/layers/vocal_parallel_embedding.py
+15
-2
@@ -8,7 +8,10 @@ from transformers import PretrainedConfig
8
9
from vllm.config.lora import LoRAConfig
10
from vllm.model_executor.custom_op import maybe_get_oot_by_class
11
-from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
11
+from vllm.model_executor.layers.vocab_parallel_embedding import (
12
+ ParallelLMHead,
13
+ VocabParallelEmbedding,
14
+)
15
from vllm.platforms import current_platform
16
17
from .base import BaseLayerWithLoRA
@@ -133,7 +136,17 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
136
packed_modules_list: list,
137
model_config: PretrainedConfig | None = None,
138
) -> bool:
136
- return type(source_layer) is maybe_get_oot_by_class(VocabParallelEmbedding)
139
+ # Accept both the in-tree and any OOT class, because layers built on top of
140
+ # `VocabParallelEmbedding` (e.g. by the Transformers modelling backend) use the
141
+ # former while `PluggableLayer.__new__` instantiates the latter
142
+ embedding_classes = (
143
+ VocabParallelEmbedding,
144
+ maybe_get_oot_by_class(VocabParallelEmbedding),
145
+ )
146
+ lm_head_classes = (ParallelLMHead, maybe_get_oot_by_class(ParallelLMHead))
147
+ return isinstance(source_layer, embedding_classes) and not isinstance(
148
+ source_layer, lm_head_classes
149
+ )
150
151
@property
152
def weight(self):
vllm/model_executor/layers/vocab_parallel_embedding.py
+1
-1
@@ -290,7 +290,7 @@ class VocabParallelEmbedding(PluggableLayer):
290
# If we are making an embedding layer, then our quantization linear
291
# method must implement the embedding operation. If we are another
292
# layer type like ParallelLMHead, this is not important.
293
- is_embedding_layer = type(self) is VocabParallelEmbedding
293
+ is_embedding_layer = not isinstance(self, ParallelLMHead)
294
quant_method_implements_embedding = method_has_implemented_embedding(
295
type(quant_method)
296
)
vllm/model_executor/models/transformers/base.py
+4
-39
@@ -46,7 +46,6 @@ from vllm.model_executor.layers.attention import (
46
from vllm.model_executor.layers.attention.mla_attention import get_mla_dims
47
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
48
from vllm.model_executor.layers.fused_moe import MoERunner
49
-from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
49
from vllm.model_executor.models.interfaces import (
50
SupportsEagle,
51
SupportsEagle3,
@@ -58,12 +57,14 @@ from vllm.model_executor.models.interfaces_base import VllmModel
57
from vllm.model_executor.models.transformers.fuser import BaseFuser, Fusers
58
from vllm.model_executor.models.transformers.fusers import MLAFuser
59
from vllm.model_executor.models.transformers.utils import (
60
+ attrsetter,
61
can_enable_torch_compile,
62
get_feature_request_tip,
63
init_on_device_without_buffers,
64
log_replacement,
65
named_state,
66
replace_conv_class,
67
+ replace_embedding_class,
68
replace_linear_class,
69
)
70
from vllm.model_executor.models.utils import (
@@ -84,17 +85,6 @@ if TYPE_CHECKING:
85
logger = init_logger(__name__)
86
87
87
-class ScaledVocabParallelEmbedding(VocabParallelEmbedding):
88
- """`VocabParallelEmbedding` that scales its output."""
89
-
90
- def __init__(self, *args, embed_scale: float, **kwargs):
91
- super().__init__(*args, **kwargs)
92
- self.embed_scale = embed_scale
93
-
94
- def forward(self, input_: torch.Tensor) -> torch.Tensor:
95
- return super().forward(input_) * self.embed_scale
96
-
97
-
88
class Base(
89
nn.Module,
90
VllmModel,
@@ -179,24 +169,9 @@ class Base(
169
# Input embeddings
170
input_embeddings = self.model.get_input_embeddings()
171
if not isinstance(input_embeddings, PPMissingLayer):
182
- names = ("embedding_size", "hidden_size")
183
- embedding_dim = getattr_iter(self.text_config, names, None)
184
- assert embedding_dim is not None
185
- embedding_kwargs = dict(
186
- num_embeddings=self.text_config.vocab_size,
187
- embedding_dim=embedding_dim,
188
- org_num_embeddings=self.text_config.vocab_size,
189
- quant_config=self.quant_config,
172
+ self.model.set_input_embeddings(
173
+ replace_embedding_class(input_embeddings, self.quant_config)
174
)
191
- embed_scale = getattr(input_embeddings, "embed_scale", None)
192
- if embed_scale is not None:
193
- # Some models scale embeddings inside the input embedding layer
194
- new_input_embeddings = ScaledVocabParallelEmbedding(
195
- **embedding_kwargs, embed_scale=float(embed_scale)
196
- )
197
- else:
198
- new_input_embeddings = VocabParallelEmbedding(**embedding_kwargs)
199
- self.model.set_input_embeddings(new_input_embeddings)
175
176
# Initialize any parameters that have not had their modules replaced
177
self.init_parameters(self.model)
@@ -380,16 +355,6 @@ class Base(
355
tip,
356
)
357
383
- def attrsetter(attr: str) -> Callable[[object, object], None]:
384
- """Set a possibly nested attribute, like the inverse of attrgetter."""
385
- parent, _, name = attr.rpartition(".")
386
-
387
- def setter(obj: object, value: object):
388
- attr_parent = attrgetter(parent)(obj) if parent else obj
389
- setattr(attr_parent, name, value)
390
-
391
- return setter
392
-
358
module_lists = []
359
module_list_idx = None
360
for i, name in enumerate(names):
vllm/model_executor/models/transformers/causal.py
+8
-4
@@ -20,7 +20,10 @@ from collections.abc import Iterable
20
from typing import TYPE_CHECKING
21
22
from vllm.model_executor.layers.logits_processor import LogitsProcessor
23
-from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
23
+from vllm.model_executor.layers.vocab_parallel_embedding import (
24
+ ParallelLMHead,
25
+ VocabParallelEmbedding,
26
+)
27
from vllm.model_executor.models.interfaces_base import VllmModelForTextGeneration
28
from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix
29
@@ -51,9 +54,10 @@ class CausalMixin(VllmModelForTextGeneration):
54
prefix=maybe_prefix(prefix, "lm_head"),
55
)
56
if tie_word_embeddings:
54
- self.lm_head = self.lm_head.tie_weights(
55
- self.model.get_input_embeddings()
56
- )
57
+ for module in self.model.get_input_embeddings().modules():
58
+ if isinstance(module, VocabParallelEmbedding):
59
+ self.lm_head = self.lm_head.tie_weights(module)
60
+ break
61
62
logit_scale = getattr(self.text_config, "logit_scale", 1.0)
63
self.logits_processor = LogitsProcessor(
vllm/model_executor/models/transformers/utils.py
+96
-1
@@ -16,9 +16,11 @@
16
# limitations under the License.
17
"""Transformers modeling backend utilities."""
18
19
-from collections.abc import Iterator
19
+from collections.abc import Callable, Iterator
20
from contextlib import contextmanager
21
+from functools import lru_cache
22
from itertools import chain
23
+from operator import attrgetter
24
from pathlib import Path
25
from typing import TYPE_CHECKING, Literal
26
@@ -32,6 +34,7 @@ from vllm.model_executor.layers.linear import (
34
ReplicatedLinear,
35
RowParallelLinear,
36
)
37
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
38
from vllm.model_executor.models.utils import maybe_prefix
39
from vllm.transformers_utils.config import is_rope_parameters_nested
40
@@ -183,6 +186,98 @@ def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv:
186
)
187
188
189
+def attrsetter(attr: str) -> Callable[[object, object], None]:
190
+ """Set a possibly nested attribute, like the inverse of attrgetter."""
191
+ parent, _, name = attr.rpartition(".")
192
+
193
+ def setter(obj: object, value: object):
194
+ attr_parent = attrgetter(parent)(obj) if parent else obj
195
+ setattr(attr_parent, name, value)
196
+
197
+ return setter
198
+
199
+
200
+class _UninitializedEmbedding(nn.Embedding):
201
+ """Make `__init__` inert, so that `VocabParallelEmbedding.__init__`'
202
+ call to `super().__init__` does not invoke `nn.Embedding.__init__`."""
203
+
204
+ def __init__(self):
205
+ pass
206
+
207
+
208
+class _VocabParallelEmbeddingBase(VocabParallelEmbedding, _UninitializedEmbedding):
209
+ """Orders `VocabParallelEmbedding` ahead of `nn.Embedding` in the MRO, so that
210
+ `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding."""
211
+
212
+
213
+@lru_cache
214
+def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbedding]:
215
+ """Subclass `cls` so that `VocabParallelEmbedding` supersedes its `nn.Embedding`.
216
+
217
+ Args:
218
+ cls: The `nn.Embedding` subclass to rebase. Cached, so a given `cls` always
219
+ maps to the same class.
220
+ Returns:
221
+ The new class, to assign to `__class__` of an instance of `cls`.
222
+ """
223
+ return type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {})
224
+
225
+
226
+def replace_embedding_class(
227
+ embedding: nn.Module,
228
+ quant_config: "QuantizationConfig | None" = None,
229
+ *,
230
+ prefix: str = "",
231
+) -> nn.Module:
232
+ """Replace the `nn.Embedding` in `embedding` with `VocabParallelEmbedding`.
233
+
234
+ Args:
235
+ embedding: The module returned by `model.get_input_embeddings()`.
236
+ quant_config: Quantization config for the new embedding.
237
+ prefix: Qualname of `embedding`, used to look up its quantization method.
238
+ Returns:
239
+ The module to install with `model.set_input_embeddings()`. Composing and
240
+ inheriting modules are mutated in place and returned as-is.
241
+ Raises:
242
+ ValueError: If `embedding` composes anything other than one `nn.Embedding`,
243
+ which would leave the input embedding weights ambiguous.
244
+ """
245
+ # If `embedding` composes its `nn.Embedding`, recurse into it
246
+ if not isinstance(embedding, nn.Embedding):
247
+ composed = [
248
+ (name, module)
249
+ for name, module in embedding.named_modules()
250
+ if isinstance(module, nn.Embedding)
251
+ ]
252
+ if len(composed) != 1:
253
+ raise ValueError(
254
+ f"Expected {type(embedding).__name__} to be an `nn.Embedding` or to "
255
+ f"compose exactly one, but found {len(composed)}."
256
+ )
257
+ name, module = composed[0]
258
+ new_embedding = replace_embedding_class(
259
+ module, quant_config, prefix=maybe_prefix(prefix, name)
260
+ )
261
+ attrsetter(name)(embedding, new_embedding)
262
+ return embedding
263
+
264
+ kwargs = dict(
265
+ num_embeddings=embedding.num_embeddings,
266
+ embedding_dim=embedding.embedding_dim,
267
+ params_dtype=embedding.weight.dtype,
268
+ quant_config=quant_config,
269
+ prefix=prefix,
270
+ )
271
+ # If `embedding` is a bare `nn.Embedding`, simple replace
272
+ if type(embedding) is nn.Embedding:
273
+ return VocabParallelEmbedding(**kwargs)
274
+
275
+ # Otherwise `embedding` inherits `nn.Embedding`, rebase it in place
276
+ embedding.__class__ = _rebase_on_vocab_parallel(type(embedding))
277
+ VocabParallelEmbedding.__init__(embedding, **kwargs)
278
+ return embedding
279
+
280
+
281
def recursive_replace_linear(
282
model: nn.Module,
283
quant_config: "QuantizationConfig | None",