Skip to content

vllm.model_executor.model_loader.weight_tying

Reconcile word embedding tying with what the checkpoint actually contains.

Functions:

_get_untied_lm_head(model)

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.

Source code in vllm/model_executor/model_loader/weight_tying.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
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)

maybe_retie_word_embeddings(model, model_config)

Re-tie word embeddings that 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.

Source code in vllm/model_executor/model_loader/weight_tying.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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)