@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f1e5fdd7f2

[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>

Harry Mellor committed Sep 1, 2026 at 16:35 UTC f1e5fdd7f2f4381cd3bd4e4b821fad4604bd72ab

3 files changed +128 -70

tests/models/transformers/test_backend.py

+89 -37

@@ -10,7 +10,7 @@ from typing import Any

10 import pytest

11 import torch

12 import torch.nn as nn

13 -from transformers import AutoConfig, AutoModel

13 +from transformers import AutoConfig, AutoModel, PretrainedConfig

15 from vllm.config import ModelConfig, VllmConfig

16 from vllm.model_executor.models.interfaces import SupportsMultiModal

@@ -316,6 +316,8 @@ def test_pooling(

316

317

318 VOCAB_SIZE = 64

319 +PER_LAYER_VOCAB_SIZE = 32

320 +NUM_POSITIONS = 16

321 HIDDEN_SIZE = 8

322 EMBED_SCALE = 3.0

323

@@ -397,10 +399,10 @@ def replace(embedding):

399 return new_embedding

400

401

400 -def assert_scaled(vpe, module, embedding=None):

401 - """`module`'s output is `embedding`'s unscaled output times `EMBED_SCALE`."""

402 +def assert_scaled(vpe, module):

403 + """`module`'s output is its own unscaled output times `EMBED_SCALE`."""

404 input_ids = torch.arange(VOCAB_SIZE)

403 - unscaled = vpe.forward(embedding if embedding is not None else module, input_ids)

405 + unscaled = vpe.forward(module, input_ids)

406 torch.testing.assert_close(module(input_ids), unscaled * EMBED_SCALE)

407

408

@@ -434,56 +436,106 @@ def test_replace_inherited_embedding(vpe):

436 assert_scaled(vpe, new_embedding)

437

438

437 -def test_replace_composed_embedding(vpe):

438 - """Wrappers are left alone; only the `nn.Embedding` they hold is replaced."""

439 - embedding = ComposedWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE)

440 - new_embedding = replace(embedding)

439 +def test_replace_is_idempotent(tp_init):

440 + """A tied table is reached twice by the recursion, and must not be rebased twice."""

441 + from vllm.model_executor.models.transformers.utils import replace_embedding_class

442 +

443 + new_embedding = replace(ScaledWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=2))

444 + again = replace_embedding_class(new_embedding)

445 +

446 + assert again is new_embedding

447 + assert type(again) is type(new_embedding)

448 +

449 +

450 +class FakeVocabModel(nn.Module):

451 + """The embedding tables a model can hold, only one of which is a vocab table.

452 +

453 + `embed_tokens_per_layer` mirrors Gemma 3n, whose per-layer table is sized by

454 + `vocab_size_per_layer_input` rather than `vocab_size`.

455 + """

456 +

457 + def __init__(self, per_layer: bool = False, composed: bool = False):

458 + super().__init__()

459 + embed_cls = ComposedWordEmbedding if composed else ScaledWordEmbedding

460 + self.embed_tokens = embed_cls(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE)

461 + if per_layer:

462 + self.embed_tokens_per_layer = ScaledWordEmbedding(

463 + PER_LAYER_VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE

464 + )

465 + # Not a vocab table: indexed by position, and read as a whole by models

466 + # which interpolate it (e.g. `SiglipVisionEmbeddings`)

467 + self.position_embedding = nn.Embedding(NUM_POSITIONS, HIDDEN_SIZE)

468 +

469 + def get_input_embeddings(self):

470 + return self.embed_tokens

471

442 - assert new_embedding is embedding

443 - assert type(embedding.embed) is vpe

444 - assert_scaled(vpe, new_embedding, embedding.embed)

472

473 +def replace_vocab_embeddings(model, **config_kwargs):

474 + """Replace `model`'s vocab tables as `recursive_replace` would, and return them."""

475 + from vllm.model_executor.models.transformers.utils import attrsetter

476

447 -def test_replace_nested_embedding(vpe):

448 - """The composed `nn.Embedding` is found and set however deeply it is nested."""

449 - wrapper = nn.Module()

450 - wrapper.add_module("inner", nn.Module())

451 - wrapper.inner.add_module(

452 - "embed", ScaledWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, embed_scale=EMBED_SCALE)

477 + stub = nn.Module()

478 + stub.model = model

479 + stub.config = PretrainedConfig(

480 + vocab_size=VOCAB_SIZE, num_positions=NUM_POSITIONS, **config_kwargs

481 )

454 - replace(wrapper)

482 + ids = Base._vocab_embedding_ids(stub)

483 +

484 + replaced = []

485 + for name, module in list(model.named_modules()):

486 + if id(module) in ids:

487 + replaced.append(replace(module))

488 + attrsetter(name)(model, replaced[-1])

489 + return replaced

490 +

491 +

492 +def test_only_vocab_tables_are_replaced(vpe):

493 + """Position embeddings share the `nn.Embedding` type, but not the treatment."""

494 + model = FakeVocabModel()

495

456 - assert isinstance(wrapper.inner.embed, vpe)

457 - assert_scaled(vpe, wrapper.inner.embed)

496 + assert replace_vocab_embeddings(model) == [model.embed_tokens]

497 + assert isinstance(model.embed_tokens, vpe)

498 + assert type(model.position_embedding) is nn.Embedding

499

500

460 -@pytest.mark.parametrize("num_embeddings", [0, 2])

461 -def test_replace_ambiguous_embedding(tp_init, num_embeddings):

462 - """Composing anything but one `nn.Embedding` is an error, not a silent guess."""

463 - wrapper = nn.Module()

464 - for i in range(num_embeddings):

465 - wrapper.add_module(f"embed_{i}", nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE))

501 +def test_per_layer_vocab_table_is_replaced(vpe):

502 + """A second vocab table is found by its size, without a bespoke accessor."""

503 + model = FakeVocabModel(per_layer=True)

504 + replaced = replace_vocab_embeddings(

505 + model, vocab_size_per_layer_input=PER_LAYER_VOCAB_SIZE

506 + )

507

467 - with pytest.raises(ValueError, match=f"found {num_embeddings}"):

468 - replace(wrapper)

508 + assert replaced == [model.embed_tokens, model.embed_tokens_per_layer]

509 + assert all(isinstance(module, vpe) for module in replaced)

510

511

471 -def test_replaced_embedding_exposes_one_vpe(vpe):

472 - """`CausalMixin` ties `lm_head` to the one `VocabParallelEmbedding` it can find.

512 +def test_composed_input_embeddings_are_replaced(vpe):

513 + """A wrapper is left alone; only the `nn.Embedding` it composes is replaced.

514

474 - `tie_weights` reads `.weight`, which a composing module does not have, so it must

475 - be handed the composed embedding instead.

515 + `CausalMixin` ties `lm_head` to whichever `VocabParallelEmbedding` it finds under

516 + `get_input_embeddings()`, reading a `.weight` the wrapper does not have.

517 """

518 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

519

479 - inherited = replace(ScaledWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE))

480 - composed = replace(ComposedWordEmbedding(VOCAB_SIZE, HIDDEN_SIZE, EMBED_SCALE))

520 + model = FakeVocabModel(composed=True)

521 + wrapper = model.embed_tokens

522 + replaced = replace_vocab_embeddings(model)

523

482 - assert [m for m in inherited.modules() if isinstance(m, vpe)] == [inherited]

483 - assert [m for m in composed.modules() if isinstance(m, vpe)] == [composed.embed]

524 + assert model.embed_tokens is wrapper

525 + assert replaced == [m for m in wrapper.modules() if isinstance(m, vpe)]

526

527 lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE)

486 - assert lm_head.tie_weights(composed.embed).weight is composed.embed.weight

528 + assert lm_head.tie_weights(replaced[0]).weight is replaced[0].weight

529 +

530 +

531 +def test_missing_input_embeddings_are_skipped(tp_init):

532 + """Pipeline ranks without the embeddings have a `PPMissingLayer` in their place."""

533 + from vllm.model_executor.models.utils import PPMissingLayer

534 +

535 + model = FakeVocabModel()

536 + model.embed_tokens = PPMissingLayer()

537 +

538 + assert replace_vocab_embeddings(model) == []

539

540

541 MULTIMODAL_MODEL = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf"

vllm/model_executor/models/transformers/base.py

+33 -8

@@ -29,7 +29,7 @@ import torch

29 import transformers

30 from packaging.version import Version

31 from torch import nn

32 -from transformers import AutoModel

32 +from transformers import AutoModel, PretrainedConfig

33 from transformers.conversion_mapping import (

34 WeightRenaming,

35 get_model_conversion_mapping,

@@ -169,13 +169,6 @@ class Base(

169 # Create attention instances for KV cache allocation

170 self.attention_instances = self.create_attention_instances()

171

172 - # Input embeddings

173 - input_embeddings = self.model.get_input_embeddings()

174 - if not isinstance(input_embeddings, PPMissingLayer):

175 - self.model.set_input_embeddings(

176 - replace_embedding_class(input_embeddings, self.quant_config)

177 - )

178 -

172 # Initialize any parameters that have not had their modules replaced

173 self.init_parameters(self.model)

174

@@ -423,6 +416,31 @@ class Base(

416 # attrsetter in case the module is nested (e.g. "text_model.norm")

417 attrsetter(name)(module, PPMissingLayer())

418

419 + def _vocab_embedding_ids(self) -> set[int]:

420 + """ids of the `nn.Embedding`s in `self.model` which hold a vocab table."""

421 +

422 + def vocab_sizes(config: PretrainedConfig):

423 + for key, value in vars(config).items():

424 + if isinstance(value, PretrainedConfig):

425 + yield from vocab_sizes(value)

426 + elif "vocab_size" in key and isinstance(value, int):

427 + yield value

428 +

429 + sizes = set(vocab_sizes(self.config))

430 + # `get_input_embeddings` may return a module which composes the `nn.Embedding`,

431 + # or a `PPMissingLayer` if the embeddings are not on this pipeline stage

432 + ids = {

433 + id(module)

434 + for module in self.model.get_input_embeddings().modules()

435 + if isinstance(module, nn.Embedding)

436 + }

437 + ids.update(

438 + id(module)

439 + for module in self.model.modules()

440 + if isinstance(module, nn.Embedding) and module.num_embeddings in sizes

441 + )

442 + return ids

443 +

444 def recursive_replace(self):

445 """Recursively replace modules in the model as needed.

446

@@ -432,6 +450,7 @@ class Base(

450 - Attention QKV projections with a fused `QKVParallelLinear` + split

451 - `nn.Linear` with vLLM's tensor parallel linear classes

452 - `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d`

453 + - Vocab `nn.Embedding`s with vLLM's `VocabParallelEmbedding`

454 - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm`

455 """

456 tp_plan = self.model.tp_plan or {}

@@ -454,6 +473,8 @@ class Base(

473 # Detect fusable patterns once per module class (cached, so this is cheap)

474 fusers = Fusers(self.model, self.vllm_config)

475

476 + vocab_embedding_ids = self._vocab_embedding_ids()

477 +

478 def register_fusion(fuser: BaseFuser, prefix: str):

479 """Register a fused layer's mappings just before it is built."""

480 self.fusers[prefix] = fuser

@@ -503,6 +524,10 @@ class Base(

524 )

525 elif isinstance(child_module, (nn.Conv2d, nn.Conv3d)):

526 new_module = replace_conv_class(child_module)

527 + elif id(child_module) in vocab_embedding_ids:

528 + new_module = replace_embedding_class(

529 + child_module, self.quant_config, prefix=qual_name

530 + )

531 elif (fuser := fusers[child_module]) is not None:

532 register_fusion(fuser, qual_name)

533 new_module = fuser.fuse(child_module, qual_name, self.vllm_config)

vllm/model_executor/models/transformers/utils.py

+6 -25

@@ -224,41 +224,22 @@ def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbe

224

225

226 def replace_embedding_class(

227 - embedding: nn.Module,

227 + embedding: nn.Embedding,

228 quant_config: "QuantizationConfig | None" = None,

229 *,

230 prefix: str = "",

231 ) -> nn.Module:

232 - """Replace the `nn.Embedding` in `embedding` with `VocabParallelEmbedding`.

232 + """Replace `nn.Embedding` with `VocabParallelEmbedding`.

233

234 Args:

235 - embedding: The module returned by `model.get_input_embeddings()`.

235 + embedding: The `nn.Embedding` holding a vocab table.

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.

239 + The module to install in `embedding`'s place.

240 """

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)

241 + # Tied tables are reached more than once, so don't rebase an already new embedding

242 + if isinstance(embedding, VocabParallelEmbedding):

243 return embedding

244

245 kwargs = dict(