100 lines
3.8 KiB
Python
100 lines
3.8 KiB
Python
# 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)
|