| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from vllm.model_executor.model_loader.checkpoint_weight_patch import ( |
| 8 | CheckpointWeightPatch, |
| 9 | load_checkpoint_weight_patches, |
| 10 | ) |
| 11 | |
| 12 | pytestmark = pytest.mark.cpu_test |
| 13 | |
| 14 | |
| 15 | class _PackedTPModel(torch.nn.Module): |
| 16 | """Model that packs the TP rank 1 shards of Q and K into one parameter.""" |
| 17 | |
| 18 | def __init__(self): |
| 19 | super().__init__() |
| 20 | self.packed_weight = torch.nn.Parameter(torch.full((4,), -1.0)) |
| 21 | self.load_calls: list[list[str]] = [] |
| 22 | |
| 23 | def load_weights(self, weights): |
| 24 | weights = list(weights) |
| 25 | self.load_calls.append([name for name, _ in weights]) |
| 26 | loaded_names = set() |
| 27 | for name, loaded_weight in weights: |
| 28 | if name == "q_proj.weight": |
| 29 | destination = self.packed_weight.data[:2] |
| 30 | elif name == "k_proj.weight": |
| 31 | destination = self.packed_weight.data[2:] |
| 32 | else: |
| 33 | raise AssertionError(name) |
| 34 | # Slice each full checkpoint tensor and load its second half for TP rank 1. |
| 35 | destination.copy_(loaded_weight.narrow(0, 2, 2)) |
| 36 | loaded_names.add(name) |
| 37 | return loaded_names |
| 38 | |
| 39 | |
| 40 | def _make_patch( |
| 41 | name: str, |
| 42 | *, |
| 43 | values: list[float], |
| 44 | indices: list[int] | None = None, |
| 45 | ) -> CheckpointWeightPatch: |
| 46 | return CheckpointWeightPatch( |
| 47 | name=name, |
| 48 | shape=(4,), |
| 49 | dtype=torch.float32, |
| 50 | values=torch.tensor(values), |
| 51 | indices=None if indices is None else torch.tensor(indices, dtype=torch.int32), |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | def test_dense_and_sparse_patches_follow_packed_tp_loader(): |
| 56 | model = _PackedTPModel() |
| 57 | |
| 58 | # Seed the TP-local packed weights through the dense path. |
| 59 | dense_loaded = load_checkpoint_weight_patches( |
| 60 | model, |
| 61 | [ |
| 62 | _make_patch("q_proj.weight", values=[0.0, 1.0, 2.0, 3.0]), |
| 63 | _make_patch("k_proj.weight", values=[10.0, 11.0, 12.0, 13.0]), |
| 64 | ], |
| 65 | ) |
| 66 | assert dense_loaded == {"q_proj.weight", "k_proj.weight"} |
| 67 | assert torch.equal(model.packed_weight, torch.tensor([2.0, 3.0, 12.0, 13.0])) |
| 68 | |
| 69 | # Sparse indices address the full checkpoint tensor; index 0 is outside this shard. |
| 70 | model.load_calls.clear() |
| 71 | original_copy = torch.Tensor.copy_ |
| 72 | sparse_loaded = load_checkpoint_weight_patches( |
| 73 | model, |
| 74 | [ |
| 75 | _make_patch( |
| 76 | "q_proj.weight", |
| 77 | indices=[0, 3], |
| 78 | values=[100.0, 30.0], |
| 79 | ), |
| 80 | _make_patch("k_proj.weight", indices=[2], values=[20.0]), |
| 81 | _make_patch("q_proj.weight", indices=[2], values=[22.0]), |
| 82 | ], |
| 83 | ) |
| 84 | |
| 85 | assert sparse_loaded == {"q_proj.weight", "k_proj.weight"} |
| 86 | assert model.load_calls == [ |
| 87 | ["q_proj.weight", "k_proj.weight"], |
| 88 | ["q_proj.weight"], |
| 89 | ] |
| 90 | assert torch.equal(model.packed_weight, torch.tensor([22.0, 30.0, 20.0, 13.0])) |
| 91 | assert torch.Tensor.copy_ is original_copy |
| 92 | |
| 93 | # The empty result tells the helper that the loader intentionally made no write. |
| 94 | class NoLocalWeightModel(torch.nn.Module): |
| 95 | """Model that deliberately owns none of the supplied weights.""" |
| 96 | |
| 97 | def __init__(self): |
| 98 | super().__init__() |
| 99 | self.weight = torch.nn.Parameter(torch.tensor([1.0])) |
| 100 | |
| 101 | def load_weights(self, _weights): |
| 102 | return set() |
| 103 | |
| 104 | no_local_weight_model = NoLocalWeightModel() |
| 105 | no_local_weight_loaded = load_checkpoint_weight_patches( |
| 106 | no_local_weight_model, |
| 107 | [_make_patch("q_proj.weight", indices=[3], values=[30.0])], |
| 108 | ) |
| 109 | assert no_local_weight_loaded == set() |
| 110 | assert torch.equal(no_local_weight_model.weight, torch.tensor([1.0])) |