| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.parameter import ModelWeightParameter, PackedvLLMParameter |
| 9 | from vllm.model_executor.utils import replace_parameter |
| 10 | |
| 11 | |
| 12 | @pytest.fixture |
| 13 | def single_rank_tp(monkeypatch: pytest.MonkeyPatch) -> None: |
| 14 | """`BasevLLMParameter.__init__` queries the TP group, which is not |
| 15 | initialized in a unit test. Pin it to a single rank. |
| 16 | """ |
| 17 | monkeypatch.setattr( |
| 18 | "vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0 |
| 19 | ) |
| 20 | monkeypatch.setattr( |
| 21 | "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1 |
| 22 | ) |
| 23 | |
| 24 | |
| 25 | @pytest.mark.parametrize("prefer_copy", [False, True]) |
| 26 | @pytest.mark.parametrize("wrap_in_parameter", [False, True]) |
| 27 | def test_replace_parameter_preserves_custom_attribute( |
| 28 | prefer_copy: bool, wrap_in_parameter: bool |
| 29 | ) -> None: |
| 30 | """`replace_parameter` must carry over attributes attached to the |
| 31 | replacement tensor (e.g. the `is_shuffled` flag set by AITER weight |
| 32 | preprocessing in `Fp8MoEMethod.process_weights_after_loading`). |
| 33 | |
| 34 | The replacement is passed both as a plain `torch.Tensor` and as a |
| 35 | `torch.nn.Parameter`, since the latter is unwrapped through `.data` |
| 36 | internally, which does not carry attributes over either. |
| 37 | """ |
| 38 | layer = torch.nn.Module() |
| 39 | layer.register_parameter( |
| 40 | "weight", torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) |
| 41 | ) |
| 42 | original_data_ptr = layer.weight.data_ptr() |
| 43 | |
| 44 | new_data: torch.Tensor = torch.ones(4, 4).t() |
| 45 | if wrap_in_parameter: |
| 46 | new_data = torch.nn.Parameter(new_data, requires_grad=False) |
| 47 | new_data.is_shuffled = True |
| 48 | |
| 49 | # Sanity check on the assumption above: tensor internals are not |
| 50 | # reachable through `__dict__`, only user-set custom attributes are. |
| 51 | assert new_data.__dict__.keys() == {"is_shuffled"} |
| 52 | |
| 53 | replace_parameter(layer, "weight", new_data, prefer_copy=prefer_copy) |
| 54 | |
| 55 | assert isinstance(layer.weight, torch.nn.Parameter) |
| 56 | assert layer.weight.is_shuffled is True |
| 57 | assert torch.equal(layer.weight.data, new_data) |
| 58 | assert layer.weight.device == new_data.device |
| 59 | |
| 60 | if prefer_copy: |
| 61 | # The existing storage is reused, so addresses captured in CUDA graphs |
| 62 | # stay valid across the update. |
| 63 | assert layer.weight.data_ptr() == original_data_ptr |
| 64 | else: |
| 65 | assert layer.weight.stride() == new_data.stride() |
| 66 | assert layer.weight.data_ptr() == new_data.data_ptr() |
| 67 | |
| 68 | |
| 69 | @pytest.mark.parametrize("prefer_copy", [False, True]) |
| 70 | def test_replace_parameter_preserves_weight_loader(prefer_copy: bool) -> None: |
| 71 | """The reload path must survive replacement: the old parameter's |
| 72 | `weight_loader` is carried over as-is, so it is still invoked as |
| 73 | `weight_loader(param, loaded_weight)`. |
| 74 | |
| 75 | Real loaders are frequently bound methods of the layer (`RoutedExperts` |
| 76 | hands `self.weight_loader` to `create_weights`); what must not happen is |
| 77 | the carry-over re-binding the loader to the new parameter, which would |
| 78 | shift `loaded_weight` into the `param` slot. |
| 79 | """ |
| 80 | calls: list[tuple[torch.Tensor, torch.Tensor]] = [] |
| 81 | |
| 82 | def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: |
| 83 | calls.append((param, loaded_weight)) |
| 84 | |
| 85 | layer = torch.nn.Module() |
| 86 | old_param = torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) |
| 87 | old_param.weight_loader = weight_loader |
| 88 | layer.register_parameter("weight", old_param) |
| 89 | |
| 90 | replace_parameter(layer, "weight", torch.ones(4, 4), prefer_copy=prefer_copy) |
| 91 | |
| 92 | # Re-binding to the new parameter would shift `loaded_weight` into the |
| 93 | # `param` slot, so the plain function must not have grown a `__self__`. |
| 94 | assert not hasattr(layer.weight.weight_loader, "__self__") |
| 95 | |
| 96 | loaded_weight = torch.full((4, 4), 2.0) |
| 97 | layer.weight.weight_loader(layer.weight, loaded_weight) |
| 98 | |
| 99 | assert len(calls) == 1 |
| 100 | assert calls[0][0] is layer.weight |
| 101 | assert calls[0][1] is loaded_weight |
| 102 | |
| 103 | |
| 104 | def test_replace_parameter_weight_loader_comes_from_old_parameter() -> None: |
| 105 | """`weight_loader` is deliberately excluded from the attribute carry-over: |
| 106 | the old parameter's loader is authoritative. A loader riding along on the |
| 107 | replacement tensor must neither shadow it nor trip the overwrite assertion |
| 108 | in `set_weight_attrs`, and the other attributes must still be carried over. |
| 109 | |
| 110 | `_weight_loader` is excluded for the same reason: it is the backing field |
| 111 | of `BasevLLMParameter.weight_loader`, so leaving it in would smuggle a |
| 112 | stale loader past the `weight_loader` exclusion. |
| 113 | """ |
| 114 | calls: list[str] = [] |
| 115 | |
| 116 | def old_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: |
| 117 | calls.append("old") |
| 118 | |
| 119 | def stale_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: |
| 120 | calls.append("stale") |
| 121 | |
| 122 | layer = torch.nn.Module() |
| 123 | old_param = torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) |
| 124 | old_param.weight_loader = old_weight_loader |
| 125 | layer.register_parameter("weight", old_param) |
| 126 | |
| 127 | new_data = torch.ones(4, 4) |
| 128 | new_data.weight_loader = stale_weight_loader |
| 129 | new_data._weight_loader = stale_weight_loader |
| 130 | new_data.is_shuffled = True |
| 131 | |
| 132 | replace_parameter(layer, "weight", new_data) |
| 133 | |
| 134 | layer.weight.weight_loader(layer.weight, torch.full((4, 4), 2.0)) |
| 135 | |
| 136 | assert calls == ["old"] |
| 137 | assert not hasattr(layer.weight, "_weight_loader") |
| 138 | assert layer.weight.is_shuffled is True |
| 139 | |
| 140 | |
| 141 | def test_replace_parameter_does_not_rebind_plain_function_attribute() -> None: |
| 142 | """A plain function stored on the replacement tensor must be carried over |
| 143 | verbatim; re-binding it as a method of the new parameter would silently |
| 144 | shift the arguments of every subsequent call. |
| 145 | """ |
| 146 | |
| 147 | def scale_for(group_size: int) -> int: |
| 148 | return group_size * 2 |
| 149 | |
| 150 | layer = torch.nn.Module() |
| 151 | layer.register_parameter( |
| 152 | "weight", torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) |
| 153 | ) |
| 154 | |
| 155 | new_data = torch.ones(4, 4) |
| 156 | new_data.scale_for = scale_for |
| 157 | |
| 158 | replace_parameter(layer, "weight", new_data) |
| 159 | |
| 160 | assert layer.weight.scale_for is scale_for |
| 161 | assert layer.weight.scale_for(3) == 6 |
| 162 | |
| 163 | |
| 164 | def test_replace_parameter_preserves_bound_method_attribute() -> None: |
| 165 | """A callable already bound to another object must keep pointing at that |
| 166 | object after replacement. |
| 167 | """ |
| 168 | |
| 169 | class KernelDispatcher: |
| 170 | def __init__(self) -> None: |
| 171 | self.calls: list[int] = [] |
| 172 | |
| 173 | def record(self, value: int) -> None: |
| 174 | self.calls.append(value) |
| 175 | |
| 176 | dispatcher = KernelDispatcher() |
| 177 | |
| 178 | layer = torch.nn.Module() |
| 179 | layer.register_parameter( |
| 180 | "weight", torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) |
| 181 | ) |
| 182 | |
| 183 | new_data = torch.ones(4, 4) |
| 184 | new_data.record = dispatcher.record |
| 185 | |
| 186 | replace_parameter(layer, "weight", new_data) |
| 187 | |
| 188 | assert layer.weight.record.__self__ is dispatcher |
| 189 | layer.weight.record(7) |
| 190 | assert dispatcher.calls == [7] |
| 191 | |
| 192 | |
| 193 | @pytest.mark.parametrize("param_kind", ["plain", "model_weight", "packed"]) |
| 194 | def test_replace_parameter_attributes_from_the_layers_own_parameter( |
| 195 | single_rank_tp: None, param_kind: str |
| 196 | ) -> None: |
| 197 | """Several callers hand back the parameter they were given, after mutating |
| 198 | its `.data` in place: the HUMMING branch of |
| 199 | `convert_to_fp8_moe_kernel_format` (`w13 = layer.w13_weight`), the |
| 200 | `AITER_MXFP4_BF16` branch of |
| 201 | `convert_gpt_oss_weight_to_mxfp4_moe_kernel_format`, the no-transpose |
| 202 | branch of `XPUFP8ScaledMM` (`layer_weight = w`), and `auto_awq`/`auto_gptq`, |
| 203 | which pass whatever is currently registered -- possibly still a |
| 204 | `BasevLLMParameter` subclass rather than a plain `Parameter`. |
| 205 | |
| 206 | `new_data is old_param` there, so every attribute of the replacement is by |
| 207 | definition an attribute of the old parameter; pin exactly which survive. |
| 208 | |
| 209 | For the vLLM parameter classes only the private backing fields come along. |
| 210 | The public names weight loading branches on -- `output_dim`, `input_dim`, |
| 211 | `packed_dim` (`getattr(param, "output_dim", None)` in `linear.py`) -- are |
| 212 | class-level properties, so they do not survive onto the plain replacement |
| 213 | and cannot be read back stale after the layout has changed. |
| 214 | """ |
| 215 | |
| 216 | class Layer(torch.nn.Module): |
| 217 | def weight_loader( |
| 218 | self, param: torch.Tensor, loaded_weight: torch.Tensor |
| 219 | ) -> None: |
| 220 | pass |
| 221 | |
| 222 | layer = Layer() |
| 223 | data = torch.zeros(4, 4) |
| 224 | loader = layer.weight_loader |
| 225 | vllm_param_attrs = { |
| 226 | "_weight_loader": loader, |
| 227 | "_input_dim": 1, |
| 228 | "_output_dim": 0, |
| 229 | "tp_rank": 0, |
| 230 | "tp_size": 1, |
| 231 | } |
| 232 | |
| 233 | old_param: torch.nn.Parameter |
| 234 | if param_kind == "plain": |
| 235 | old_param = torch.nn.Parameter(data, requires_grad=False) |
| 236 | old_param.weight_loader = loader |
| 237 | # MoE scale marker, read back by `RoutedExperts.weight_loader`. |
| 238 | old_param.quant_method = "block" |
| 239 | source_attrs = {"weight_loader": loader, "quant_method": "block"} |
| 240 | elif param_kind == "model_weight": |
| 241 | old_param = ModelWeightParameter( |
| 242 | data=data, input_dim=1, output_dim=0, weight_loader=loader |
| 243 | ) |
| 244 | source_attrs = dict(vllm_param_attrs) |
| 245 | else: |
| 246 | old_param = PackedvLLMParameter( |
| 247 | data=data, |
| 248 | input_dim=1, |
| 249 | output_dim=0, |
| 250 | packed_dim=0, |
| 251 | packed_factor=8, |
| 252 | weight_loader=loader, |
| 253 | ) |
| 254 | source_attrs = dict(vllm_param_attrs) |
| 255 | source_attrs |= { |
| 256 | "_packed_dim": 0, |
| 257 | "_packed_factor": 8, |
| 258 | "_marlin_tile_size": None, |
| 259 | } |
| 260 | |
| 261 | layer.register_parameter("weight", old_param) |
| 262 | old_param.data = torch.ones(4, 4) |
| 263 | |
| 264 | # Pin what the caller hands back: on the vLLM classes `weight_loader` is a |
| 265 | # property, so the loader sits under `_weight_loader` instead. |
| 266 | assert dict(old_param.__dict__) == source_attrs |
| 267 | |
| 268 | replace_parameter(layer, "weight", old_param) |
| 269 | |
| 270 | # Everything carries over by value, except that the loader is re-read from |
| 271 | # the old parameter and lands under its public name. Comparing values (not |
| 272 | # just keys) is what rules out the loader being re-bound to the new |
| 273 | # parameter, which would shift `loaded_weight` into the `param` slot. |
| 274 | expected = dict(source_attrs) |
| 275 | expected.pop("_weight_loader", None) |
| 276 | expected["weight_loader"] = loader |
| 277 | |
| 278 | assert type(layer.weight) is torch.nn.Parameter |
| 279 | assert dict(layer.weight.__dict__) == expected |
| 280 | |
| 281 | if param_kind != "plain": |
| 282 | for public_name in ("output_dim", "input_dim", "packed_dim", "packed_factor"): |
| 283 | assert getattr(layer.weight, public_name, None) is None |