[Bugfix] Retain vocab embeddings during replacement (#55083)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Sep 3, 2026 at 06:17 UTC
758c79e1c5fc228a78d154f59dc0521e063e7a93
2 files changed
+11
-11
tests/models/transformers/test_backend.py
+2
-2
@@ -479,11 +479,11 @@ def replace_vocab_embeddings(model, **config_kwargs):
479
stub.config = PretrainedConfig(
480
vocab_size=VOCAB_SIZE, num_positions=NUM_POSITIONS, **config_kwargs
481
)
482
- ids = Base._vocab_embedding_ids(stub)
482
+ embeddings = Base._vocab_embeddings(stub)
483
484
replaced = []
485
for name, module in list(model.named_modules()):
486
- if id(module) in ids:
486
+ if module in embeddings:
487
replaced.append(replace(module))
488
attrsetter(name)(model, replaced[-1])
489
return replaced
vllm/model_executor/models/transformers/base.py
+9
-9
@@ -416,8 +416,8 @@ 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."""
419
+ def _vocab_embeddings(self) -> set[nn.Embedding]:
420
+ """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():
@@ -429,17 +429,17 @@ class Base(
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)
432
+ embeddings = {
433
+ module
434
for module in self.model.get_input_embeddings().modules()
435
if isinstance(module, nn.Embedding)
436
}
437
- ids.update(
438
- id(module)
437
+ embeddings.update(
438
+ module
439
for module in self.model.modules()
440
if isinstance(module, nn.Embedding) and module.num_embeddings in sizes
441
)
442
- return ids
442
+ return embeddings
443
444
def recursive_replace(self):
445
"""Recursively replace modules in the model as needed.
@@ -473,7 +473,7 @@ 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()
476
+ vocab_embeddings = self._vocab_embeddings()
477
478
def register_fusion(fuser: BaseFuser, prefix: str):
479
"""Register a fused layer's mappings just before it is built."""
@@ -524,7 +524,7 @@ 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:
527
+ elif child_module in vocab_embeddings:
528
new_module = replace_embedding_class(
529
child_module, self.quant_config, prefix=qual_name
530
)