| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | |
| 10 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 11 | ParallelLMHead, |
| 12 | VocabParallelEmbedding, |
| 13 | ) |
| 14 | from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings |
| 15 | |
| 16 | VOCAB_SIZE = 16 |
| 17 | HIDDEN_SIZE = 4 |
| 18 | |
| 19 | |
| 20 | class UntiedModel(nn.Module): |
| 21 | """Nests the head like a multimodal model, which is the harder case.""" |
| 22 | |
| 23 | def __init__(self): |
| 24 | super().__init__() |
| 25 | self.language_model = nn.Module() |
| 26 | self.language_model.model = nn.Module() |
| 27 | self.language_model.model.embed_tokens = VocabParallelEmbedding( |
| 28 | VOCAB_SIZE, HIDDEN_SIZE |
| 29 | ) |
| 30 | self.language_model.lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE) |
| 31 | self.embed_tokens.weight.data.fill_(1.0) |
| 32 | self.lm_head.weight.data.fill_(1.0) |
| 33 | |
| 34 | @property |
| 35 | def embed_tokens(self) -> VocabParallelEmbedding: |
| 36 | return self.language_model.model.embed_tokens |
| 37 | |
| 38 | @property |
| 39 | def lm_head(self) -> ParallelLMHead: |
| 40 | return self.language_model.lm_head |
| 41 | |
| 42 | |
| 43 | def make_model_config(untied_by_checkpoint=False): |
| 44 | return SimpleNamespace( |
| 45 | model="dummy-model", |
| 46 | word_embeddings_untied_by_checkpoint=untied_by_checkpoint, |
| 47 | ) |
| 48 | |
| 49 | |
| 50 | @pytest.mark.cpu_test |
| 51 | @pytest.mark.usefixtures("dist_init") |
| 52 | @pytest.mark.parametrize("identical", [True, False]) |
| 53 | def test_retie_only_when_identical(identical: bool): |
| 54 | """A redundant copy of a tied lm_head is shared again to reclaim memory.""" |
| 55 | model = UntiedModel() |
| 56 | if not identical: |
| 57 | model.lm_head.weight.data.fill_(2.0) |
| 58 | |
| 59 | maybe_retie_word_embeddings(model, make_model_config(untied_by_checkpoint=True)) |
| 60 | |
| 61 | assert (model.lm_head.weight is model.embed_tokens.weight) is identical |
| 62 | if not identical: |
| 63 | assert torch.all(model.lm_head.weight == 2.0) |
| 64 | |
| 65 | |
| 66 | @pytest.mark.cpu_test |
| 67 | @pytest.mark.usefixtures("dist_init") |
| 68 | def test_quantized_lm_head_is_left_alone(): |
| 69 | """A quantized head may store its weights packed under another name.""" |
| 70 | model = UntiedModel() |
| 71 | model.lm_head.quant_method = SimpleNamespace() |
| 72 | |
| 73 | maybe_retie_word_embeddings(model, make_model_config(untied_by_checkpoint=True)) |
| 74 | |
| 75 | assert model.lm_head.weight is not model.embed_tokens.weight |
| 76 | |
| 77 | |
| 78 | @pytest.mark.cpu_test |
| 79 | @pytest.mark.usefixtures("dist_init") |
| 80 | def test_no_retie_without_checkpoint_override(): |
| 81 | """Word embeddings the config genuinely unties are left alone.""" |
| 82 | model = UntiedModel() |
| 83 | |
| 84 | maybe_retie_word_embeddings(model, make_model_config()) |
| 85 | |
| 86 | assert model.lm_head.weight is not model.embed_tokens.weight |