diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py index 5fd39a5b9d..621c249137 100644 --- a/tests/models/transformers/test_backend.py +++ b/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" diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 1d8fa5b5d7..0377b5708a 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/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) diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index ed6c47a86e..116dd7b516 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/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(