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

Fix weight tying (#51665)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed Aug 20, 2026 at 11:31 UTC f0c14b4f776bb976e654b430442067539fbdb2ea
70 files changed +526 -238
tests/model_executor/model_loader/test_weight_tying.py new
+86
@@ -0,0 +1,86 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +
6 +import pytest
7 +import torch
8 +from torch import nn
9 +
10 +from vllm.model_executor.layers.vocab_parallel_embedding import (
11 + ParallelLMHead,
12 + VocabParallelEmbedding,
13 +)
14 +from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings
15 +
16 +VOCAB_SIZE = 16
17 +HIDDEN_SIZE = 4
18 +
19 +
20 +class UntiedModel(nn.Module):
21 + """Nests the head like a multimodal model, which is the harder case."""
22 +
23 + def __init__(self):
24 + super().__init__()
25 + self.language_model = nn.Module()
26 + self.language_model.model = nn.Module()
27 + self.language_model.model.embed_tokens = VocabParallelEmbedding(
28 + VOCAB_SIZE, HIDDEN_SIZE
29 + )
30 + self.language_model.lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE)
31 + self.embed_tokens.weight.data.fill_(1.0)
32 + self.lm_head.weight.data.fill_(1.0)
33 +
34 + @property
35 + def embed_tokens(self) -> VocabParallelEmbedding:
36 + return self.language_model.model.embed_tokens
37 +
38 + @property
39 + def lm_head(self) -> ParallelLMHead:
40 + return self.language_model.lm_head
41 +
42 +
43 +def make_model_config(untied_by_checkpoint=False):
44 + return SimpleNamespace(
45 + model="dummy-model",
46 + word_embeddings_untied_by_checkpoint=untied_by_checkpoint,
47 + )
48 +
49 +
50 +@pytest.mark.cpu_test
51 +@pytest.mark.usefixtures("dist_init")
52 +@pytest.mark.parametrize("identical", [True, False])
53 +def test_retie_only_when_identical(identical: bool):
54 + """A redundant copy of a tied lm_head is shared again to reclaim memory."""
55 + model = UntiedModel()
56 + if not identical:
57 + model.lm_head.weight.data.fill_(2.0)
58 +
59 + maybe_retie_word_embeddings(model, make_model_config(untied_by_checkpoint=True))
60 +
61 + assert (model.lm_head.weight is model.embed_tokens.weight) is identical
62 + if not identical:
63 + assert torch.all(model.lm_head.weight == 2.0)
64 +
65 +
66 +@pytest.mark.cpu_test
67 +@pytest.mark.usefixtures("dist_init")
68 +def test_quantized_lm_head_is_left_alone():
69 + """A quantized head may store its weights packed under another name."""
70 + model = UntiedModel()
71 + model.lm_head.quant_method = SimpleNamespace()
72 +
73 + maybe_retie_word_embeddings(model, make_model_config(untied_by_checkpoint=True))
74 +
75 + assert model.lm_head.weight is not model.embed_tokens.weight
76 +
77 +
78 +@pytest.mark.cpu_test
79 +@pytest.mark.usefixtures("dist_init")
80 +def test_no_retie_without_checkpoint_override():
81 + """Word embeddings the config genuinely unties are left alone."""
82 + model = UntiedModel()
83 +
84 + maybe_retie_word_embeddings(model, make_model_config())
85 +
86 + assert model.lm_head.weight is not model.embed_tokens.weight
tests/models/test_utils.py
+78
@@ -4,6 +4,10 @@
4 import pytest
5 import torch
6
7 +from vllm.model_executor.layers.vocab_parallel_embedding import (
8 + ParallelLMHead,
9 + VocabParallelEmbedding,
10 +)
11 from vllm.model_executor.models.utils import (
12 AutoWeightsLoader,
13 _merge_multimodal_embeddings,
@@ -165,6 +169,80 @@ def test_module_skip_substr():
169 assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1
170
171
172 +VOCAB_SIZE = 16
173 +HIDDEN_SIZE = 2
174 +
175 +
176 +class ModuleWithTiedWeights(torch.nn.Module):
177 + """Mimics how models tie `lm_head` to the input embeddings."""
178 +
179 + def __init__(self, tie: bool):
180 + super().__init__()
181 + self.model = torch.nn.Module()
182 + self.model.embed_tokens = VocabParallelEmbedding(VOCAB_SIZE, HIDDEN_SIZE)
183 + self.lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE)
184 + if tie:
185 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
186 +
187 +
188 +def make_embedding_weights(value: float) -> torch.Tensor:
189 + return torch.full((VOCAB_SIZE, HIDDEN_SIZE), value)
190 +
191 +
192 +@pytest.mark.cpu_test
193 +@pytest.mark.usefixtures("dist_init")
194 +@pytest.mark.parametrize("tie", [True, False])
195 +def test_module_skip_tied_weights(tie: bool):
196 + """Tied weights must be loaded once, under the first of their names."""
197 + mod = ModuleWithTiedWeights(tie)
198 +
199 + weights = [
200 + ("model.embed_tokens.weight", make_embedding_weights(1.0)),
201 + ("lm_head.weight", make_embedding_weights(2.0)),
202 + ]
203 + loaded = AutoWeightsLoader(mod).load_weights(iter(weights))
204 +
205 + if tie:
206 + assert loaded == {"model.embed_tokens.weight"}
207 + assert torch.all(mod.lm_head.weight[:VOCAB_SIZE] == 1.0)
208 + else:
209 + assert loaded == {"model.embed_tokens.weight", "lm_head.weight"}
210 + assert torch.all(mod.lm_head.weight[:VOCAB_SIZE] == 2.0)
211 +
212 +
213 +@pytest.mark.cpu_test
214 +@pytest.mark.usefixtures("dist_init")
215 +def test_module_skip_tied_weights_without_canonical():
216 + """Skipping a tied weight must not leave the shared weight uninitialized."""
217 + mod = ModuleWithTiedWeights(tie=True)
218 +
219 + weights = [("lm_head.weight", make_embedding_weights(2.0))]
220 + with pytest.raises(ValueError, match="model.embed_tokens.weight"):
221 + AutoWeightsLoader(mod).load_weights(iter(weights))
222 +
223 +
224 +class ModuleWithSharedParam(torch.nn.Module):
225 + """Mimics an MoE router shared between the MLP and its fused experts."""
226 +
227 + def __init__(self):
228 + super().__init__()
229 + self.experts = torch.nn.Module()
230 + self.experts.gate = torch.nn.Linear(2, 2, bias=False)
231 + self.gate = self.experts.gate
232 +
233 +
234 +@pytest.mark.cpu_test
235 +def test_module_load_shared_params_that_are_not_tied_embeddings():
236 + """Only tied embeddings are skipped; other shared params must still load."""
237 + mod = ModuleWithSharedParam()
238 +
239 + weights = [("gate.weight", torch.Tensor([[1, 2], [3, 4]]))]
240 + loaded = AutoWeightsLoader(mod).load_weights(iter(weights))
241 +
242 + assert loaded == {"gate.weight"}
243 + assert torch.all(mod.gate.weight == torch.Tensor([[1, 2], [3, 4]]))
244 +
245 +
246 class raise_if_cuda_sync:
247 def __enter__(self):
248 self.previous_debug_mode = torch.cuda.get_sync_debug_mode()
tests/test_config.py
+39
@@ -604,6 +604,45 @@ def test_with_hf_config_leaves_unknown_model_type_without_architectures(
604 assert updated.model_config.hf_config.architectures is None
605
606
607 +@pytest.mark.parametrize(
608 + "checkpoint_tensors,tied",
609 + [
610 + # The checkpoint has an lm_head of its own, so it must win over the config
611 + (["model.embed_tokens.weight", "lm_head.weight"], False),
612 + (["model.embed_tokens.weight"], True),
613 + # Contents unknown (not safetensors), so the config must be left alone
614 + ([], True),
615 + ],
616 +)
617 +def test_maybe_untie_word_embeddings(tmp_path, checkpoint_tensors, tied):
618 + import torch
619 + from safetensors.torch import save_file
620 +
621 + if checkpoint_tensors:
622 + save_file(
623 + {name: torch.zeros(2, 2) for name in checkpoint_tensors},
624 + tmp_path / "model.safetensors",
625 + )
626 +
627 + text_config = SimpleNamespace(tie_word_embeddings=True)
628 + model_config = SimpleNamespace(
629 + model=str(tmp_path),
630 + revision=None,
631 + hf_config=SimpleNamespace(
632 + tie_word_embeddings=True,
633 + get_text_config=lambda: text_config,
634 + ),
635 + word_embeddings_untied_by_checkpoint=False,
636 + )
637 +
638 + ModelConfig.maybe_untie_word_embeddings(model_config)
639 +
640 + # Both levels must agree, since different callers read different ones
641 + assert model_config.hf_config.tie_word_embeddings is tied
642 + assert text_config.tie_word_embeddings is tied
643 + assert model_config.word_embeddings_untied_by_checkpoint is not tied
644 +
645 +
646 def test_async_scheduling_with_pipeline_parallelism_is_allowed():
647 cfg = VllmConfig(
648 scheduler_config=SchedulerConfig(
tests/transformers_utils/test_config.py
+24 -2
@@ -8,14 +8,17 @@ only get the `eos_token_id` from the tokenizer as defined by
8
9 from types import SimpleNamespace
10 from typing import cast
11 -from unittest.mock import patch
11 +from unittest.mock import MagicMock, patch
12
13 from transformers import PretrainedConfig
14
15 from vllm.config.model import ModelConfig
16 from vllm.tokenizers import get_tokenizer
17 from vllm.transformers_utils import config as config_module
18 -from vllm.transformers_utils.config import try_get_generation_config
18 +from vllm.transformers_utils.config import (
19 + get_safetensors_params_metadata,
20 + try_get_generation_config,
21 +)
22
23
24 def test_get_llama3_eos_token():
@@ -77,3 +80,22 @@ def test_model_config_generation_fallback_forwards_code_revision():
80 config_format="auto",
81 token=None,
82 )
83 +
84 +
85 +def test_safetensors_metadata_of_repo_without_safetensors():
86 + """A repo storing its weights in another format is an answer, not a failure,
87 + so it must not be retried."""
88 + from huggingface_hub.errors import LocalEntryNotFoundError, NotASafetensorsRepoError
89 +
90 + get_safetensors_metadata = MagicMock(
91 + side_effect=NotASafetensorsRepoError("not a safetensors repo")
92 + )
93 + api = SimpleNamespace(
94 + get_safetensors_metadata=get_safetensors_metadata,
95 + snapshot_download=MagicMock(side_effect=LocalEntryNotFoundError("no cache")),
96 + )
97 +
98 + with patch.object(config_module, "hf_api", lambda: api):
99 + assert get_safetensors_params_metadata("some/pytorch-only-model") == {}
100 +
101 + get_safetensors_metadata.assert_called_once()
tests/transformers_utils/test_repo_utils.py
+11
@@ -14,6 +14,7 @@ from vllm.transformers_utils.repo_utils import (
14 get_hf_file_to_dict,
15 is_mistral_model_repo,
16 list_filtered_repo_files,
17 + with_retry,
18 )
19
20
@@ -185,3 +186,13 @@ def test_is_mistral_model_repo(files: list[str], expected_bool: bool):
186 repo_type="model",
187 token="token",
188 )
189 +
190 +
191 +def test_with_retry_does_not_retry_fatal_errors():
192 + """A definitive answer must not be retried, which costs a delay per attempt."""
193 + func = MagicMock(side_effect=FileNotFoundError("no such file"))
194 +
195 + with pytest.raises(FileNotFoundError):
196 + with_retry(func, "Error", fatal_errors=(FileNotFoundError,))
197 +
198 + func.assert_called_once()
vllm/config/model.py
+33
@@ -30,6 +30,7 @@ from vllm.platforms import current_platform
30 from vllm.tasks import PoolingTask, ScoreType, SupportedTask
31 from vllm.transformers_utils.config import (
32 ConfigFormat,
33 + checkpoint_has_lm_head,
34 get_config,
35 get_hf_image_processor_config,
36 get_hf_text_config,
@@ -185,6 +186,10 @@ class ModelConfig:
186 """The Hugging Face config of the model."""
187 hf_text_config: PretrainedConfig = field(init=False)
188 """The Hugging Face config of the text model (same as hf_config for text models)."""
189 + word_embeddings_untied_by_checkpoint: bool = field(default=False, init=False)
190 + """Whether `tie_word_embeddings` was overridden to `False` because the checkpoint
191 + contains an `lm_head` of its own. The two may still turn out to be identical, in
192 + which case they are re-tied once the weights have been loaded."""
193 hf_config_path: str | None = None
194 """Name or path of the Hugging Face config to use. If unspecified, model
195 name or path will be used."""
@@ -1109,6 +1114,34 @@ class ModelConfig:
1114 def _get_encoder_config(self) -> dict[str, Any] | None:
1115 return get_sentence_transformer_tokenizer_config(self.model, self.revision)
1116
1117 + def maybe_untie_word_embeddings(self) -> None:
1118 + """Stop trusting `tie_word_embeddings` when the checkpoint disagrees.
1119 +
1120 + A config may claim the word embeddings are tied while the checkpoint
1121 + ships an `lm_head` of its own. Tying regardless would silently discard
1122 + that tensor, so build the `lm_head` as if untied and let it load. The
1123 + two are compared once loaded, and re-tied if they turn out to match, by
1124 + [maybe_retie_word_embeddings][vllm.model_executor.model_loader.weight_tying.maybe_retie_word_embeddings].
1125 +
1126 + Transformers makes the same decision in `PreTrainedModel.tie_weights`,
1127 + where it can compare the two tensors directly.
1128 + """
1129 + if not getattr(self.hf_config, "tie_word_embeddings", False):
1130 + return
1131 + if not checkpoint_has_lm_head(self.model, revision=self.revision):
1132 + return
1133 +
1134 + logger.debug(
1135 + "The config for %s says the word embeddings are tied, but the checkpoint "
1136 + "contains an lm_head. Loading it to find out whether they really are tied.",
1137 + self.model,
1138 + )
1139 + # Both levels must be updated: `VllmConfig.with_hf_config` reads the top
1140 + # level, and the text config is what the language model itself sees.
1141 + self.hf_config.tie_word_embeddings = False
1142 + self.hf_config.get_text_config().tie_word_embeddings = False
1143 + self.word_embeddings_untied_by_checkpoint = True
1144 +
1145 def _get_default_runner_type(
1146 self,
1147 architectures: list[str],
vllm/config/vllm.py
+10
@@ -1170,6 +1170,16 @@ class VllmConfig:
1170 self.model_config, self.load_config
1171 )
1172
1173 + # "dummy" reads no weights at all, and the sharded formats read a vLLM
1174 + # state dict, which stores tied word embeddings under the lm_head only.
1175 + # Neither can tell us what the original checkpoint contained.
1176 + if self.model_config is not None and self.load_config.load_format not in (
1177 + "dummy",
1178 + "sharded_state",
1179 + "runai_streamer_sharded",
1180 + ):
1181 + self.model_config.maybe_untie_word_embeddings()
1182 +
1183 if (
1184 self.quant_config is not None
1185 and self.model_config is not None
vllm/model_executor/model_loader/utils.py
+5
@@ -24,6 +24,7 @@ from vllm.model_executor.model_loader.reload import (
24 record_metadata_for_reloading,
25 set_torchao_reload_attrs,
26 )
27 +from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings
28 from vllm.model_executor.models.interfaces import SupportsQuant
29 from vllm.tracing import instrument
30 from vllm.utils.mem_utils import release_device_memory_under_pressure
@@ -96,6 +97,10 @@ def initialize_model(
97 def process_weights_after_loading(
98 model: nn.Module, model_config: ModelConfig, target_device: torch.device
99 ) -> None:
100 + # Reclaim memory when an explicit lm_head has been
101 + # loaded, but it is identical to the input embeddings.
102 + maybe_retie_word_embeddings(model, model_config)
103 +
104 for _, module in model.named_modules():
105 quant_method = getattr(module, "quant_method", None)
106 if isinstance(quant_method, QuantizeMethodBase):
vllm/model_executor/model_loader/weight_tying.py new
+99
@@ -0,0 +1,99 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Reconcile word embedding tying with what the checkpoint actually contains."""
4 +
5 +from dataclasses import dataclass
6 +
7 +import torch
8 +from torch import nn
9 +
10 +from vllm.config import ModelConfig
11 +from vllm.logger import init_logger
12 +from vllm.model_executor.layers.vocab_parallel_embedding import (
13 + ParallelLMHead,
14 + UnquantizedEmbeddingMethod,
15 + VocabParallelEmbedding,
16 +)
17 +
18 +logger = init_logger(__name__)
19 +
20 +
21 +@dataclass
22 +class _UntiedLMHead:
23 + name: str
24 + lm_head: ParallelLMHead
25 + embed_tokens: VocabParallelEmbedding
26 +
27 + @property
28 + def weight_name(self) -> str:
29 + return f"{self.name}.weight"
30 +
31 + def tie(self, model: nn.Module) -> None:
32 + parent_name, _, attr = self.name.rpartition(".")
33 + parent = model.get_submodule(parent_name) if parent_name else model
34 + setattr(parent, attr, self.lm_head.tie_weights(self.embed_tokens))
35 +
36 +
37 +def _get_untied_lm_head(model: nn.Module) -> _UntiedLMHead | None:
38 + """Locate an `lm_head` that could be tied to the input embeddings.
39 +
40 + Returns `None` unless the model has exactly one of each, so that models with
41 + several heads (such as MTP) or with already tied weights are left alone.
42 + Both are found by type rather than by name, because `get_input_embeddings`
43 + takes different arguments on multimodal models and the `lm_head` of a
44 + multimodal model is nested inside its language model.
45 +
46 + A quantized `lm_head` is also left alone. Its weights may be packed under
47 + another name, and online quantization creates them after loading, so
48 + neither their contents nor whether they were loaded can be established here.
49 + Note that this is a property of the layer, not of the model: most quantized
50 + models leave the `lm_head` and the input embeddings unquantized.
51 + """
52 + heads = list[tuple[str, ParallelLMHead]]()
53 + embeddings = list[VocabParallelEmbedding]()
54 + for name, module in model.named_modules():
55 + if not isinstance(module, VocabParallelEmbedding):
56 + continue
57 + if not isinstance(module.quant_method, UnquantizedEmbeddingMethod):
58 + continue
59 + if isinstance(module, ParallelLMHead):
60 + heads.append((name, module))
61 + else:
62 + embeddings.append(module)
63 +
64 + if len(heads) != 1 or len(embeddings) != 1:
65 + return None
66 +
67 + (name, lm_head), embed_tokens = heads[0], embeddings[0]
68 + if lm_head.weight.shape != embed_tokens.weight.shape:
69 + return None
70 + return _UntiedLMHead(name, lm_head, embed_tokens)
71 +
72 +
73 +def maybe_retie_word_embeddings(model: nn.Module, model_config: ModelConfig) -> None:
74 + """Re-tie word embeddings that
75 + [ModelConfig.maybe_untie_word_embeddings][vllm.config.ModelConfig.maybe_untie_word_embeddings]
76 + untied, if the loaded `lm_head` turned out to be identical to the input embeddings
77 + after all.
78 +
79 + Checkpoints produced by quantization or fine-tuning tooling often keep a redundant
80 + copy of the tied `lm_head`. Sharing the storage again reclaims the memory it would
81 + otherwise cost."""
82 + if not model_config.word_embeddings_untied_by_checkpoint:
83 + return
84 + if (untied := _get_untied_lm_head(model)) is None:
85 + return
86 +
87 + # On device, torch.equal segfaults on ROCm when sleep mode is enabled
88 + if not torch.equal(untied.lm_head.weight.cpu(), untied.embed_tokens.weight.cpu()):
89 + logger.warning(
90 + "The config for %s says the word embeddings are tied, but the checkpoint "
91 + "contains a different %s, which has been used instead of tying. "
92 + "Set `tie_word_embeddings=False` in the config to silence this warning.",
93 + model_config.model,
94 + untied.weight_name,
95 + )
96 + return
97 +
98 + logger.debug("Re-tying %s, which is identical to the input embeddings", untied.name)
99 + untied.tie(model)
vllm/model_executor/models/apertus.py
+1 -4
@@ -491,8 +491,5 @@ class ApertusForCausalLM(
491 return logits
492
493 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
494 - loader = AutoWeightsLoader(
495 - self,
496 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
497 - )
494 + loader = AutoWeightsLoader(self)
495 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/arcee.py
+1 -5
@@ -356,11 +356,7 @@ class ArceeForCausalLM(
356 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
357 """Load weights into the model (delegates to inner model and handles
358 tied embeddings)."""
359 - loader = AutoWeightsLoader(
360 - self,
361 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
362 - skip_substrs=["gate_proj"],
363 - )
359 + loader = AutoWeightsLoader(self, skip_substrs=["gate_proj"])
360 # AutoWeightLoader handles weight name remapping, including fusing
361 # separate q_proj, k_proj, v_proj into qkv_proj
362 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/arctic.py
+1 -4
@@ -589,8 +589,5 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
589 return logits
590
591 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
592 - loader = AutoWeightsLoader(
593 - self,
594 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
595 - )
592 + loader = AutoWeightsLoader(self)
593 return loader.load_weights(weights)
vllm/model_executor/models/bailing_moe.py
+1 -4
@@ -557,10 +557,7 @@ class BailingMoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
557 yield name, loaded_weight
558
559 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
560 - loader = AutoWeightsLoader(
561 - self,
562 - skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
563 - )
560 + loader = AutoWeightsLoader(self)
561 return loader.load_weights(self._normalize_lm_head(weights))
562
563
vllm/model_executor/models/bloom.py
+1 -1
@@ -360,7 +360,7 @@ class BloomForCausalLM(nn.Module, SupportsPP, SupportsQuant):
360 return logits
361
362 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
363 - loader = AutoWeightsLoader(self, skip_prefixes=["lm_head.weight"])
363 + loader = AutoWeightsLoader(self)
364 weights = _add_transformer_prefix(weights)
365 return loader.load_weights(weights)
366
vllm/model_executor/models/chameleon.py
+1 -2
@@ -1059,6 +1059,5 @@ class ChameleonForConditionalGeneration(
1059 return logits
1060
1061 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1062 - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
1063 - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
1062 + loader = AutoWeightsLoader(self)
1063 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/cohere2_moe.py
-1
@@ -493,7 +493,6 @@ class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
493 config = vllm_config.model_config.hf_config
494 quant_config = vllm_config.quant_config
495 self.config = config
496 - assert getattr(config, "tie_word_embeddings", True)
496 self.unpadded_vocab_size = config.vocab_size
497 self.quant_config = quant_config
498 self.logits_scale = config.logit_scale
vllm/model_executor/models/commandr.py
-3
@@ -429,9 +429,6 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
429 quant_config = vllm_config.quant_config
430
431 self.config = config
432 - # currently all existing command R models have `tie_word_embeddings`
433 - # enabled
434 - assert config.tie_word_embeddings
432
433 self.quant_config = quant_config
434 self.logits_processor = LogitsProcessor(
vllm/model_executor/models/ernie45_moe.py
+1 -4
@@ -629,8 +629,5 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe
629 return logits
630
631 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
632 - loader = AutoWeightsLoader(
633 - self,
634 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
635 - )
632 + loader = AutoWeightsLoader(self)
633 return loader.load_weights(weights)
vllm/model_executor/models/ernie_mtp.py
+1 -2
@@ -214,6 +214,5 @@ class ErnieMTP(nn.Module):
214 if any(k in name for k in ("mtp", "embed_tokens", "lm_head")):
215 yield name, weight
216
217 - skip_prefixes = ["lm_head"] if self.config.tie_word_embeddings else []
218 - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
217 + loader = AutoWeightsLoader(self)
218 return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/exaone.py
+1 -7
@@ -446,11 +446,5 @@ class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
446 return logits
447
448 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
449 - loader = AutoWeightsLoader(
450 - self,
451 - # With tie_word_embeddings, we can skip lm_head.weight
452 - # The weight might appear unnecessarily in the files if the model is
453 - # processed with quantization, LoRA, fine-tuning, etc.
454 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
455 - )
449 + loader = AutoWeightsLoader(self)
450 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/exaone4.py
+1 -7
@@ -446,11 +446,5 @@ class Exaone4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
446 return logits
447
448 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
449 - loader = AutoWeightsLoader(
450 - self,
451 - # With tie_word_embeddings, we can skip lm_head.weight
452 - # The weight might appear unnecessarily in the files if the model is
453 - # processed with quantization, LoRA, fine-tuning, etc.
454 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
455 - )
449 + loader = AutoWeightsLoader(self)
450 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/exaone_moe.py
+1 -6
@@ -560,12 +560,7 @@ class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
560 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
561 loader = AutoWeightsLoader(
562 self,
563 - # With tie_word_embeddings, we can skip lm_head.weight
564 - # The weight might appear unnecessarily in the files if the model is
565 - # processed with quantization, LoRA, fine-tuning, etc.
566 - skip_prefixes=(
567 - ["lm_head.", "mtp."] if self.config.tie_word_embeddings else ["mtp."]
568 - ),
563 + skip_prefixes=["mtp."],
564 # Skip loading extra parameters for GPTQ/modelopt models.
565 ignore_unexpected_suffixes=[
566 ".bias",
vllm/model_executor/models/fairseq2_llama.py
+1 -4
@@ -74,10 +74,7 @@ class Fairseq2LlamaForCausalLM(LlamaForCausalLM):
74
75 params = dict(self.named_parameters())
76
77 - loader = AutoWeightsLoader(
78 - self,
79 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
80 - )
77 + loader = AutoWeightsLoader(self)
78 return loader.load_weights(
79 self.reshape_fairseq2_weights(name, loaded_weight, params)
80 for name, loaded_weight in weights
vllm/model_executor/models/falcon.py
+1 -4
@@ -535,8 +535,5 @@ class FalconForCausalLM(nn.Module, SupportsPP):
535 return logits
536
537 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
538 - loader = AutoWeightsLoader(
539 - self,
540 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
541 - )
538 + loader = AutoWeightsLoader(self)
539 return loader.load_weights(weights)
vllm/model_executor/models/falcon_h1.py
+1 -4
@@ -643,8 +643,5 @@ class FalconH1ForCausalLM(
643 return logits
644
645 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
646 - loader = AutoWeightsLoader(
647 - self,
648 - skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
649 - )
646 + loader = AutoWeightsLoader(self)
647 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/gemma.py
+1 -6
@@ -346,8 +346,6 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
346 quant_config = vllm_config.quant_config
347
348 self.config = config
349 - # currently all existing Gemma models have `tie_word_embeddings` enabled
350 - assert config.tie_word_embeddings
349
350 self.quant_config = quant_config
351 self.model = GemmaModel(
@@ -381,8 +379,5 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
379 return logits
380
381 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
384 - loader = AutoWeightsLoader(
385 - self,
386 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
387 - )
382 + loader = AutoWeightsLoader(self)
383 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/gemma2.py
+1 -6
@@ -335,8 +335,6 @@ class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
335
336 super().__init__()
337 self.config = config
338 - # currently all existing Gemma models have `tie_word_embeddings` enabled
339 - assert config.tie_word_embeddings
338 self.quant_config = quant_config
339 self.model = Gemma2Model(
340 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
@@ -371,8 +369,5 @@ class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
369 return logits
370
371 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
374 - loader = AutoWeightsLoader(
375 - self,
376 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
377 - )
372 + loader = AutoWeightsLoader(self)
373 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/gemma3.py
+1 -4
@@ -442,8 +442,5 @@ class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
442 return logits
443
444 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
445 - loader = AutoWeightsLoader(
446 - self,
447 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
448 - )
445 + loader = AutoWeightsLoader(self)
446 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/gemma4.py
-3
@@ -1694,15 +1694,12 @@ class Gemma4ForCausalLM(
1694 yield name, weight
1695
1696 # Skip multimodal weights — handled by the multimodal wrapper.
1697 - # Also skip lm_head when weights are tied.
1697 skip = [
1698 "audio_tower.",
1699 "vision_tower.",
1700 "embed_audio.",
1701 "embed_vision.",
1702 ]
1704 - if self.config.tie_word_embeddings:
1705 - skip.append("lm_head.")
1703
1704 loader = AutoWeightsLoader(self, skip_substrs=skip)
1705 return loader.load_weights(_weight_iterator())
vllm/model_executor/models/glm4.py
+1 -2
@@ -292,11 +292,10 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
292 return logits
293
294 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
295 - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else []
295 # Skip the speculative (MTP) layers, which are loaded by the
296 # draft model instead.
297 num_nextn_layers = getattr(self.config, "num_nextn_predict_layers", 0)
299 - skip_prefixes += [
298 + skip_prefixes = [
299 f"model.layers.{self.config.num_hidden_layers + i}."
300 for i in range(num_nextn_layers)
301 ]
vllm/model_executor/models/gpt_oss.py
+1 -4
@@ -1242,8 +1242,5 @@ class GptOssForCausalLM(
1242 return logits
1243
1244 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1245 - loader = AutoWeightsLoader(
1246 - self,
1247 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
1248 - )
1245 + loader = AutoWeightsLoader(self)
1246 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/granite.py
+1 -6
@@ -467,10 +467,5 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
467 )
468
469 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
470 - # With tie_word_embeddings, we can skip lm_head.weight
471 - # The weight might appear unnecessarily in the files if the model is
472 - # processed with quantization, LoRA, fine-tuning, etc.
473 - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
474 -
475 - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
470 + loader = AutoWeightsLoader(self)
471 return loader.load_weights(weights)
vllm/model_executor/models/granitemoe.py
+1 -2
@@ -452,6 +452,5 @@ class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
452 )
453
454 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
455 - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
456 - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
455 + loader = AutoWeightsLoader(self)
456 return loader.load_weights(weights)
vllm/model_executor/models/granitemoeshared.py
+1 -2
@@ -289,6 +289,5 @@ class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
289 )
290
291 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
292 - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
293 - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
292 + loader = AutoWeightsLoader(self)
293 return loader.load_weights(weights)
vllm/model_executor/models/hrm_text.py
+1 -2
@@ -521,6 +521,5 @@ class HrmTextForCausalLM(nn.Module):
521 return self.logits_processor(self.lm_head, hidden_states)
522
523 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
524 - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
525 - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
524 + loader = AutoWeightsLoader(self)
525 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/hunyuan_v1.py
+1 -4
@@ -970,10 +970,7 @@ class HunyuanV1ModelBase(
970 )
971
972 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
973 - loader = AutoWeightsLoader(
974 - self,
975 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
976 - )
973 + loader = AutoWeightsLoader(self)
974 return loader.load_weights(weights)
975
976 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
vllm/model_executor/models/hunyuan_vision.py
+1 -4
@@ -1034,10 +1034,7 @@ class HunYuanVLForConditionalGeneration(
1034 return self.language_model.compute_logits(hidden_states)
1035
1036 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1037 - loader = AutoWeightsLoader(
1038 - self,
1039 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
1040 - )
1037 + loader = AutoWeightsLoader(self)
1038 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1039
1040 def get_mm_mapping(self) -> MultiModelKeys:
vllm/model_executor/models/hy_v3.py
+1 -4
@@ -710,10 +710,7 @@ class HYV3ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
710 continue
711 yield name, weight
712
713 - loader = AutoWeightsLoader(
714 - self,
715 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
716 - )
713 + loader = AutoWeightsLoader(self)
714 return loader.load_weights(_filter_weights(weights))
715
716 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
vllm/model_executor/models/hyperclovax.py
+1 -4
@@ -478,8 +478,5 @@ class HyperCLOVAXForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
478 self,
479 weights: Iterable[tuple[str, torch.Tensor]],
480 ) -> set[str]:
481 - loader = AutoWeightsLoader(
482 - self,
483 - skip_prefixes=["lm_head."] if self.config.tie_word_embeddings else None,
484 - )
481 + loader = AutoWeightsLoader(self)
482 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/internlm2.py
+1 -4
@@ -381,10 +381,7 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsQuant):
381 return logits
382
383 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
384 - loader = AutoWeightsLoader(
385 - self,
386 - skip_prefixes=(["output."] if self.config.tie_word_embeddings else None),
387 - )
384 + loader = AutoWeightsLoader(self)
385 return loader.load_weights(weights)
386
387
vllm/model_executor/models/iquest_loopcoder.py
+1 -4
@@ -575,8 +575,5 @@ class IQuestLoopCoderForCausalLM(nn.Module):
575 return logits
576
577 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
578 - loader = AutoWeightsLoader(
579 - self,
580 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
581 - )
578 + loader = AutoWeightsLoader(self)
579 return loader.load_weights(weights)
vllm/model_executor/models/jais2.py
+1 -4
@@ -437,8 +437,5 @@ class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
437 return logits
438
439 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
440 - loader = AutoWeightsLoader(
441 - self,
442 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
443 - )
440 + loader = AutoWeightsLoader(self)
441 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/laguna.py
+1 -4
@@ -763,8 +763,5 @@ class LagunaForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3):
763 return logits
764
765 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
766 - loader = AutoWeightsLoader(
767 - self,
768 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
769 - )
766 + loader = AutoWeightsLoader(self)
767 return loader.load_weights(weights)
vllm/model_executor/models/lfm2.py
+1 -4
@@ -508,8 +508,5 @@ class Lfm2ForCausalLM(
508 return logits
509
510 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
511 - loader = AutoWeightsLoader(
512 - self,
513 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
514 - )
511 + loader = AutoWeightsLoader(self)
512 return loader.load_weights(weights)
vllm/model_executor/models/lfm2_moe.py
+1 -4
@@ -671,8 +671,5 @@ class Lfm2MoeForCausalLM(
671 return logits
672
673 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
674 - loader = AutoWeightsLoader(
675 - self,
676 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
677 - )
674 + loader = AutoWeightsLoader(self)
675 return loader.load_weights(weights)
vllm/model_executor/models/llama.py
+1 -4
@@ -534,10 +534,7 @@ class LlamaForCausalLM(
534 return logits
535
536 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
537 - loader = AutoWeightsLoader(
538 - self,
539 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
540 - )
537 + loader = AutoWeightsLoader(self)
538 return loader.load_weights(weights)
539
540
vllm/model_executor/models/llama4.py
+1 -4
@@ -792,10 +792,7 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
792 )
793
794 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
795 - loader = AutoWeightsLoader(
796 - self,
797 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
798 - )
795 + loader = AutoWeightsLoader(self)
796 # Use a generator (not a list comprehension) so the weights iterator is
797 # consumed lazily by AutoWeightsLoader. Materializing it here would hold
798 # the entire language-model checkpoint in host memory at once, which can
vllm/model_executor/models/mimo.py
+1 -3
@@ -119,10 +119,8 @@ class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module):
119 )
120
121 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
122 - skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else []
122 # MTP layers are loaded by the draft model, not the main model.
124 - skip_prefixes.append("model.mtp_layers.")
125 - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
123 + loader = AutoWeightsLoader(self, skip_prefixes=["model.mtp_layers."])
124 return loader.load_weights(weights)
125
126 def compute_logits(
vllm/model_executor/models/minicpm.py
+1 -4
@@ -648,8 +648,5 @@ class MiniCPMForCausalLM(
648 return logits
649
650 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
651 - loader = AutoWeightsLoader(
652 - self,
653 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
654 - )
651 + loader = AutoWeightsLoader(self)
652 return loader.load_weights(weights)
vllm/model_executor/models/minicpm_eagle.py
+1 -4
@@ -387,8 +387,5 @@ class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle
387 process_eagle_weight(self, name)
388 return name, loaded_weight
389
390 - loader = AutoWeightsLoader(
391 - self,
392 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
393 - )
390 + loader = AutoWeightsLoader(self)
391 return loader.load_weights(map(transform, weights))
vllm/model_executor/models/mistral.py
+1 -4
@@ -277,10 +277,7 @@ class MistralForCausalLM(LlamaForCausalLM):
277 )
278
279 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
280 - loader = AutoWeightsLoader(
281 - self,
282 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
283 - )
280 + loader = AutoWeightsLoader(self)
281 return loader.load_weights(
282 self.maybe_remap_mistral(name, loaded_weight)
283 for name, loaded_weight in weights
vllm/model_executor/models/mpt.py
-1
@@ -279,7 +279,6 @@ class MPTForCausalLM(nn.Module, SupportsPP):
279 config = vllm_config.model_config.hf_config
280 quant_config = vllm_config.quant_config
281 self.config = config
282 - assert config.tie_word_embeddings
282 self.quant_config = quant_config
283
284 self.transformer = MPTModel(
vllm/model_executor/models/nemotron_nas.py
+1 -4
@@ -412,8 +412,5 @@ class DeciLMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, HasNoOps):
412 return logits
413
414 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
415 - loader = AutoWeightsLoader(
416 - self,
417 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
418 - )
415 + loader = AutoWeightsLoader(self)
416 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/olmo3.py
+1 -6
@@ -408,10 +408,5 @@ class Olmo3ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
408 return logits
409
410 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
411 - loader = AutoWeightsLoader(
412 - self,
413 - skip_prefixes=(
414 - ["lm_head.weight"] if self.config.tie_word_embeddings else None
415 - ),
416 - )
411 + loader = AutoWeightsLoader(self)
412 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/olmo_hybrid.py
+1 -6
@@ -474,10 +474,5 @@ class OlmoHybridForCausalLM(
474 return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
475
476 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
477 - loader = AutoWeightsLoader(
478 - self,
479 - skip_prefixes=(
480 - ["lm_head.weight"] if self.config.tie_word_embeddings else None
481 - ),
482 - )
477 + loader = AutoWeightsLoader(self)
478 return loader.load_weights(weights)
vllm/model_executor/models/openpangu.py
+1 -4
@@ -1153,10 +1153,7 @@ class OpenPanguModelBase(nn.Module, SupportsPP, SupportsLoRA):
1153 return logits
1154
1155 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1156 - loader = AutoWeightsLoader(
1157 - self,
1158 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
1159 - )
1156 + loader = AutoWeightsLoader(self)
1157 return loader.load_weights(weights)
1158
1159
vllm/model_executor/models/opt.py
+1 -6
@@ -384,10 +384,5 @@ class OPTForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
384 return logits
385
386 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
387 - loader = AutoWeightsLoader(
388 - self,
389 - skip_prefixes=(
390 - ["lm_head.weight"] if self.config.tie_word_embeddings else None
391 - ),
392 - )
387 + loader = AutoWeightsLoader(self)
388 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/plamo3.py
+1 -4
@@ -430,8 +430,5 @@ class Plamo3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
430 return logits
431
432 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
433 - loader = AutoWeightsLoader(
434 - self,
435 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
436 - )
433 + loader = AutoWeightsLoader(self)
434 return loader.load_weights(weights)
vllm/model_executor/models/qwen2.py
+1 -4
@@ -500,8 +500,5 @@ class Qwen2ForCausalLM(
500 return logits
501
502 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
503 - loader = AutoWeightsLoader(
504 - self,
505 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
506 - )
503 + loader = AutoWeightsLoader(self)
504 return loader.load_weights(weights)
vllm/model_executor/models/qwen3.py
+1 -4
@@ -335,8 +335,5 @@ class Qwen3ForCausalLM(
335 return logits
336
337 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
338 - loader = AutoWeightsLoader(
339 - self,
340 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
341 - )
338 + loader = AutoWeightsLoader(self)
339 return loader.load_weights(weights)
vllm/model_executor/models/rnj1.py
+1 -4
@@ -395,8 +395,5 @@ class Rnj1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
395 return logits
396
397 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
398 - loader = AutoWeightsLoader(
399 - self,
400 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
401 - )
398 + loader = AutoWeightsLoader(self)
399 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/sarvam.py
+1 -4
@@ -679,10 +679,7 @@ class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfE
679 self,
680 weights: Iterable[tuple[str, torch.Tensor]],
681 ) -> set[str]:
682 - loader = AutoWeightsLoader(
683 - self,
684 - skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
685 - )
682 + loader = AutoWeightsLoader(self)
683 return loader.load_weights(weights)
684
685
vllm/model_executor/models/seed_oss.py
+1 -4
@@ -428,8 +428,5 @@ class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
428 return logits
429
430 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
431 - loader = AutoWeightsLoader(
432 - self,
433 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
434 - )
431 + loader = AutoWeightsLoader(self)
432 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/telechat2.py
+1 -4
@@ -122,8 +122,5 @@ class TeleChat2ForCausalLM(LlamaForCausalLM):
122 return TeleChat2Model(vllm_config=vllm_config, prefix=prefix)
123
124 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
125 - loader = AutoWeightsLoader(
126 - self,
127 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
128 - )
125 + loader = AutoWeightsLoader(self)
126 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/transformers/causal.py
-4
@@ -40,11 +40,7 @@ class CausalMixin(VllmModelForTextGeneration):
40 vllm_config=vllm_config, prefix=prefix
41 )
42
43 - # Tell `Base.load_weights` to skip
44 - # `lm_head` if the model has tied word embeddings
43 tie_word_embeddings = self._get_tie_word_embeddings()
46 - if tie_word_embeddings:
47 - self.skip_prefixes.append("lm_head.")
44
45 if self.pp_group.is_last_rank:
46 self.lm_head = ParallelLMHead(
vllm/model_executor/models/utils.py
+54 -6
@@ -18,6 +18,7 @@ from vllm.distributed import (
18 get_tensor_model_parallel_world_size,
19 )
20 from vllm.logger import init_logger
21 +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
22 from vllm.model_executor.model_loader.reload import (
23 support_quantized_model_reload_from_hp_weights,
24 )
@@ -169,6 +170,20 @@ class WeightsMapper:
170 return replace(self, orig_to_new_stacked={})
171
172
173 +def _get_tied_embedding_params(module: nn.Module) -> dict[str, str]:
174 + """Map each tied word embedding qualname to the first name it aliases."""
175 + canonical = dict[int, str]()
176 + aliased = dict[str, str]()
177 + for prefix, submodule in module.named_modules(remove_duplicate=False):
178 + if not isinstance(submodule, VocabParallelEmbedding):
179 + continue
180 + for name, param in submodule.named_parameters(remove_duplicate=False):
181 + qualname = f"{prefix}.{name}" if prefix else name
182 + if (first_name := canonical.setdefault(id(param), qualname)) != qualname:
183 + aliased[qualname] = first_name
184 + return aliased
185 +
186 +
187 class AutoWeightsLoader:
188 """
189 Helper class to load weights into a [`torch.nn.Module`][]. It is able
@@ -213,6 +228,14 @@ class AutoWeightsLoader:
228 # update default skip_substrs
229 self.skip_substrs += self.ROTARY_EMBEDS_UNUSED_WEIGHTS
230
231 + # Weight tying makes two qualnames point at the same `nn.Parameter`
232 + # (e.g. `lm_head.weight` and `model.embed_tokens.weight`). Loading both
233 + # would write the same storage twice, so only the first name reached by
234 + # module traversal is loaded and the rest are skipped.
235 + self.aliased_params = _get_tied_embedding_params(module)
236 + self._skipped_aliases = dict[str, str]()
237 + self._loaded_params_are_complete = True
238 +
239 def _groupby_prefix(
240 self,
241 weights: Iterable[tuple[str, torch.Tensor]],
@@ -242,8 +265,10 @@ class AutoWeightsLoader:
265 return ".".join((prefix, rest))
266
267 def _can_skip(self, qualname: str) -> bool:
245 - return any(qualname.startswith(p) for p in self.skip_prefixes) or any(
246 - substr in qualname for substr in self.skip_substrs
268 + return (
269 + qualname in self.aliased_params
270 + or any(qualname.startswith(p) for p in self.skip_prefixes)
271 + or any(substr in qualname for substr in self.skip_substrs)
272 )
273
274 def _can_ignore_unexpected(self, qualname: str) -> bool:
@@ -332,6 +357,7 @@ class AutoWeightsLoader:
357 logger.warning(
358 "Unable to collect loaded parameters for module %s", module
359 )
360 + self._loaded_params_are_complete = False
361 else:
362 yield from map(
363 lambda x: self._get_qualname(base_prefix, x),
@@ -415,14 +441,36 @@ class AutoWeightsLoader:
441 self.ignore_unexpected_suffixes.extend(ignore_unexpected_suffixes)
442 if mapper is not None:
443 weights = mapper.apply(weights)
418 - # filter out weights with first-prefix/substr to skip in name
419 - weights = (
420 - (name, weight) for name, weight in weights if not self._can_skip(name)
421 - )
444 + weights = self._filter_skipped(weights)
445
446 autoloaded_weights = set(self._load_module("", self.module, weights))
447 + self._check_skipped_aliases(autoloaded_weights)
448 return autoloaded_weights
449
450 + def _filter_skipped(
451 + self, weights: Iterable[tuple[str, torch.Tensor]]
452 + ) -> Iterable[tuple[str, torch.Tensor]]:
453 + for name, weight in weights:
454 + if (canonical := self.aliased_params.get(name)) is not None:
455 + self._skipped_aliases[name] = canonical
456 + if self._can_skip(name):
457 + continue
458 +
459 + yield name, weight
460 +
461 + def _check_skipped_aliases(self, autoloaded_weights: set[str]) -> None:
462 + """Guard against skipping an alias whose canonical name never loads."""
463 + if not self._loaded_params_are_complete:
464 + return
465 + for alias, canonical in self._skipped_aliases.items():
466 + if canonical not in autoloaded_weights:
467 + raise ValueError(
468 + f"{alias!r} was skipped because it is tied to {canonical!r} "
469 + f"in {self.module._get_name()}, but {canonical!r} was not "
470 + "found in the checkpoint, so the tied weight is "
471 + "uninitialized."
472 + )
473 +
474
475 def maybe_fuse_shared_experts(
476 weights: Iterable[tuple[str, torch.Tensor]],
vllm/models/kimi_k3/amd/linear.py
+1 -4
@@ -1082,8 +1082,5 @@ class KimiLinearForCausalLM(
1082 return self.logits_processor(self.lm_head, hidden_states)
1083
1084 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1085 - loader = AutoWeightsLoader(
1086 - self,
1087 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
1088 - )
1085 + loader = AutoWeightsLoader(self)
1086 return loader.load_weights(weights)
vllm/models/kimi_k3/nvidia/model.py
+1 -4
@@ -1676,10 +1676,7 @@ class KimiLinearForCausalLM(
1676 return self.logits_processor(self.lm_head, hidden_states)
1677
1678 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1679 - loader = AutoWeightsLoader(
1680 - self,
1681 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
1682 - )
1679 + loader = AutoWeightsLoader(self)
1680 loaded = loader.load_weights(weights)
1681 self.model.finalize_mega_moe_weights()
1682 # The fused MultiHeadLatentAttention's process_weights_after_loading
vllm/transformers_utils/config.py
+21 -1
@@ -1112,7 +1112,9 @@ def try_get_safetensors_metadata(
1112
1113 try:
1114 return with_retry(
1115 - get_safetensors_metadata_partial, "Error retrieving safetensors"
1115 + get_safetensors_metadata_partial,
1116 + "Error retrieving safetensors",
1117 + fatal_errors=(huggingface_hub.errors.NotASafetensorsRepoError,),
1118 )
1119 except Exception:
1120 return None
@@ -1218,6 +1220,24 @@ def get_safetensors_params_metadata(
1220 return _read_safetensors_metadata_in_dir(Path(local_dir))
1221
1222
1223 +@cache
1224 +def checkpoint_has_lm_head(model: str, *, revision: str | None = None) -> bool | None:
1225 + """Whether the checkpoint contains an `lm_head` tensor of its own.
1226 +
1227 + Args:
1228 + model: Name or path of the model repository.
1229 + revision: The specific model version to use.
1230 +
1231 + Returns:
1232 + `None` if the checkpoint contents could not be determined, for example
1233 + because it is not stored as safetensors.
1234 + """
1235 + metadata = get_safetensors_params_metadata(model, revision=revision)
1236 + if not metadata:
1237 + return None
1238 + return any(name.endswith("lm_head.weight") for name in metadata)
1239 +
1240 +
1241 def _download_mistral_config_file(model, revision) -> dict:
1242 config_file_name = "params.json"
1243 config_dict = get_hf_file_to_dict(config_file_name, model, revision)
vllm/transformers_utils/repo_utils.py
+13
@@ -95,10 +95,23 @@ def with_retry(
95 log_msg: str,
96 max_retries: int = 2,
97 retry_delay: int = 2,
98 + fatal_errors: tuple[type[Exception], ...] = (),
99 ) -> _R:
100 + """Call `func`, retrying transient failures.
101 +
102 + Args:
103 + func: The call to make.
104 + log_msg: Prefix for the message logged when a call fails.
105 + max_retries: How many times to call `func` before giving up.
106 + retry_delay: Seconds to wait after the first failure, doubled each time.
107 + fatal_errors: Exceptions that are a definitive answer rather than a
108 + transient failure. These are raised immediately, without logging.
109 + """
110 for attempt in range(max_retries):
111 try:
112 return func()
113 + except fatal_errors:
114 + raise
115 except Exception as e:
116 if attempt == max_retries - 1:
117 logger.error("%s: %s", log_msg, e)