| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Reconcile word embedding tying with what the checkpoint actually contains.""" |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | |
| 10 | from vllm.config import ModelConfig |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 13 | ParallelLMHead, |
| 14 | UnquantizedEmbeddingMethod, |
| 15 | VocabParallelEmbedding, |
| 16 | ) |
| 17 | |
| 18 | logger = init_logger(__name__) |
| 19 | |
| 20 | |
| 21 | @dataclass |
| 22 | class _UntiedLMHead: |
| 23 | name: str |
| 24 | lm_head: ParallelLMHead |
| 25 | embed_tokens: VocabParallelEmbedding |
| 26 | |
| 27 | @property |
| 28 | def weight_name(self) -> str: |
| 29 | return f"{self.name}.weight" |
| 30 | |
| 31 | def tie(self, model: nn.Module) -> None: |
| 32 | parent_name, _, attr = self.name.rpartition(".") |
| 33 | parent = model.get_submodule(parent_name) if parent_name else model |
| 34 | setattr(parent, attr, self.lm_head.tie_weights(self.embed_tokens)) |
| 35 | |
| 36 | |
| 37 | def _get_untied_lm_head(model: nn.Module) -> _UntiedLMHead | None: |
| 38 | """Locate an `lm_head` that could be tied to the input embeddings. |
| 39 | |
| 40 | Returns `None` unless the model has exactly one of each, so that models with |
| 41 | several heads (such as MTP) or with already tied weights are left alone. |
| 42 | Both are found by type rather than by name, because `get_input_embeddings` |
| 43 | takes different arguments on multimodal models and the `lm_head` of a |
| 44 | multimodal model is nested inside its language model. |
| 45 | |
| 46 | A quantized `lm_head` is also left alone. Its weights may be packed under |
| 47 | another name, and online quantization creates them after loading, so |
| 48 | neither their contents nor whether they were loaded can be established here. |
| 49 | Note that this is a property of the layer, not of the model: most quantized |
| 50 | models leave the `lm_head` and the input embeddings unquantized. |
| 51 | """ |
| 52 | heads = list[tuple[str, ParallelLMHead]]() |
| 53 | embeddings = list[VocabParallelEmbedding]() |
| 54 | for name, module in model.named_modules(): |
| 55 | if not isinstance(module, VocabParallelEmbedding): |
| 56 | continue |
| 57 | if not isinstance(module.quant_method, UnquantizedEmbeddingMethod): |
| 58 | continue |
| 59 | if isinstance(module, ParallelLMHead): |
| 60 | heads.append((name, module)) |
| 61 | else: |
| 62 | embeddings.append(module) |
| 63 | |
| 64 | if len(heads) != 1 or len(embeddings) != 1: |
| 65 | return None |
| 66 | |
| 67 | (name, lm_head), embed_tokens = heads[0], embeddings[0] |
| 68 | if lm_head.weight.shape != embed_tokens.weight.shape: |
| 69 | return None |
| 70 | return _UntiedLMHead(name, lm_head, embed_tokens) |
| 71 | |
| 72 | |
| 73 | def maybe_retie_word_embeddings(model: nn.Module, model_config: ModelConfig) -> None: |
| 74 | """Re-tie word embeddings that |
| 75 | [ModelConfig.maybe_untie_word_embeddings][vllm.config.ModelConfig.maybe_untie_word_embeddings] |
| 76 | untied, if the loaded `lm_head` turned out to be identical to the input embeddings |
| 77 | after all. |
| 78 | |
| 79 | Checkpoints produced by quantization or fine-tuning tooling often keep a redundant |
| 80 | copy of the tied `lm_head`. Sharing the storage again reclaims the memory it would |
| 81 | otherwise cost.""" |
| 82 | if not model_config.word_embeddings_untied_by_checkpoint: |
| 83 | return |
| 84 | if (untied := _get_untied_lm_head(model)) is None: |
| 85 | return |
| 86 | |
| 87 | # On device, torch.equal segfaults on ROCm when sleep mode is enabled |
| 88 | if not torch.equal(untied.lm_head.weight.cpu(), untied.embed_tokens.weight.cpu()): |
| 89 | logger.warning( |
| 90 | "The config for %s says the word embeddings are tied, but the checkpoint " |
| 91 | "contains a different %s, which has been used instead of tying. " |
| 92 | "Set `tie_word_embeddings=False` in the config to silence this warning.", |
| 93 | model_config.model, |
| 94 | untied.weight_name, |
| 95 | ) |
| 96 | return |
| 97 | |
| 98 | logger.debug("Re-tying %s, which is identical to the input embeddings", untied.name) |
| 99 | untied.tie(model) |