| 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.models.kimi_k3.nvidia.model import ( |
| 11 | KimiK3ForConditionalGeneration, |
| 12 | KimiLinearForCausalLM, |
| 13 | ) |
| 14 | |
| 15 | pytestmark = pytest.mark.cpu_test |
| 16 | |
| 17 | |
| 18 | class _FakeKimiLinearModel(nn.Module): |
| 19 | def __init__(self) -> None: |
| 20 | super().__init__() |
| 21 | self.tensor_a = nn.Parameter(torch.zeros(1)) |
| 22 | self.tensor_c = nn.Parameter(torch.zeros(1)) |
| 23 | self.finalized_values: list[tuple[float, float]] = [] |
| 24 | |
| 25 | def load_weights(self, weights): |
| 26 | params = dict(self.named_parameters()) |
| 27 | loaded = set() |
| 28 | for name, value in weights: |
| 29 | params[name].data.copy_(value) |
| 30 | loaded.add(name) |
| 31 | return loaded |
| 32 | |
| 33 | def finalize_mega_moe_weights(self) -> None: |
| 34 | self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item())) |
| 35 | # MegaMoE finalization replaces its original weight parameters. |
| 36 | self.tensor_a = None |
| 37 | self.tensor_c = None |
| 38 | |
| 39 | |
| 40 | def test_interleaved_composite_weights_finalize_kimi_once_after_loading() -> None: |
| 41 | language_model = object.__new__(KimiLinearForCausalLM) |
| 42 | nn.Module.__init__(language_model) |
| 43 | language_model.config = SimpleNamespace(tie_word_embeddings=False) |
| 44 | language_model.model = _FakeKimiLinearModel() |
| 45 | |
| 46 | model = object.__new__(KimiK3ForConditionalGeneration) |
| 47 | nn.Module.__init__(model) |
| 48 | model.language_model = language_model |
| 49 | model.vision_tower = nn.Module() |
| 50 | model.vision_tower.tensor_b = nn.Parameter(torch.zeros(1)) |
| 51 | |
| 52 | loaded = model.load_weights( |
| 53 | iter( |
| 54 | [ |
| 55 | ("language_model.model.tensor_a", torch.tensor([1.0])), |
| 56 | ("vision_tower.tensor_b", torch.tensor([2.0])), |
| 57 | ("language_model.model.tensor_c", torch.tensor([3.0])), |
| 58 | ] |
| 59 | ) |
| 60 | ) |
| 61 | |
| 62 | assert loaded == { |
| 63 | "language_model.model.tensor_a", |
| 64 | "vision_tower.tensor_b", |
| 65 | "language_model.model.tensor_c", |
| 66 | } |
| 67 | assert language_model.model.finalized_values == [] |
| 68 | |
| 69 | model.process_weights_after_loading() |
| 70 | |
| 71 | assert language_model.model.finalized_values == [(1.0, 3.0)] |
| 72 | assert model.vision_tower.tensor_b.item() == 2.0 |