Fix weight tying (#51665)

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

tests/model_executor/model_loader/test_weight_tying.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings
+
+VOCAB_SIZE = 16
+HIDDEN_SIZE = 4
+
+
+class UntiedModel(nn.Module):
+    """Nests the head like a multimodal model, which is the harder case."""
+
+    def __init__(self):
+        super().__init__()
+        self.language_model = nn.Module()
+        self.language_model.model = nn.Module()
+        self.language_model.model.embed_tokens = VocabParallelEmbedding(
+            VOCAB_SIZE, HIDDEN_SIZE
+        )
+        self.language_model.lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE)
+        self.embed_tokens.weight.data.fill_(1.0)
+        self.lm_head.weight.data.fill_(1.0)
+
+    @property
+    def embed_tokens(self) -> VocabParallelEmbedding:
+        return self.language_model.model.embed_tokens
+
+    @property
+    def lm_head(self) -> ParallelLMHead:
+        return self.language_model.lm_head
+
+
+def make_model_config(untied_by_checkpoint=False):
+    return SimpleNamespace(
+        model="dummy-model",
+        word_embeddings_untied_by_checkpoint=untied_by_checkpoint,
+    )
+
+
+@pytest.mark.cpu_test
+@pytest.mark.usefixtures("dist_init")
+@pytest.mark.parametrize("identical", [True, False])
+def test_retie_only_when_identical(identical: bool):
+    """A redundant copy of a tied lm_head is shared again to reclaim memory."""
+    model = UntiedModel()
+    if not identical:
+        model.lm_head.weight.data.fill_(2.0)
+
+    maybe_retie_word_embeddings(model, make_model_config(untied_by_checkpoint=True))
+
+    assert (model.lm_head.weight is model.embed_tokens.weight) is identical
+    if not identical:
+        assert torch.all(model.lm_head.weight == 2.0)
+
+
+@pytest.mark.cpu_test
+@pytest.mark.usefixtures("dist_init")
+def test_quantized_lm_head_is_left_alone():
+    """A quantized head may store its weights packed under another name."""
+    model = UntiedModel()
+    model.lm_head.quant_method = SimpleNamespace()
+
+    maybe_retie_word_embeddings(model, make_model_config(untied_by_checkpoint=True))
+
+    assert model.lm_head.weight is not model.embed_tokens.weight
+
+
+@pytest.mark.cpu_test
+@pytest.mark.usefixtures("dist_init")
+def test_no_retie_without_checkpoint_override():
+    """Word embeddings the config genuinely unties are left alone."""
+    model = UntiedModel()
+
+    maybe_retie_word_embeddings(model, make_model_config())
+
+    assert model.lm_head.weight is not model.embed_tokens.weight

tests/models/test_utils.py

@@ -4,6 +4,10 @@
 import pytest
 import torch
 
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     _merge_multimodal_embeddings,
@@ -165,6 +169,80 @@ def test_module_skip_substr():
     assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1
 
 
+VOCAB_SIZE = 16
+HIDDEN_SIZE = 2
+
+
+class ModuleWithTiedWeights(torch.nn.Module):
+    """Mimics how models tie `lm_head` to the input embeddings."""
+
+    def __init__(self, tie: bool):
+        super().__init__()
+        self.model = torch.nn.Module()
+        self.model.embed_tokens = VocabParallelEmbedding(VOCAB_SIZE, HIDDEN_SIZE)
+        self.lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE)
+        if tie:
+            self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
+
+
+def make_embedding_weights(value: float) -> torch.Tensor:
+    return torch.full((VOCAB_SIZE, HIDDEN_SIZE), value)
+
+
+@pytest.mark.cpu_test
+@pytest.mark.usefixtures("dist_init")
+@pytest.mark.parametrize("tie", [True, False])
+def test_module_skip_tied_weights(tie: bool):
+    """Tied weights must be loaded once, under the first of their names."""
+    mod = ModuleWithTiedWeights(tie)
+
+    weights = [
+        ("model.embed_tokens.weight", make_embedding_weights(1.0)),
+        ("lm_head.weight", make_embedding_weights(2.0)),
+    ]
+    loaded = AutoWeightsLoader(mod).load_weights(iter(weights))
+
+    if tie:
+        assert loaded == {"model.embed_tokens.weight"}
+        assert torch.all(mod.lm_head.weight[:VOCAB_SIZE] == 1.0)
+    else:
+        assert loaded == {"model.embed_tokens.weight", "lm_head.weight"}
+        assert torch.all(mod.lm_head.weight[:VOCAB_SIZE] == 2.0)
+
+
+@pytest.mark.cpu_test
+@pytest.mark.usefixtures("dist_init")
+def test_module_skip_tied_weights_without_canonical():
+    """Skipping a tied weight must not leave the shared weight uninitialized."""
+    mod = ModuleWithTiedWeights(tie=True)
+
+    weights = [("lm_head.weight", make_embedding_weights(2.0))]
+    with pytest.raises(ValueError, match="model.embed_tokens.weight"):
+        AutoWeightsLoader(mod).load_weights(iter(weights))
+
+
+class ModuleWithSharedParam(torch.nn.Module):
+    """Mimics an MoE router shared between the MLP and its fused experts."""
+
+    def __init__(self):
+        super().__init__()
+        self.experts = torch.nn.Module()
+        self.experts.gate = torch.nn.Linear(2, 2, bias=False)
+        self.gate = self.experts.gate
+
+
+@pytest.mark.cpu_test
+def test_module_load_shared_params_that_are_not_tied_embeddings():
+    """Only tied embeddings are skipped; other shared params must still load."""
+    mod = ModuleWithSharedParam()
+
+    weights = [("gate.weight", torch.Tensor([[1, 2], [3, 4]]))]
+    loaded = AutoWeightsLoader(mod).load_weights(iter(weights))
+
+    assert loaded == {"gate.weight"}
+    assert torch.all(mod.gate.weight == torch.Tensor([[1, 2], [3, 4]]))
+
+
 class raise_if_cuda_sync:
     def __enter__(self):
         self.previous_debug_mode = torch.cuda.get_sync_debug_mode()

tests/test_config.py

@@ -604,6 +604,45 @@ def test_with_hf_config_leaves_unknown_model_type_without_architectures(
     assert updated.model_config.hf_config.architectures is None
 
 
+@pytest.mark.parametrize(
+    "checkpoint_tensors,tied",
+    [
+        # The checkpoint has an lm_head of its own, so it must win over the config
+        (["model.embed_tokens.weight", "lm_head.weight"], False),
+        (["model.embed_tokens.weight"], True),
+        # Contents unknown (not safetensors), so the config must be left alone
+        ([], True),
+    ],
+)
+def test_maybe_untie_word_embeddings(tmp_path, checkpoint_tensors, tied):
+    import torch
+    from safetensors.torch import save_file
+
+    if checkpoint_tensors:
+        save_file(
+            {name: torch.zeros(2, 2) for name in checkpoint_tensors},
+            tmp_path / "model.safetensors",
+        )
+
+    text_config = SimpleNamespace(tie_word_embeddings=True)
+    model_config = SimpleNamespace(
+        model=str(tmp_path),
+        revision=None,
+        hf_config=SimpleNamespace(
+            tie_word_embeddings=True,
+            get_text_config=lambda: text_config,
+        ),
+        word_embeddings_untied_by_checkpoint=False,
+    )
+
+    ModelConfig.maybe_untie_word_embeddings(model_config)
+
+    # Both levels must agree, since different callers read different ones
+    assert model_config.hf_config.tie_word_embeddings is tied
+    assert text_config.tie_word_embeddings is tied
+    assert model_config.word_embeddings_untied_by_checkpoint is not tied
+
+
 def test_async_scheduling_with_pipeline_parallelism_is_allowed():
     cfg = VllmConfig(
         scheduler_config=SchedulerConfig(

tests/transformers_utils/test_config.py

@@ -8,14 +8,17 @@ only get the `eos_token_id` from the tokenizer as defined by
 
 from types import SimpleNamespace
 from typing import cast
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 from transformers import PretrainedConfig
 
 from vllm.config.model import ModelConfig
 from vllm.tokenizers import get_tokenizer
 from vllm.transformers_utils import config as config_module
-from vllm.transformers_utils.config import try_get_generation_config
+from vllm.transformers_utils.config import (
+    get_safetensors_params_metadata,
+    try_get_generation_config,
+)
 
 
 def test_get_llama3_eos_token():
@@ -77,3 +80,22 @@ def test_model_config_generation_fallback_forwards_code_revision():
         config_format="auto",
         token=None,
     )
+
+
+def test_safetensors_metadata_of_repo_without_safetensors():
+    """A repo storing its weights in another format is an answer, not a failure,
+    so it must not be retried."""
+    from huggingface_hub.errors import LocalEntryNotFoundError, NotASafetensorsRepoError
+
+    get_safetensors_metadata = MagicMock(
+        side_effect=NotASafetensorsRepoError("not a safetensors repo")
+    )
+    api = SimpleNamespace(
+        get_safetensors_metadata=get_safetensors_metadata,
+        snapshot_download=MagicMock(side_effect=LocalEntryNotFoundError("no cache")),
+    )
+
+    with patch.object(config_module, "hf_api", lambda: api):
+        assert get_safetensors_params_metadata("some/pytorch-only-model") == {}
+
+    get_safetensors_metadata.assert_called_once()

tests/transformers_utils/test_repo_utils.py

@@ -14,6 +14,7 @@ from vllm.transformers_utils.repo_utils import (
     get_hf_file_to_dict,
     is_mistral_model_repo,
     list_filtered_repo_files,
+    with_retry,
 )
 
 
@@ -185,3 +186,13 @@ def test_is_mistral_model_repo(files: list[str], expected_bool: bool):
             repo_type="model",
             token="token",
         )
+
+
+def test_with_retry_does_not_retry_fatal_errors():
+    """A definitive answer must not be retried, which costs a delay per attempt."""
+    func = MagicMock(side_effect=FileNotFoundError("no such file"))
+
+    with pytest.raises(FileNotFoundError):
+        with_retry(func, "Error", fatal_errors=(FileNotFoundError,))
+
+    func.assert_called_once()

vllm/config/model.py

@@ -30,6 +30,7 @@ from vllm.platforms import current_platform
 from vllm.tasks import PoolingTask, ScoreType, SupportedTask
 from vllm.transformers_utils.config import (
     ConfigFormat,
+    checkpoint_has_lm_head,
     get_config,
     get_hf_image_processor_config,
     get_hf_text_config,
@@ -185,6 +186,10 @@ class ModelConfig:
     """The Hugging Face config of the model."""
     hf_text_config: PretrainedConfig = field(init=False)
     """The Hugging Face config of the text model (same as hf_config for text models)."""
+    word_embeddings_untied_by_checkpoint: bool = field(default=False, init=False)
+    """Whether `tie_word_embeddings` was overridden to `False` because the checkpoint
+    contains an `lm_head` of its own. The two may still turn out to be identical, in
+    which case they are re-tied once the weights have been loaded."""
     hf_config_path: str | None = None
     """Name or path of the Hugging Face config to use. If unspecified, model
     name or path will be used."""
@@ -1109,6 +1114,34 @@ class ModelConfig:
     def _get_encoder_config(self) -> dict[str, Any] | None:
         return get_sentence_transformer_tokenizer_config(self.model, self.revision)
 
+    def maybe_untie_word_embeddings(self) -> None:
+        """Stop trusting `tie_word_embeddings` when the checkpoint disagrees.
+
+        A config may claim the word embeddings are tied while the checkpoint
+        ships an `lm_head` of its own. Tying regardless would silently discard
+        that tensor, so build the `lm_head` as if untied and let it load. The
+        two are compared once loaded, and re-tied if they turn out to match, by
+        [maybe_retie_word_embeddings][vllm.model_executor.model_loader.weight_tying.maybe_retie_word_embeddings].
+
+        Transformers makes the same decision in `PreTrainedModel.tie_weights`,
+        where it can compare the two tensors directly.
+        """
+        if not getattr(self.hf_config, "tie_word_embeddings", False):
+            return
+        if not checkpoint_has_lm_head(self.model, revision=self.revision):
+            return
+
+        logger.debug(
+            "The config for %s says the word embeddings are tied, but the checkpoint "
+            "contains an lm_head. Loading it to find out whether they really are tied.",
+            self.model,
+        )
+        # Both levels must be updated: `VllmConfig.with_hf_config` reads the top
+        # level, and the text config is what the language model itself sees.
+        self.hf_config.tie_word_embeddings = False
+        self.hf_config.get_text_config().tie_word_embeddings = False
+        self.word_embeddings_untied_by_checkpoint = True
+
     def _get_default_runner_type(
         self,
         architectures: list[str],

vllm/config/vllm.py

@@ -1170,6 +1170,16 @@ class VllmConfig:
                 self.model_config, self.load_config
             )
 
+        # "dummy" reads no weights at all, and the sharded formats read a vLLM
+        # state dict, which stores tied word embeddings under the lm_head only.
+        # Neither can tell us what the original checkpoint contained.
+        if self.model_config is not None and self.load_config.load_format not in (
+            "dummy",
+            "sharded_state",
+            "runai_streamer_sharded",
+        ):
+            self.model_config.maybe_untie_word_embeddings()
+
         if (
             self.quant_config is not None
             and self.model_config is not None

vllm/model_executor/model_loader/utils.py

@@ -24,6 +24,7 @@ from vllm.model_executor.model_loader.reload import (
     record_metadata_for_reloading,
     set_torchao_reload_attrs,
 )
+from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings
 from vllm.model_executor.models.interfaces import SupportsQuant
 from vllm.tracing import instrument
 from vllm.utils.mem_utils import release_device_memory_under_pressure
@@ -96,6 +97,10 @@ def initialize_model(
 def process_weights_after_loading(
     model: nn.Module, model_config: ModelConfig, target_device: torch.device
 ) -> None:
+    # Reclaim memory when an explicit lm_head has been
+    # loaded, but it is identical to the input embeddings.
+    maybe_retie_word_embeddings(model, model_config)
+
     for _, module in model.named_modules():
         quant_method = getattr(module, "quant_method", None)
         if isinstance(quant_method, QuantizeMethodBase):

vllm/model_executor/model_loader/weight_tying.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reconcile word embedding tying with what the checkpoint actually contains."""
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from vllm.config import ModelConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    UnquantizedEmbeddingMethod,
+    VocabParallelEmbedding,
+)
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class _UntiedLMHead:
+    name: str
+    lm_head: ParallelLMHead
+    embed_tokens: VocabParallelEmbedding
+
+    @property
+    def weight_name(self) -> str:
+        return f"{self.name}.weight"
+
+    def tie(self, model: nn.Module) -> None:
+        parent_name, _, attr = self.name.rpartition(".")
+        parent = model.get_submodule(parent_name) if parent_name else model
+        setattr(parent, attr, self.lm_head.tie_weights(self.embed_tokens))
+
+
+def _get_untied_lm_head(model: nn.Module) -> _UntiedLMHead | None:
+    """Locate an `lm_head` that could be tied to the input embeddings.
+
+    Returns `None` unless the model has exactly one of each, so that models with
+    several heads (such as MTP) or with already tied weights are left alone.
+    Both are found by type rather than by name, because `get_input_embeddings`
+    takes different arguments on multimodal models and the `lm_head` of a
+    multimodal model is nested inside its language model.
+
+    A quantized `lm_head` is also left alone. Its weights may be packed under
+    another name, and online quantization creates them after loading, so
+    neither their contents nor whether they were loaded can be established here.
+    Note that this is a property of the layer, not of the model: most quantized
+    models leave the `lm_head` and the input embeddings unquantized.
+    """
+    heads = list[tuple[str, ParallelLMHead]]()
+    embeddings = list[VocabParallelEmbedding]()
+    for name, module in model.named_modules():
+        if not isinstance(module, VocabParallelEmbedding):
+            continue
+        if not isinstance(module.quant_method, UnquantizedEmbeddingMethod):
+            continue
+        if isinstance(module, ParallelLMHead):
+            heads.append((name, module))
+        else:
+            embeddings.append(module)
+
+    if len(heads) != 1 or len(embeddings) != 1:
+        return None
+
+    (name, lm_head), embed_tokens = heads[0], embeddings[0]
+    if lm_head.weight.shape != embed_tokens.weight.shape:
+        return None
+    return _UntiedLMHead(name, lm_head, embed_tokens)
+
+
+def maybe_retie_word_embeddings(model: nn.Module, model_config: ModelConfig) -> None:
+    """Re-tie word embeddings that
+    [ModelConfig.maybe_untie_word_embeddings][vllm.config.ModelConfig.maybe_untie_word_embeddings]
+    untied, if the loaded `lm_head` turned out to be identical to the input embeddings
+    after all.
+
+    Checkpoints produced by quantization or fine-tuning tooling often keep a redundant
+    copy of the tied `lm_head`. Sharing the storage again reclaims the memory it would
+    otherwise cost."""
+    if not model_config.word_embeddings_untied_by_checkpoint:
+        return
+    if (untied := _get_untied_lm_head(model)) is None:
+        return
+
+    # On device, torch.equal segfaults on ROCm when sleep mode is enabled
+    if not torch.equal(untied.lm_head.weight.cpu(), untied.embed_tokens.weight.cpu()):
+        logger.warning(
+            "The config for %s says the word embeddings are tied, but the checkpoint "
+            "contains a different %s, which has been used instead of tying. "
+            "Set `tie_word_embeddings=False` in the config to silence this warning.",
+            model_config.model,
+            untied.weight_name,
+        )
+        return
+
+    logger.debug("Re-tying %s, which is identical to the input embeddings", untied.name)
+    untied.tie(model)

vllm/model_executor/models/apertus.py

@@ -491,8 +491,5 @@ class ApertusForCausalLM(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/arcee.py

@@ -356,11 +356,7 @@ class ArceeForCausalLM(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights into the model (delegates to inner model and handles
         tied embeddings)."""
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-            skip_substrs=["gate_proj"],
-        )
+        loader = AutoWeightsLoader(self, skip_substrs=["gate_proj"])
         # AutoWeightLoader handles weight name remapping, including fusing
         # separate q_proj, k_proj, v_proj into qkv_proj
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/arctic.py

@@ -589,8 +589,5 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/bailing_moe.py

@@ -557,10 +557,7 @@ class BailingMoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
             yield name, loaded_weight
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(self._normalize_lm_head(weights))
 
 

vllm/model_executor/models/bloom.py

@@ -360,7 +360,7 @@ class BloomForCausalLM(nn.Module, SupportsPP, SupportsQuant):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(self, skip_prefixes=["lm_head.weight"])
+        loader = AutoWeightsLoader(self)
         weights = _add_transformer_prefix(weights)
         return loader.load_weights(weights)
 

vllm/model_executor/models/chameleon.py

@@ -1059,6 +1059,5 @@ class ChameleonForConditionalGeneration(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
-        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/cohere2_moe.py

@@ -493,7 +493,6 @@ class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
         config = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
         self.config = config
-        assert getattr(config, "tie_word_embeddings", True)
         self.unpadded_vocab_size = config.vocab_size
         self.quant_config = quant_config
         self.logits_scale = config.logit_scale

vllm/model_executor/models/commandr.py

@@ -429,9 +429,6 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
         quant_config = vllm_config.quant_config
 
         self.config = config
-        # currently all existing command R models have `tie_word_embeddings`
-        # enabled
-        assert config.tie_word_embeddings
 
         self.quant_config = quant_config
         self.logits_processor = LogitsProcessor(

vllm/model_executor/models/ernie45_moe.py

@@ -629,8 +629,5 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/ernie_mtp.py

@@ -214,6 +214,5 @@ class ErnieMTP(nn.Module):
                 if any(k in name for k in ("mtp", "embed_tokens", "lm_head")):
                     yield name, weight
 
-        skip_prefixes = ["lm_head"] if self.config.tie_word_embeddings else []
-        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/exaone.py

@@ -446,11 +446,5 @@ class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            # With tie_word_embeddings, we can skip lm_head.weight
-            # The weight might appear unnecessarily in the files if the model is
-            # processed with quantization, LoRA, fine-tuning, etc.
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/exaone4.py

@@ -446,11 +446,5 @@ class Exaone4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            # With tie_word_embeddings, we can skip lm_head.weight
-            # The weight might appear unnecessarily in the files if the model is
-            # processed with quantization, LoRA, fine-tuning, etc.
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/exaone_moe.py

@@ -560,12 +560,7 @@ class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(
             self,
-            # With tie_word_embeddings, we can skip lm_head.weight
-            # The weight might appear unnecessarily in the files if the model is
-            # processed with quantization, LoRA, fine-tuning, etc.
-            skip_prefixes=(
-                ["lm_head.", "mtp."] if self.config.tie_word_embeddings else ["mtp."]
-            ),
+            skip_prefixes=["mtp."],
             # Skip loading extra parameters for GPTQ/modelopt models.
             ignore_unexpected_suffixes=[
                 ".bias",

vllm/model_executor/models/fairseq2_llama.py

@@ -74,10 +74,7 @@ class Fairseq2LlamaForCausalLM(LlamaForCausalLM):
 
         params = dict(self.named_parameters())
 
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(
             self.reshape_fairseq2_weights(name, loaded_weight, params)
             for name, loaded_weight in weights

vllm/model_executor/models/falcon.py

@@ -535,8 +535,5 @@ class FalconForCausalLM(nn.Module, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/falcon_h1.py

@@ -643,8 +643,5 @@ class FalconH1ForCausalLM(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/gemma.py

@@ -346,8 +346,6 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
         quant_config = vllm_config.quant_config
 
         self.config = config
-        # currently all existing Gemma models have `tie_word_embeddings` enabled
-        assert config.tie_word_embeddings
 
         self.quant_config = quant_config
         self.model = GemmaModel(
@@ -381,8 +379,5 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/gemma2.py

@@ -335,8 +335,6 @@ class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
 
         super().__init__()
         self.config = config
-        # currently all existing Gemma models have `tie_word_embeddings` enabled
-        assert config.tie_word_embeddings
         self.quant_config = quant_config
         self.model = Gemma2Model(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
@@ -371,8 +369,5 @@ class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/gemma3.py

@@ -442,8 +442,5 @@ class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/gemma4.py

@@ -1694,15 +1694,12 @@ class Gemma4ForCausalLM(
                 yield name, weight
 
         # Skip multimodal weights — handled by the multimodal wrapper.
-        # Also skip lm_head when weights are tied.
         skip = [
             "audio_tower.",
             "vision_tower.",
             "embed_audio.",
             "embed_vision.",
         ]
-        if self.config.tie_word_embeddings:
-            skip.append("lm_head.")
 
         loader = AutoWeightsLoader(self, skip_substrs=skip)
         return loader.load_weights(_weight_iterator())

vllm/model_executor/models/glm4.py

@@ -292,11 +292,10 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else []
         # Skip the speculative (MTP) layers, which are loaded by the
         # draft model instead.
         num_nextn_layers = getattr(self.config, "num_nextn_predict_layers", 0)
-        skip_prefixes += [
+        skip_prefixes = [
             f"model.layers.{self.config.num_hidden_layers + i}."
             for i in range(num_nextn_layers)
         ]

vllm/model_executor/models/gpt_oss.py

@@ -1242,8 +1242,5 @@ class GptOssForCausalLM(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/granite.py

@@ -467,10 +467,5 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # With tie_word_embeddings, we can skip lm_head.weight
-        # The weight might appear unnecessarily in the files if the model is
-        # processed with quantization, LoRA, fine-tuning, etc.
-        skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
-
-        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/granitemoe.py

@@ -452,6 +452,5 @@ class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
-        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/granitemoeshared.py

@@ -289,6 +289,5 @@ class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
-        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/hrm_text.py

@@ -521,6 +521,5 @@ class HrmTextForCausalLM(nn.Module):
         return self.logits_processor(self.lm_head, hidden_states)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
-        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/hunyuan_v1.py

@@ -970,10 +970,7 @@ class HunyuanV1ModelBase(
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

vllm/model_executor/models/hunyuan_vision.py

@@ -1034,10 +1034,7 @@ class HunYuanVLForConditionalGeneration(
         return self.language_model.compute_logits(hidden_states)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:

vllm/model_executor/models/hy_v3.py

@@ -710,10 +710,7 @@ class HYV3ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
                     continue
                 yield name, weight
 
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(_filter_weights(weights))
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

vllm/model_executor/models/hyperclovax.py

@@ -478,8 +478,5 @@ class HyperCLOVAXForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         self,
         weights: Iterable[tuple[str, torch.Tensor]],
     ) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=["lm_head."] if self.config.tie_word_embeddings else None,
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/internlm2.py

@@ -381,10 +381,7 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsQuant):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["output."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
 
 

vllm/model_executor/models/iquest_loopcoder.py

@@ -575,8 +575,5 @@ class IQuestLoopCoderForCausalLM(nn.Module):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/jais2.py

@@ -437,8 +437,5 @@ class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/laguna.py

@@ -763,8 +763,5 @@ class LagunaForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/lfm2.py

@@ -508,8 +508,5 @@ class Lfm2ForCausalLM(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/lfm2_moe.py

@@ -671,8 +671,5 @@ class Lfm2MoeForCausalLM(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/llama.py

@@ -534,10 +534,7 @@ class LlamaForCausalLM(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
 
 

vllm/model_executor/models/llama4.py

@@ -792,10 +792,7 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         # Use a generator (not a list comprehension) so the weights iterator is
         # consumed lazily by AutoWeightsLoader. Materializing it here would hold
         # the entire language-model checkpoint in host memory at once, which can

vllm/model_executor/models/mimo.py

@@ -119,10 +119,8 @@ class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else []
         # MTP layers are loaded by the draft model, not the main model.
-        skip_prefixes.append("model.mtp_layers.")
-        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        loader = AutoWeightsLoader(self, skip_prefixes=["model.mtp_layers."])
         return loader.load_weights(weights)
 
     def compute_logits(

vllm/model_executor/models/minicpm.py

@@ -648,8 +648,5 @@ class MiniCPMForCausalLM(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/minicpm_eagle.py

@@ -387,8 +387,5 @@ class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle
             process_eagle_weight(self, name)
             return name, loaded_weight
 
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(map(transform, weights))

vllm/model_executor/models/mistral.py

@@ -277,10 +277,7 @@ class MistralForCausalLM(LlamaForCausalLM):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(
             self.maybe_remap_mistral(name, loaded_weight)
             for name, loaded_weight in weights

vllm/model_executor/models/mpt.py

@@ -279,7 +279,6 @@ class MPTForCausalLM(nn.Module, SupportsPP):
         config = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
         self.config = config
-        assert config.tie_word_embeddings
         self.quant_config = quant_config
 
         self.transformer = MPTModel(

vllm/model_executor/models/nemotron_nas.py

@@ -412,8 +412,5 @@ class DeciLMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, HasNoOps):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/olmo3.py

@@ -408,10 +408,5 @@ class Olmo3ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(
-                ["lm_head.weight"] if self.config.tie_word_embeddings else None
-            ),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/olmo_hybrid.py

@@ -474,10 +474,5 @@ class OlmoHybridForCausalLM(
         return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(
-                ["lm_head.weight"] if self.config.tie_word_embeddings else None
-            ),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/openpangu.py

@@ -1153,10 +1153,7 @@ class OpenPanguModelBase(nn.Module, SupportsPP, SupportsLoRA):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
 
 

vllm/model_executor/models/opt.py

@@ -384,10 +384,5 @@ class OPTForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(
-                ["lm_head.weight"] if self.config.tie_word_embeddings else None
-            ),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/plamo3.py

@@ -430,8 +430,5 @@ class Plamo3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/qwen2.py

@@ -500,8 +500,5 @@ class Qwen2ForCausalLM(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/qwen3.py

@@ -335,8 +335,5 @@ class Qwen3ForCausalLM(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/rnj1.py

@@ -395,8 +395,5 @@ class Rnj1ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/sarvam.py

@@ -679,10 +679,7 @@ class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfE
         self,
         weights: Iterable[tuple[str, torch.Tensor]],
     ) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
 
 

vllm/model_executor/models/seed_oss.py

@@ -428,8 +428,5 @@ class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/telechat2.py

@@ -122,8 +122,5 @@ class TeleChat2ForCausalLM(LlamaForCausalLM):
         return TeleChat2Model(vllm_config=vllm_config, prefix=prefix)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/transformers/causal.py

@@ -40,11 +40,7 @@ class CausalMixin(VllmModelForTextGeneration):
             vllm_config=vllm_config, prefix=prefix
         )
 
-        # Tell `Base.load_weights` to skip
-        # `lm_head` if the model has tied word embeddings
         tie_word_embeddings = self._get_tie_word_embeddings()
-        if tie_word_embeddings:
-            self.skip_prefixes.append("lm_head.")
 
         if self.pp_group.is_last_rank:
             self.lm_head = ParallelLMHead(

vllm/model_executor/models/utils.py

@@ -18,6 +18,7 @@ from vllm.distributed import (
     get_tensor_model_parallel_world_size,
 )
 from vllm.logger import init_logger
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.model_executor.model_loader.reload import (
     support_quantized_model_reload_from_hp_weights,
 )
@@ -169,6 +170,20 @@ class WeightsMapper:
         return replace(self, orig_to_new_stacked={})
 
 
+def _get_tied_embedding_params(module: nn.Module) -> dict[str, str]:
+    """Map each tied word embedding qualname to the first name it aliases."""
+    canonical = dict[int, str]()
+    aliased = dict[str, str]()
+    for prefix, submodule in module.named_modules(remove_duplicate=False):
+        if not isinstance(submodule, VocabParallelEmbedding):
+            continue
+        for name, param in submodule.named_parameters(remove_duplicate=False):
+            qualname = f"{prefix}.{name}" if prefix else name
+            if (first_name := canonical.setdefault(id(param), qualname)) != qualname:
+                aliased[qualname] = first_name
+    return aliased
+
+
 class AutoWeightsLoader:
     """
     Helper class to load weights into a [`torch.nn.Module`][]. It is able
@@ -213,6 +228,14 @@ class AutoWeightsLoader:
         # update default skip_substrs
         self.skip_substrs += self.ROTARY_EMBEDS_UNUSED_WEIGHTS
 
+        # Weight tying makes two qualnames point at the same `nn.Parameter`
+        # (e.g. `lm_head.weight` and `model.embed_tokens.weight`). Loading both
+        # would write the same storage twice, so only the first name reached by
+        # module traversal is loaded and the rest are skipped.
+        self.aliased_params = _get_tied_embedding_params(module)
+        self._skipped_aliases = dict[str, str]()
+        self._loaded_params_are_complete = True
+
     def _groupby_prefix(
         self,
         weights: Iterable[tuple[str, torch.Tensor]],
@@ -242,8 +265,10 @@ class AutoWeightsLoader:
         return ".".join((prefix, rest))
 
     def _can_skip(self, qualname: str) -> bool:
-        return any(qualname.startswith(p) for p in self.skip_prefixes) or any(
-            substr in qualname for substr in self.skip_substrs
+        return (
+            qualname in self.aliased_params
+            or any(qualname.startswith(p) for p in self.skip_prefixes)
+            or any(substr in qualname for substr in self.skip_substrs)
         )
 
     def _can_ignore_unexpected(self, qualname: str) -> bool:
@@ -332,6 +357,7 @@ class AutoWeightsLoader:
                     logger.warning(
                         "Unable to collect loaded parameters for module %s", module
                     )
+                    self._loaded_params_are_complete = False
                 else:
                     yield from map(
                         lambda x: self._get_qualname(base_prefix, x),
@@ -415,14 +441,36 @@ class AutoWeightsLoader:
             self.ignore_unexpected_suffixes.extend(ignore_unexpected_suffixes)
         if mapper is not None:
             weights = mapper.apply(weights)
-        # filter out weights with first-prefix/substr to skip in name
-        weights = (
-            (name, weight) for name, weight in weights if not self._can_skip(name)
-        )
+        weights = self._filter_skipped(weights)
 
         autoloaded_weights = set(self._load_module("", self.module, weights))
+        self._check_skipped_aliases(autoloaded_weights)
         return autoloaded_weights
 
+    def _filter_skipped(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[tuple[str, torch.Tensor]]:
+        for name, weight in weights:
+            if (canonical := self.aliased_params.get(name)) is not None:
+                self._skipped_aliases[name] = canonical
+            if self._can_skip(name):
+                continue
+
+            yield name, weight
+
+    def _check_skipped_aliases(self, autoloaded_weights: set[str]) -> None:
+        """Guard against skipping an alias whose canonical name never loads."""
+        if not self._loaded_params_are_complete:
+            return
+        for alias, canonical in self._skipped_aliases.items():
+            if canonical not in autoloaded_weights:
+                raise ValueError(
+                    f"{alias!r} was skipped because it is tied to {canonical!r} "
+                    f"in {self.module._get_name()}, but {canonical!r} was not "
+                    "found in the checkpoint, so the tied weight is "
+                    "uninitialized."
+                )
+
 
 def maybe_fuse_shared_experts(
     weights: Iterable[tuple[str, torch.Tensor]],

vllm/models/kimi_k3/amd/linear.py

@@ -1082,8 +1082,5 @@ class KimiLinearForCausalLM(
         return self.logits_processor(self.lm_head, hidden_states)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/models/kimi_k3/nvidia/model.py

@@ -1676,10 +1676,7 @@ class KimiLinearForCausalLM(
         return self.logits_processor(self.lm_head, hidden_states)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
+        loader = AutoWeightsLoader(self)
         loaded = loader.load_weights(weights)
         self.model.finalize_mega_moe_weights()
         # The fused MultiHeadLatentAttention's process_weights_after_loading

vllm/transformers_utils/config.py

@@ -1112,7 +1112,9 @@ def try_get_safetensors_metadata(
 
     try:
         return with_retry(
-            get_safetensors_metadata_partial, "Error retrieving safetensors"
+            get_safetensors_metadata_partial,
+            "Error retrieving safetensors",
+            fatal_errors=(huggingface_hub.errors.NotASafetensorsRepoError,),
         )
     except Exception:
         return None
@@ -1218,6 +1220,24 @@ def get_safetensors_params_metadata(
     return _read_safetensors_metadata_in_dir(Path(local_dir))
 
 
+@cache
+def checkpoint_has_lm_head(model: str, *, revision: str | None = None) -> bool | None:
+    """Whether the checkpoint contains an `lm_head` tensor of its own.
+
+    Args:
+        model: Name or path of the model repository.
+        revision: The specific model version to use.
+
+    Returns:
+        `None` if the checkpoint contents could not be determined, for example
+        because it is not stored as safetensors.
+    """
+    metadata = get_safetensors_params_metadata(model, revision=revision)
+    if not metadata:
+        return None
+    return any(name.endswith("lm_head.weight") for name in metadata)
+
+
 def _download_mistral_config_file(model, revision) -> dict:
     config_file_name = "params.json"
     config_dict = get_hf_file_to_dict(config_file_name, model, revision)

vllm/transformers_utils/repo_utils.py

@@ -95,10 +95,23 @@ def with_retry(
     log_msg: str,
     max_retries: int = 2,
     retry_delay: int = 2,
+    fatal_errors: tuple[type[Exception], ...] = (),
 ) -> _R:
+    """Call `func`, retrying transient failures.
+
+    Args:
+        func: The call to make.
+        log_msg: Prefix for the message logged when a call fails.
+        max_retries: How many times to call `func` before giving up.
+        retry_delay: Seconds to wait after the first failure, doubled each time.
+        fatal_errors: Exceptions that are a definitive answer rather than a
+            transient failure. These are raised immediately, without logging.
+    """
     for attempt in range(max_retries):
         try:
             return func()
+        except fatal_errors:
+            raise
         except Exception as e:
             if attempt == max_retries - 1:
                 logger.error("%s: %s", log_msg, e)