[Transformers backend] Replace vocab embeddings in recursive_replace (#54760)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

tests/models/transformers/test_backend.py

@@ -10,7 +10,7 @@ from typing import Any
 import pytest
 import torch
 import torch.nn as nn
-from transformers import AutoConfig, AutoModel
+from transformers import AutoConfig, AutoModel, PretrainedConfig
 
 from vllm.config import ModelConfig, VllmConfig
 from vllm.model_executor.models.interfaces import SupportsMultiModal
@@ -316,6 +316,8 @@ def test_pooling(
 
 
 VOCAB_SIZE = 64
+PER_LAYER_VOCAB_SIZE = 32
+NUM_POSITIONS = 16
 HIDDEN_SIZE = 8
 EMBED_SCALE = 3.0
 
@@ -397,10 +399,10 @@ def replace(embedding):
     return new_embedding
 
 
-def assert_scaled(vpe, module, embedding=None):
-    """`module`'s output is `embedding`'s unscaled output times `EMBED_SCALE`."""
+def assert_scaled(vpe, module):
+    """`module`'s output is its own unscaled output times `EMBED_SCALE`."""
     input_ids = torch.arange(VOCAB_SIZE)
-    unscaled = vpe.forward(embedding if embedding is not None else module, input_ids)
+    unscaled = vpe.forward(module, input_ids)
     torch.testing.assert_close(module(input_ids), unscaled * EMBED_SCALE)
 
 
@@ -434,56 +436,106 @@ def test_replace_inherited_embedding(vpe):
     assert_scaled(vpe, new_embedding)
 
 
-def test_replace_composed_embedding(vpe):
-    """Wrappers are left alone; only the `nn.Embedding` they hold is replaced."""
-    embedding = ComposedWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE)
-    new_embedding = replace(embedding)
+def test_replace_is_idempotent(tp_init):
+    """A tied table is reached twice by the recursion, and must not be rebased twice."""
+    from vllm.model_executor.models.transformers.utils import replace_embedding_class
 
-    assert new_embedding is embedding
-    assert type(embedding.embed) is vpe
-    assert_scaled(vpe, new_embedding, embedding.embed)
+    new_embedding = replace(ScaledWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=2))
+    again = replace_embedding_class(new_embedding)
+
+    assert again is new_embedding
+    assert type(again) is type(new_embedding)
 
 
-def test_replace_nested_embedding(vpe):
-    """The composed `nn.Embedding` is found and set however deeply it is nested."""
-    wrapper = nn.Module()
-    wrapper.add_module("inner", nn.Module())
-    wrapper.inner.add_module(
-        "embed", ScaledWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE)
+class FakeVocabModel(nn.Module):
+    """The embedding tables a model can hold, only one of which is a vocab table.
+
+    `embed_tokens_per_layer` mirrors Gemma 3n, whose per-layer table is sized by
+    `vocab_size_per_layer_input` rather than `vocab_size`.
+    """
+
+    def __init__(self, per_layer: bool = False, composed: bool = False):
+        super().__init__()
+        embed_cls = ComposedWordEmbedding if composed else ScaledWordEmbedding
+        self.embed_tokens = embed_cls(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE)
+        if per_layer:
+            self.embed_tokens_per_layer = ScaledWordEmbedding(
+                PER_LAYER_VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE
+            )
+        # Not a vocab table: indexed by position, and read as a whole by models
+        # which interpolate it (e.g. `SiglipVisionEmbeddings`)
+        self.position_embedding = nn.Embedding(NUM_POSITIONS, HIDDEN_SIZE)
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+
+def replace_vocab_embeddings(model, **config_kwargs):
+    """Replace `model`'s vocab tables as `recursive_replace` would, and return them."""
+    from vllm.model_executor.models.transformers.utils import attrsetter
+
+    stub = nn.Module()
+    stub.model = model
+    stub.config = PretrainedConfig(
+        vocab_size=VOCAB_SIZE, num_positions=NUM_POSITIONS, **config_kwargs
     )
-    replace(wrapper)
+    ids = Base._vocab_embedding_ids(stub)
 
-    assert isinstance(wrapper.inner.embed, vpe)
-    assert_scaled(vpe, wrapper.inner.embed)
+    replaced = []
+    for name, module in list(model.named_modules()):
+        if id(module) in ids:
+            replaced.append(replace(module))
+            attrsetter(name)(model, replaced[-1])
+    return replaced
 
 
-@pytest.mark.parametrize("num_embeddings", [0, 2])
-def test_replace_ambiguous_embedding(tp_init, num_embeddings):
-    """Composing anything but one `nn.Embedding` is an error, not a silent guess."""
-    wrapper = nn.Module()
-    for i in range(num_embeddings):
-        wrapper.add_module(f"embed_{i}", nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE))
+def test_only_vocab_tables_are_replaced(vpe):
+    """Position embeddings share the `nn.Embedding` type, but not the treatment."""
+    model = FakeVocabModel()
 
-    with pytest.raises(ValueError, match=f"found {num_embeddings}"):
-        replace(wrapper)
+    assert replace_vocab_embeddings(model) == [model.embed_tokens]
+    assert isinstance(model.embed_tokens, vpe)
+    assert type(model.position_embedding) is nn.Embedding
 
 
-def test_replaced_embedding_exposes_one_vpe(vpe):
-    """`CausalMixin` ties `lm_head` to the one `VocabParallelEmbedding` it can find.
+def test_per_layer_vocab_table_is_replaced(vpe):
+    """A second vocab table is found by its size, without a bespoke accessor."""
+    model = FakeVocabModel(per_layer=True)
+    replaced = replace_vocab_embeddings(
+        model, vocab_size_per_layer_input=PER_LAYER_VOCAB_SIZE
+    )
 
-    `tie_weights` reads `.weight`, which a composing module does not have, so it must
-    be handed the composed embedding instead.
+    assert replaced == [model.embed_tokens, model.embed_tokens_per_layer]
+    assert all(isinstance(module, vpe) for module in replaced)
+
+
+def test_composed_input_embeddings_are_replaced(vpe):
+    """A wrapper is left alone; only the `nn.Embedding` it composes is replaced.
+
+    `CausalMixin` ties `lm_head` to whichever `VocabParallelEmbedding` it finds under
+    `get_input_embeddings()`, reading a `.weight` the wrapper does not have.
     """
     from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 
-    inherited = replace(ScaledWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE))
-    composed = replace(ComposedWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, EMBED_SCALE))
+    model = FakeVocabModel(composed=True)
+    wrapper = model.embed_tokens
+    replaced = replace_vocab_embeddings(model)
 
-    assert [m for m in inherited.modules() if isinstance(m, vpe)] == [inherited]
-    assert [m for m in composed.modules() if isinstance(m, vpe)] == [composed.embed]
+    assert model.embed_tokens is wrapper
+    assert replaced == [m for m in wrapper.modules() if isinstance(m, vpe)]
 
     lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE)
-    assert lm_head.tie_weights(composed.embed).weight is composed.embed.weight
+    assert lm_head.tie_weights(replaced[0]).weight is replaced[0].weight
+
+
+def test_missing_input_embeddings_are_skipped(tp_init):
+    """Pipeline ranks without the embeddings have a `PPMissingLayer` in their place."""
+    from vllm.model_executor.models.utils import PPMissingLayer
+
+    model = FakeVocabModel()
+    model.embed_tokens = PPMissingLayer()
+
+    assert replace_vocab_embeddings(model) == []
 
 
 MULTIMODAL_MODEL = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf"

vllm/model_executor/models/transformers/base.py

@@ -29,7 +29,7 @@ import torch
 import transformers
 from packaging.version import Version
 from torch import nn
-from transformers import AutoModel
+from transformers import AutoModel, PretrainedConfig
 from transformers.conversion_mapping import (
     WeightRenaming,
     get_model_conversion_mapping,
@@ -169,13 +169,6 @@ class Base(
         # Create attention instances for KV cache allocation
         self.attention_instances = self.create_attention_instances()
 
-        # Input embeddings
-        input_embeddings = self.model.get_input_embeddings()
-        if not isinstance(input_embeddings, PPMissingLayer):
-            self.model.set_input_embeddings(
-                replace_embedding_class(input_embeddings, self.quant_config)
-            )
-
         # Initialize any parameters that have not had their modules replaced
         self.init_parameters(self.model)
 
@@ -423,6 +416,31 @@ class Base(
                 # attrsetter in case the module is nested (e.g. "text_model.norm")
                 attrsetter(name)(module, PPMissingLayer())
 
+    def _vocab_embedding_ids(self) -> set[int]:
+        """ids of the `nn.Embedding`s in `self.model` which hold a vocab table."""
+
+        def vocab_sizes(config: PretrainedConfig):
+            for key, value in vars(config).items():
+                if isinstance(value, PretrainedConfig):
+                    yield from vocab_sizes(value)
+                elif "vocab_size" in key and isinstance(value, int):
+                    yield value
+
+        sizes = set(vocab_sizes(self.config))
+        # `get_input_embeddings` may return a module which composes the `nn.Embedding`,
+        # or a `PPMissingLayer` if the embeddings are not on this pipeline stage
+        ids = {
+            id(module)
+            for module in self.model.get_input_embeddings().modules()
+            if isinstance(module, nn.Embedding)
+        }
+        ids.update(
+            id(module)
+            for module in self.model.modules()
+            if isinstance(module, nn.Embedding) and module.num_embeddings in sizes
+        )
+        return ids
+
     def recursive_replace(self):
         """Recursively replace modules in the model as needed.
 
@@ -432,6 +450,7 @@ class Base(
         - Attention QKV projections with a fused `QKVParallelLinear` + split
         - `nn.Linear` with vLLM's tensor parallel linear classes
         - `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d`
+        - Vocab `nn.Embedding`s with vLLM's `VocabParallelEmbedding`
         - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm`
         """
         tp_plan = self.model.tp_plan or {}
@@ -454,6 +473,8 @@ class Base(
         # Detect fusable patterns once per module class (cached, so this is cheap)
         fusers = Fusers(self.model, self.vllm_config)
 
+        vocab_embedding_ids = self._vocab_embedding_ids()
+
         def register_fusion(fuser: BaseFuser, prefix: str):
             """Register a fused layer's mappings just before it is built."""
             self.fusers[prefix] = fuser
@@ -503,6 +524,10 @@ class Base(
                     )
                 elif isinstance(child_module, (nn.Conv2d, nn.Conv3d)):
                     new_module = replace_conv_class(child_module)
+                elif id(child_module) in vocab_embedding_ids:
+                    new_module = replace_embedding_class(
+                        child_module, self.quant_config, prefix=qual_name
+                    )
                 elif (fuser := fusers[child_module]) is not None:
                     register_fusion(fuser, qual_name)
                     new_module = fuser.fuse(child_module, qual_name, self.vllm_config)

vllm/model_executor/models/transformers/utils.py

@@ -224,41 +224,22 @@ def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbe
 
 
 def replace_embedding_class(
-    embedding: nn.Module,
+    embedding: nn.Embedding,
     quant_config: "QuantizationConfig | None" = None,
     *,
     prefix: str = "",
 ) -> nn.Module:
-    """Replace the `nn.Embedding` in `embedding` with `VocabParallelEmbedding`.
+    """Replace `nn.Embedding` with `VocabParallelEmbedding`.
 
     Args:
-        embedding: The module returned by `model.get_input_embeddings()`.
+        embedding: The `nn.Embedding` holding a vocab table.
         quant_config: Quantization config for the new embedding.
         prefix: Qualname of `embedding`, used to look up its quantization method.
     Returns:
-        The module to install with `model.set_input_embeddings()`. Composing and
-        inheriting modules are mutated in place and returned as-is.
-    Raises:
-        ValueError: If `embedding` composes anything other than one `nn.Embedding`,
-            which would leave the input embedding weights ambiguous.
+        The module to install in `embedding`'s place.
     """
-    # If `embedding` composes its `nn.Embedding`, recurse into it
-    if not isinstance(embedding, nn.Embedding):
-        composed = [
-            (name, module)
-            for name, module in embedding.named_modules()
-            if isinstance(module, nn.Embedding)
-        ]
-        if len(composed) != 1:
-            raise ValueError(
-                f"Expected {type(embedding).__name__} to be an `nn.Embedding` or to "
-                f"compose exactly one, but found {len(composed)}."
-            )
-        name, module = composed[0]
-        new_embedding = replace_embedding_class(
-            module, quant_config, prefix=maybe_prefix(prefix, name)
-        )
-        attrsetter(name)(embedding, new_embedding)
+    # Tied tables are reached more than once, so don't rebase an already new embedding
+    if isinstance(embedding, VocabParallelEmbedding):
         return embedding
 
     kwargs = dict(