@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 81be2e09ae

[Weight processing] Copy over `new_data` attributes in `replace_parameter` (#49601)

Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

fxmarty-amd committed Aug 6, 2026 at 19:03 UTC 81be2e09aebfd1c45b3ed9f73d2850da8a72984c
10 files changed +334 -51
tests/model_executor/test_utils.py new
+283
@@ -0,0 +1,283 @@
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
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py
-5
@@ -145,11 +145,6 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
145
146 from vllm._aiter_ops import rocm_aiter_ops
147
148 - # Re-tag the preshuffled weights: replace_parameter drops the
149 - # is_shuffled flag, without which aiter picks a broken CK kernel.
150 - w1.is_shuffled = True
151 - w2.is_shuffled = True
152 -
148 limit = self.quant_config.gemm1_clamp_limit
149 swiglu_limit = 0.0 if limit is None else float(limit)
150
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+4
@@ -483,10 +483,14 @@ def convert_to_fp8_moe_kernel_format(
483 )
484 elif fp8_backend == Fp8MoeBackend.AITER:
485 w13, w2 = rocm_aiter_ops.shuffle_weights(w13, w2)
486 + w13.is_shuffled = True
487 + w2.is_shuffled = True
488 elif fp8_backend == Fp8MoeBackend.AITER_MXFP8:
489 w13, w2, w13_scale, w2_scale = rocm_aiter_ops.shuffle_mxfp8_moe_weights(
490 w13, w2, w13_scale, w2_scale
491 )
492 + w13.is_shuffled = True
493 + w2.is_shuffled = True
494 elif fp8_backend == Fp8MoeBackend.HUMMING:
495 from vllm.model_executor.layers.quantization.utils.humming_utils import (
496 convert_to_humming_moe_kernel_format,
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+6
@@ -1081,6 +1081,9 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
1081 .view(-1, n)
1082 )
1083
1084 + w13_weight.is_shuffled = True
1085 + w2_weight.is_shuffled = True
1086 +
1087 return (
1088 w13_weight,
1089 w2_weight,
@@ -1504,6 +1507,9 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1507 False,
1508 )
1509
1510 + w13_weight.is_shuffled = True
1511 + w2_weight.is_shuffled = True
1512 +
1513 return (
1514 w13_weight,
1515 w2_weight,
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+2
@@ -332,6 +332,8 @@ def convert_to_unquantized_kernel_format(
332 ) -> tuple[torch.Tensor, torch.Tensor]:
333 if unquantized_backend == UnquantizedMoeBackend.AITER:
334 w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(w13_weight, w2_weight)
335 + w13_weight.is_shuffled = True
336 + w2_weight.is_shuffled = True
337
338 elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
339 if moe_config.is_act_and_mul:
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
-5
@@ -146,11 +146,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
146 replace_parameter(layer, "w13_weight", w13_new, prefer_copy=is_weight_update)
147 replace_parameter(layer, "w2_weight", w2_new, prefer_copy=is_weight_update)
148
149 - # AITER backend requires weights to be marked as shuffled.
150 - if self.unquantized_backend == UnquantizedMoeBackend.AITER:
151 - layer.w13_weight.is_shuffled = True
152 - layer.w2_weight.is_shuffled = True
153 -
149 if not is_weight_update:
150 # Setup moe kernel only on the first call. For the unquantized
151 # method, moe_quant_config carries no quantized scales -- only
vllm/model_executor/layers/quantization/fp8.py
-6
@@ -28,7 +28,6 @@ from vllm.model_executor.layers.fused_moe.config import (
28 FusedMoEQuantConfig,
29 )
30 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
31 - Fp8MoeBackend,
31 convert_to_fp8_moe_kernel_format,
32 make_fp8_moe_kernel,
33 make_fp8_moe_quant_config,
@@ -671,11 +670,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
670 replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
671 replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
672
674 - # AITER backend requires weights to be marked as shuffled.
675 - if self.fp8_backend == Fp8MoeBackend.AITER:
676 - layer.w13_weight.is_shuffled = True
677 - layer.w2_weight.is_shuffled = True
678 -
673 self.moe_quant_config = self.get_fused_moe_quant_config(layer)
674 assert self.moe_quant_config is not None
675 assert self.experts_cls is not None
vllm/model_executor/layers/quantization/mxfp4.py
+4 -13
@@ -366,11 +366,6 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
366 self.w13_precision_config = w13_scale
367 self.w2_precision_config = w2_scale
368
369 - # AITER backend requires weights to be marked as shuffled.
370 - if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
371 - layer.w13_weight.is_shuffled = True
372 - layer.w2_weight.is_shuffled = True
373 -
369 if w13_bias is not None and w2_bias is not None:
370 replace_parameter(layer, "w13_bias", w13_bias)
371 replace_parameter(layer, "w2_bias", w2_bias)
@@ -760,14 +755,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
755 self.w13_precision_config = w13_scale
756 self.w2_precision_config = w2_scale
757
763 - # AITER backend requires weights to be marked as shuffled.
764 - if (
765 - self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16
766 - and not uses_triton_weight_format
767 - ):
768 - layer.w13_weight.is_shuffled = True
769 - layer.w2_weight.is_shuffled = True
770 -
758 if w13_bias is not None and w2_bias is not None:
759 replace_parameter(layer, "w13_bias", w13_bias)
760 replace_parameter(layer, "w2_bias", w2_bias)
@@ -816,6 +803,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
803 w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), num_experts, guinterleave
804 )
805 w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1]))
806 +
807 + w13.is_shuffled = True
808 + w2.is_shuffled = True
809 +
810 return w13, w2, w13_scale, w2_scale
811
812 def process_weights_after_loading(self, layer):
vllm/model_executor/layers/quantization/quark/quark_moe.py
-9
@@ -27,7 +27,6 @@ from vllm.model_executor.layers.fused_moe.config import (
27 ocp_mx_moe_quant_config,
28 )
29 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
30 - Fp8MoeBackend,
30 convert_to_fp8_moe_kernel_format,
31 make_fp8_moe_kernel,
32 make_fp8_moe_quant_config,
@@ -447,10 +446,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
446 replace_parameter(layer, "w13_weight_scale", w13_scale)
447 replace_parameter(layer, "w2_weight_scale", w2_scale)
448
450 - if self.fp8_backend == Fp8MoeBackend.AITER:
451 - layer.w13_weight.is_shuffled = True
452 - layer.w2_weight.is_shuffled = True
453 -
449 self.moe_quant_config = self.get_fused_moe_quant_config(layer)
450 assert self.moe_quant_config is not None
451 assert self.experts_cls is not None
@@ -1328,10 +1323,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1323 replace_parameter(layer, "w13_bias", w13_bias)
1324 replace_parameter(layer, "w2_bias", w2_bias)
1325
1331 - if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:
1332 - layer.w13_weight.is_shuffled = True
1333 - layer.w2_weight.is_shuffled = True
1334 -
1326 torch.accelerator.empty_cache()
1327
1328 # Build quant config and kernel
vllm/model_executor/utils.py
+35 -13
@@ -47,21 +47,29 @@ def set_weight_attrs(
47 def replace_parameter(
48 layer: torch.nn.Module,
49 param_name: str,
50 - new_data: torch.Tensor | None,
50 + new_tensor: torch.Tensor | None,
51 prefer_copy: bool = False,
52 ):
53 """
54 Replace a parameter of a layer while maintaining the ability to reload the weight.
55 Called within implementations of the `process_weights_after_loading` method.
56
57 + Custom attributes set on ``new_tensor`` (e.g. kernel dispatch flags such as
58 + ``is_shuffled``) are carried over to the replacement parameter, except
59 + ``weight_loader``, which is always taken from the existing parameter.
60 +
61 + Attributes of the existing parameter are otherwise dropped when a new
62 + parameter is registered, but kept when ``prefer_copy`` reuses it in place.
63 +
64 This function should not be called on weights which are tied/shared
65
66 Args:
67 layer: Layer containing parameter to replace
68 param_name: Name of parameter to replace
62 - new_data: New data of the new parameter, or None to set the parameter to None
69 + new_tensor: New data of the new parameter, or None to set the parameter
70 + to None
71 prefer_copy: If True and the existing parameter is compatible with
64 - ``new_data`` (same shape, dtype, and device), copy ``new_data``
72 + ``new_tensor`` (same shape, dtype, and device), copy ``new_tensor``
73 into the existing parameter in place rather than re-registering
74 a new parameter. This preserves the parameter's storage address
75 (``data_ptr``), which is required for captured CUDA graphs to
@@ -69,27 +77,41 @@ def replace_parameter(
77 """
78 # should not be used on a tied/shared param
79
72 - # If new_data is None, set the parameter to None
73 - if new_data is None:
80 + # If new_tensor is None, set the parameter to None
81 + if new_tensor is None:
82 setattr(layer, param_name, None)
83 return
84
77 - if isinstance(new_data, torch.nn.Parameter):
78 - new_data = new_data.data
79 -
85 old_param: torch.nn.Parameter | None = getattr(layer, param_name, None)
86
87 + new_tensor_attrs = dict(new_tensor.__dict__)
88 +
89 + # `weight_loader` is the only attribute not ported over from new_tensor,
90 + # old_param.weight_loader only is supported. `_weight_loader` is
91 + # `BasevLLMParameter`'s backing field for its `weight_loader` property, so
92 + # a loader riding along under that name has to be dropped as well.
93 + new_tensor_attrs.pop("weight_loader", None)
94 + new_tensor_attrs.pop("_weight_loader", None)
95 +
96 + if isinstance(new_tensor, torch.nn.Parameter):
97 + new_tensor = new_tensor.data
98 +
99 if (
100 prefer_copy
101 and old_param is not None
85 - and old_param.shape == new_data.shape
86 - and old_param.dtype == new_data.dtype
87 - and old_param.device == new_data.device
102 + and old_param.shape == new_tensor.shape
103 + and old_param.dtype == new_tensor.dtype
104 + and old_param.device == new_tensor.device
105 ):
89 - old_param.copy_(new_data)
106 + old_param.copy_(new_tensor)
107 + for attr_name, attr in new_tensor_attrs.items():
108 + setattr(old_param, attr_name, attr)
109 return
110
92 - new_param = torch.nn.Parameter(new_data, requires_grad=False)
111 + new_param = torch.nn.Parameter(new_tensor, requires_grad=False)
112 +
113 + for attr_name, attr in new_tensor_attrs.items():
114 + setattr(new_param, attr_name, attr)
115
116 if old_param is not None and hasattr(old_param, "weight_loader"):
117 weight_loader = old_param.weight_loader