diff --git a/tests/model_executor/model_loader/test_weight_tying.py b/tests/model_executor/model_loader/test_weight_tying.py new file mode 100644 index 0000000000..8d68da5676 --- /dev/null +++ b/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 diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py index 8d47b44365..2b32a63ac6 100644 --- a/tests/models/test_utils.py +++ b/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() diff --git a/tests/test_config.py b/tests/test_config.py index 2d9c3b83f7..7ec18b3e40 100644 --- a/tests/test_config.py +++ b/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( diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index 42261feb53..35b5a697db 100644 --- a/tests/transformers_utils/test_config.py +++ b/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() diff --git a/tests/transformers_utils/test_repo_utils.py b/tests/transformers_utils/test_repo_utils.py index 36d0acccd6..4a32734c53 100644 --- a/tests/transformers_utils/test_repo_utils.py +++ b/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() diff --git a/vllm/config/model.py b/vllm/config/model.py index 56e533ea53..37d2b7a754 100644 --- a/vllm/config/model.py +++ b/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], diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a271a65413..ca009b9184 100644 --- a/vllm/config/vllm.py +++ b/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 diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index b7ccb78c04..d28051a548 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/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): diff --git a/vllm/model_executor/model_loader/weight_tying.py b/vllm/model_executor/model_loader/weight_tying.py new file mode 100644 index 0000000000..d0c309d2c4 --- /dev/null +++ b/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) diff --git a/vllm/model_executor/models/apertus.py b/vllm/model_executor/models/apertus.py index 74b17129c9..b904802a3a 100644 --- a/vllm/model_executor/models/apertus.py +++ b/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) diff --git a/vllm/model_executor/models/arcee.py b/vllm/model_executor/models/arcee.py index c32a903bba..81cf3c5e0f 100644 --- a/vllm/model_executor/models/arcee.py +++ b/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) diff --git a/vllm/model_executor/models/arctic.py b/vllm/model_executor/models/arctic.py index ccb3f1976c..555e58aa13 100644 --- a/vllm/model_executor/models/arctic.py +++ b/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) diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index a3da7e39e7..33749818b1 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/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)) diff --git a/vllm/model_executor/models/bloom.py b/vllm/model_executor/models/bloom.py index 1347b0e605..8fb6bb4aa4 100644 --- a/vllm/model_executor/models/bloom.py +++ b/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) diff --git a/vllm/model_executor/models/chameleon.py b/vllm/model_executor/models/chameleon.py index 03c9fe1179..517697669c 100644 --- a/vllm/model_executor/models/chameleon.py +++ b/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) diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index 3873b6ee00..9af1a3439b 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/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 diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py index 5481fb0abc..ca2543055d 100644 --- a/vllm/model_executor/models/commandr.py +++ b/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( diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py index 5cc351a6aa..c7f26b6dd4 100644 --- a/vllm/model_executor/models/ernie45_moe.py +++ b/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) diff --git a/vllm/model_executor/models/ernie_mtp.py b/vllm/model_executor/models/ernie_mtp.py index 88aadbdb73..0baf80e8cc 100644 --- a/vllm/model_executor/models/ernie_mtp.py +++ b/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) diff --git a/vllm/model_executor/models/exaone.py b/vllm/model_executor/models/exaone.py index 0742ae870b..fef05d386a 100644 --- a/vllm/model_executor/models/exaone.py +++ b/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) diff --git a/vllm/model_executor/models/exaone4.py b/vllm/model_executor/models/exaone4.py index 571aa3a426..31ea665e0a 100644 --- a/vllm/model_executor/models/exaone4.py +++ b/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) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index ef01330f0c..229a741dca 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/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", diff --git a/vllm/model_executor/models/fairseq2_llama.py b/vllm/model_executor/models/fairseq2_llama.py index e898034fbf..97afb760cf 100644 --- a/vllm/model_executor/models/fairseq2_llama.py +++ b/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 diff --git a/vllm/model_executor/models/falcon.py b/vllm/model_executor/models/falcon.py index cdc388e529..97e133a03c 100644 --- a/vllm/model_executor/models/falcon.py +++ b/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) diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py index 3c96d00c28..0dca542ff3 100644 --- a/vllm/model_executor/models/falcon_h1.py +++ b/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) diff --git a/vllm/model_executor/models/gemma.py b/vllm/model_executor/models/gemma.py index 949799fa65..0f1282fd96 100644 --- a/vllm/model_executor/models/gemma.py +++ b/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) diff --git a/vllm/model_executor/models/gemma2.py b/vllm/model_executor/models/gemma2.py index da5161ffa0..0bdf9b1392 100644 --- a/vllm/model_executor/models/gemma2.py +++ b/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) diff --git a/vllm/model_executor/models/gemma3.py b/vllm/model_executor/models/gemma3.py index 4f370fec14..e886c2ccdd 100644 --- a/vllm/model_executor/models/gemma3.py +++ b/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) diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index f0ac972b27..e8bb72ac37 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/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()) diff --git a/vllm/model_executor/models/glm4.py b/vllm/model_executor/models/glm4.py index 4c30e30008..5d95175ab0 100644 --- a/vllm/model_executor/models/glm4.py +++ b/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) ] diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 27982eb1f3..d5fbee384a 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/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) diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index df813fd1e9..7b8f42fdd1 100644 --- a/vllm/model_executor/models/granite.py +++ b/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) diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py index d775de50ca..3a9872a605 100644 --- a/vllm/model_executor/models/granitemoe.py +++ b/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) diff --git a/vllm/model_executor/models/granitemoeshared.py b/vllm/model_executor/models/granitemoeshared.py index ecd942da40..d77f58a5b3 100644 --- a/vllm/model_executor/models/granitemoeshared.py +++ b/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) diff --git a/vllm/model_executor/models/hrm_text.py b/vllm/model_executor/models/hrm_text.py index dd7ba93e7b..29825ece5a 100644 --- a/vllm/model_executor/models/hrm_text.py +++ b/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) diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py index b90d463627..cf1933bd7d 100644 --- a/vllm/model_executor/models/hunyuan_v1.py +++ b/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: diff --git a/vllm/model_executor/models/hunyuan_vision.py b/vllm/model_executor/models/hunyuan_vision.py index 74e4c1a6ce..252556e86c 100644 --- a/vllm/model_executor/models/hunyuan_vision.py +++ b/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: diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index 0248ac3d9c..8b77d4b1e8 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/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]]: diff --git a/vllm/model_executor/models/hyperclovax.py b/vllm/model_executor/models/hyperclovax.py index b6d66698c9..66c3b4a270 100644 --- a/vllm/model_executor/models/hyperclovax.py +++ b/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) diff --git a/vllm/model_executor/models/internlm2.py b/vllm/model_executor/models/internlm2.py index ffc750c7c5..4996723b3c 100644 --- a/vllm/model_executor/models/internlm2.py +++ b/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) diff --git a/vllm/model_executor/models/iquest_loopcoder.py b/vllm/model_executor/models/iquest_loopcoder.py index fdb90b962d..d3b670658a 100644 --- a/vllm/model_executor/models/iquest_loopcoder.py +++ b/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) diff --git a/vllm/model_executor/models/jais2.py b/vllm/model_executor/models/jais2.py index 95b8c3ee44..84254e1ee5 100644 --- a/vllm/model_executor/models/jais2.py +++ b/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) diff --git a/vllm/model_executor/models/laguna.py b/vllm/model_executor/models/laguna.py index 9431f9a0ec..115000588d 100644 --- a/vllm/model_executor/models/laguna.py +++ b/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) diff --git a/vllm/model_executor/models/lfm2.py b/vllm/model_executor/models/lfm2.py index 69a0e23be4..6f03150b84 100644 --- a/vllm/model_executor/models/lfm2.py +++ b/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) diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index 966a458dea..c2a21e52d1 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/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) diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py index cbf75593e0..e3c5272224 100644 --- a/vllm/model_executor/models/llama.py +++ b/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) diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index 9577c27f54..2a073d2dc4 100644 --- a/vllm/model_executor/models/llama4.py +++ b/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 diff --git a/vllm/model_executor/models/mimo.py b/vllm/model_executor/models/mimo.py index 56250ba2f4..7be514e990 100644 --- a/vllm/model_executor/models/mimo.py +++ b/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( diff --git a/vllm/model_executor/models/minicpm.py b/vllm/model_executor/models/minicpm.py index ee9b31f140..699e003592 100644 --- a/vllm/model_executor/models/minicpm.py +++ b/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) diff --git a/vllm/model_executor/models/minicpm_eagle.py b/vllm/model_executor/models/minicpm_eagle.py index 890d52961c..89ce01d611 100644 --- a/vllm/model_executor/models/minicpm_eagle.py +++ b/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)) diff --git a/vllm/model_executor/models/mistral.py b/vllm/model_executor/models/mistral.py index ce1332d0c9..86c4b171e5 100644 --- a/vllm/model_executor/models/mistral.py +++ b/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 diff --git a/vllm/model_executor/models/mpt.py b/vllm/model_executor/models/mpt.py index 8e509fbcb4..b8b4c5670e 100644 --- a/vllm/model_executor/models/mpt.py +++ b/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( diff --git a/vllm/model_executor/models/nemotron_nas.py b/vllm/model_executor/models/nemotron_nas.py index 5a5f0e7773..df3eb60dfb 100644 --- a/vllm/model_executor/models/nemotron_nas.py +++ b/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) diff --git a/vllm/model_executor/models/olmo3.py b/vllm/model_executor/models/olmo3.py index 9ac4e61cf4..a54dd5749e 100644 --- a/vllm/model_executor/models/olmo3.py +++ b/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) diff --git a/vllm/model_executor/models/olmo_hybrid.py b/vllm/model_executor/models/olmo_hybrid.py index 122a1c1c5b..9a24df5ec9 100644 --- a/vllm/model_executor/models/olmo_hybrid.py +++ b/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) diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 4a2030cfdd..60b5b301d5 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/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) diff --git a/vllm/model_executor/models/opt.py b/vllm/model_executor/models/opt.py index d78e51e3dd..1efe9eb8a8 100644 --- a/vllm/model_executor/models/opt.py +++ b/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) diff --git a/vllm/model_executor/models/plamo3.py b/vllm/model_executor/models/plamo3.py index 2ba38a7b1f..3f9553f64a 100644 --- a/vllm/model_executor/models/plamo3.py +++ b/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) diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py index 3ec8ac9319..3820f5e39b 100644 --- a/vllm/model_executor/models/qwen2.py +++ b/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) diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py index 7db73a749f..79434bfac9 100644 --- a/vllm/model_executor/models/qwen3.py +++ b/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) diff --git a/vllm/model_executor/models/rnj1.py b/vllm/model_executor/models/rnj1.py index 2bcd279198..b184e3fcad 100644 --- a/vllm/model_executor/models/rnj1.py +++ b/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) diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index 04590a2a91..9b5a784cef 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/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) diff --git a/vllm/model_executor/models/seed_oss.py b/vllm/model_executor/models/seed_oss.py index f3fa8b2d5d..df0ed7e72e 100644 --- a/vllm/model_executor/models/seed_oss.py +++ b/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) diff --git a/vllm/model_executor/models/telechat2.py b/vllm/model_executor/models/telechat2.py index 42fa6d6871..637c456dbe 100644 --- a/vllm/model_executor/models/telechat2.py +++ b/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) diff --git a/vllm/model_executor/models/transformers/causal.py b/vllm/model_executor/models/transformers/causal.py index a32e2b54ba..29e780557c 100644 --- a/vllm/model_executor/models/transformers/causal.py +++ b/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( diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 5e1236a7eb..2f1d6f3b5c 100644 --- a/vllm/model_executor/models/utils.py +++ b/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]], diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index 5cb916ef54..835c3be4db 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/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) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index f1166695fd..472272a357 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/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 diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 87271c2588..72fe24302e 100644 --- a/vllm/transformers_utils/config.py +++ b/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) diff --git a/vllm/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py index 4066edbbf0..ca20390aad 100644 --- a/vllm/transformers_utils/repo_utils.py +++ b/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)