| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import regex as re |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 9 | ParallelLMHead, |
| 10 | VocabParallelEmbedding, |
| 11 | ) |
| 12 | from vllm.model_executor.models.utils import ( |
| 13 | AutoWeightsLoader, |
| 14 | WeightsMapper, |
| 15 | _merge_multimodal_embeddings, |
| 16 | ) |
| 17 | from vllm.platforms import current_platform |
| 18 | |
| 19 | DEVICE_TYPE = current_platform.device_type |
| 20 | |
| 21 | |
| 22 | class ModuleWithBatchNorm(torch.nn.Module): |
| 23 | def __init__(self): |
| 24 | super().__init__() |
| 25 | self.bn = torch.nn.BatchNorm1d(2) |
| 26 | |
| 27 | def forward(self, x): |
| 28 | return self.bn(x) |
| 29 | |
| 30 | |
| 31 | class ModuleWithNestedBatchNorm(torch.nn.Module): |
| 32 | def __init__(self): |
| 33 | super().__init__() |
| 34 | self.nested_mod = ModuleWithBatchNorm() |
| 35 | |
| 36 | def forward(self, x): |
| 37 | return self.nested_mod(x) |
| 38 | |
| 39 | |
| 40 | @pytest.mark.cpu_test |
| 41 | def test_module_with_batchnorm_can_load(): |
| 42 | """Ensure the auto weight loader can load batchnorm stats.""" |
| 43 | mod = ModuleWithBatchNorm() |
| 44 | # Run some data through the module with batchnorm |
| 45 | mod(torch.Tensor([[1, 2], [3, 4]])) |
| 46 | |
| 47 | # Try to load the weights to a new instance |
| 48 | def weight_generator(): |
| 49 | yield from mod.state_dict().items() |
| 50 | |
| 51 | new_mod = ModuleWithBatchNorm() |
| 52 | |
| 53 | assert not torch.all(new_mod.bn.running_mean == mod.bn.running_mean) |
| 54 | assert not torch.all(new_mod.bn.running_var == mod.bn.running_var) |
| 55 | assert new_mod.bn.num_batches_tracked.item() == 0 |
| 56 | |
| 57 | loader = AutoWeightsLoader(new_mod) |
| 58 | loader.load_weights(weight_generator()) |
| 59 | |
| 60 | # Ensure the stats are updated |
| 61 | assert torch.all(new_mod.bn.running_mean == mod.bn.running_mean) |
| 62 | assert torch.all(new_mod.bn.running_var == mod.bn.running_var) |
| 63 | assert new_mod.bn.num_batches_tracked.item() == 1 |
| 64 | |
| 65 | |
| 66 | @pytest.mark.cpu_test |
| 67 | def test_module_with_child_containing_batchnorm_can_autoload(): |
| 68 | """Ensure the auto weight loader can load nested modules batchnorm stats.""" |
| 69 | mod = ModuleWithNestedBatchNorm() |
| 70 | # Run some data through the module with batchnorm |
| 71 | mod(torch.Tensor([[1, 2], [3, 4]])) |
| 72 | |
| 73 | # Try to load the weights to a new instance |
| 74 | def weight_generator(): |
| 75 | yield from mod.state_dict().items() |
| 76 | |
| 77 | new_mod = ModuleWithNestedBatchNorm() |
| 78 | |
| 79 | assert not torch.all( |
| 80 | new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean |
| 81 | ) |
| 82 | assert not torch.all( |
| 83 | new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var |
| 84 | ) |
| 85 | assert new_mod.nested_mod.bn.num_batches_tracked.item() == 0 |
| 86 | |
| 87 | loader = AutoWeightsLoader(new_mod) |
| 88 | loader.load_weights(weight_generator()) |
| 89 | |
| 90 | # Ensure the stats are updated |
| 91 | assert torch.all( |
| 92 | new_mod.nested_mod.bn.running_mean == mod.nested_mod.bn.running_mean |
| 93 | ) |
| 94 | assert torch.all(new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var) |
| 95 | assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1 |
| 96 | |
| 97 | |
| 98 | VOCAB_SIZE = 16 |
| 99 | HIDDEN_SIZE = 2 |
| 100 | |
| 101 | |
| 102 | class ModuleWithTiedWeights(torch.nn.Module): |
| 103 | """Mimics how models tie `lm_head` to the input embeddings.""" |
| 104 | |
| 105 | def __init__(self, tie: bool): |
| 106 | super().__init__() |
| 107 | self.model = torch.nn.Module() |
| 108 | self.model.embed_tokens = VocabParallelEmbedding(VOCAB_SIZE, HIDDEN_SIZE) |
| 109 | self.lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE) |
| 110 | if tie: |
| 111 | self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) |
| 112 | |
| 113 | |
| 114 | def make_embedding_weights(value: float) -> torch.Tensor: |
| 115 | return torch.full((VOCAB_SIZE, HIDDEN_SIZE), value) |
| 116 | |
| 117 | |
| 118 | @pytest.mark.cpu_test |
| 119 | @pytest.mark.usefixtures("dist_init") |
| 120 | @pytest.mark.parametrize("tie", [True, False]) |
| 121 | def test_module_skip_tied_weights(tie: bool): |
| 122 | """Tied weights must be loaded once, under the first of their names.""" |
| 123 | mod = ModuleWithTiedWeights(tie) |
| 124 | |
| 125 | weights = [ |
| 126 | ("model.embed_tokens.weight", make_embedding_weights(1.0)), |
| 127 | ("lm_head.weight", make_embedding_weights(2.0)), |
| 128 | ] |
| 129 | loaded = AutoWeightsLoader(mod).load_weights(iter(weights)) |
| 130 | |
| 131 | if tie: |
| 132 | assert loaded == {"model.embed_tokens.weight"} |
| 133 | assert torch.all(mod.lm_head.weight[:VOCAB_SIZE] == 1.0) |
| 134 | else: |
| 135 | assert loaded == {"model.embed_tokens.weight", "lm_head.weight"} |
| 136 | assert torch.all(mod.lm_head.weight[:VOCAB_SIZE] == 2.0) |
| 137 | |
| 138 | |
| 139 | @pytest.mark.cpu_test |
| 140 | @pytest.mark.usefixtures("dist_init") |
| 141 | def test_module_skip_tied_weights_without_canonical(): |
| 142 | """Skipping a tied weight must not leave the shared weight uninitialized.""" |
| 143 | mod = ModuleWithTiedWeights(tie=True) |
| 144 | |
| 145 | weights = [("lm_head.weight", make_embedding_weights(2.0))] |
| 146 | with pytest.raises(ValueError, match="model.embed_tokens.weight"): |
| 147 | AutoWeightsLoader(mod).load_weights(iter(weights)) |
| 148 | |
| 149 | |
| 150 | class ModuleWithSharedParam(torch.nn.Module): |
| 151 | """Mimics an MoE router shared between the MLP and its fused experts.""" |
| 152 | |
| 153 | def __init__(self): |
| 154 | super().__init__() |
| 155 | self.experts = torch.nn.Module() |
| 156 | self.experts.gate = torch.nn.Linear(2, 2, bias=False) |
| 157 | self.gate = self.experts.gate |
| 158 | |
| 159 | |
| 160 | @pytest.mark.cpu_test |
| 161 | def test_module_load_shared_params_that_are_not_tied_embeddings(): |
| 162 | """Only tied embeddings are skipped; other shared params must still load.""" |
| 163 | mod = ModuleWithSharedParam() |
| 164 | |
| 165 | weights = [("gate.weight", torch.Tensor([[1, 2], [3, 4]]))] |
| 166 | loaded = AutoWeightsLoader(mod).load_weights(iter(weights)) |
| 167 | |
| 168 | assert loaded == {"gate.weight"} |
| 169 | assert torch.all(mod.gate.weight == torch.Tensor([[1, 2], [3, 4]])) |
| 170 | |
| 171 | |
| 172 | class raise_if_cuda_sync: |
| 173 | def __enter__(self): |
| 174 | self.previous_debug_mode = torch.cuda.get_sync_debug_mode() |
| 175 | torch.cuda.set_sync_debug_mode("error") |
| 176 | |
| 177 | def __exit__(self, exception_type, exception_value, traceback): |
| 178 | torch.cuda.set_sync_debug_mode(self.previous_debug_mode) |
| 179 | |
| 180 | |
| 181 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") |
| 182 | def test_merge_multimodal_embeddings_no_sync(): |
| 183 | inputs_embeds = torch.zeros( |
| 184 | [5, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0" |
| 185 | ) |
| 186 | multimodal_embeddings = [ |
| 187 | torch.ones([3, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0") |
| 188 | ] |
| 189 | is_multimodal = torch.tensor([True, False, True, True, False], device="cpu") |
| 190 | with raise_if_cuda_sync(): |
| 191 | _merge_multimodal_embeddings( |
| 192 | inputs_embeds, multimodal_embeddings, is_multimodal |
| 193 | ) |
| 194 | |
| 195 | |
| 196 | @pytest.mark.cpu_test |
| 197 | def test_get_rename_mapper_keeps_only_renames(): |
| 198 | """`None` means "do not load", which is meaningless to the consumers of |
| 199 | this mapper (LoRA name parsing, quantization config layer lists), and |
| 200 | applying it would silently shrink their lists.""" |
| 201 | mapper = WeightsMapper( |
| 202 | orig_to_new_regex={re.compile(r"^drop_regex\."): None}, |
| 203 | orig_to_new_substr={"drop_substr": None, "keep_substr": "kept"}, |
| 204 | orig_to_new_stacked={".q_proj": (".qkv_proj", "q")}, |
| 205 | orig_to_new_prefix={"drop_prefix.": None, "keep_prefix.": "kept."}, |
| 206 | orig_to_new_suffix={".drop_suffix": None}, |
| 207 | ) |
| 208 | renames = mapper.get_rename_mapper() |
| 209 | |
| 210 | assert renames.orig_to_new_regex == {} |
| 211 | assert renames.orig_to_new_substr == {"keep_substr": "kept"} |
| 212 | assert renames.orig_to_new_stacked == {} |
| 213 | assert renames.orig_to_new_prefix == {"keep_prefix.": "kept."} |
| 214 | assert renames.orig_to_new_suffix == {} |
| 215 | |
| 216 | # Names the full mapper drops now survive unchanged. |
| 217 | for name in ("drop_regex.w", "drop_substr.w", "drop_prefix.w", "w.drop_suffix"): |
| 218 | assert mapper._map_name(name) is None |
| 219 | assert renames._map_name(name) == name |
| 220 | |
| 221 | |
| 222 | def test_weights_mapper_stacks_one_weight_into_several_shards(): |
| 223 | """One checkpoint tensor can feed more than one shard of a stacked |
| 224 | parameter, and a dropped name stays dropped rather than reappearing as a |
| 225 | shard of the stacked one.""" |
| 226 | mapper = WeightsMapper( |
| 227 | orig_to_new_substr={"layers.1.k_proj.": None}, |
| 228 | orig_to_new_stacked={".k_proj.": (".qkv_proj.", ["k", "v"])}, |
| 229 | ) |
| 230 | weight = torch.ones(2) |
| 231 | weights = [(f"layers.{i}.k_proj.weight", weight) for i in (0, 1)] |
| 232 | |
| 233 | mapped = list(mapper.apply(weights)) |
| 234 | |
| 235 | assert [(name, w.shard_id) for name, w in mapped] == [ |
| 236 | ("layers.0.qkv_proj.weight", "k"), |
| 237 | ("layers.0.qkv_proj.weight", "v"), |
| 238 | ] |
| 239 | # Each shard needs its own tensor object, but they alias one allocation. |
| 240 | assert mapped[0][1] is not mapped[1][1] |
| 241 | assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr() == weight.data_ptr() |