A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Fix layerwise reload dropping params after a composed weight loader (#44814)
Signed-off-by: hallerite <git@hallerite.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Kyle Sayers <kylesayrs@gmail.com>
hallerite committed
Jun 10, 2026 at 15:42 UTC
c9e5bf813530fb9ce06024e075da0f520b0718c8
2 files changed
+96
-1
tests/model_executor/model_loader/test_reload.py
+81
@@ -25,6 +25,10 @@ from vllm.model_executor.model_loader.reload.meta import (
25
)
26
from vllm.model_executor.model_loader.reload.types import LayerReloadingInfo
27
from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
28
+from vllm.model_executor.model_loader.weight_utils import (
29
+ composed_weight_loader,
30
+ default_weight_loader,
31
+)
32
from vllm.platforms import current_platform
33
34
@@ -178,6 +182,83 @@ def test_get_numel_loaded():
182
assert ret == "value"
183
184
185
+def test_get_numel_loaded_caps_at_param_size():
186
+ # composed_weight_loader copies into the param twice (the load and the
187
+ # in-place post-load transform), but only param.numel() distinct elements
188
+ # are loaded. get_numel_loaded must not double-count, otherwise a layer's
189
+ # loaded-element total can be reached early and trailing params get dropped.
190
+ param = torch.empty(10)
191
+ loaded_weight = torch.ones(10)
192
+ loader = composed_weight_loader(default_weight_loader, lambda x: x + 1)
193
+
194
+ args = inspect.signature(loader).bind(param, loaded_weight)
195
+ num_loaded, _ = get_numel_loaded(loader, args)
196
+ assert num_loaded == 10
197
+
198
+
199
+class _ComposedLoaderLayer(torch.nn.Module):
200
+ """Mimics a Mamba2 mixer's equal-numel direct params (A, D, dt_bias).
201
+
202
+ ``A`` uses ``composed_weight_loader`` (an extra in-place transform copy),
203
+ matching ``MambaMixer2`` where ``A`` is loaded as ``-exp(A_log)``.
204
+ """
205
+
206
+ def __init__(self):
207
+ super().__init__()
208
+ self.A = torch.nn.Parameter(torch.empty(4, dtype=torch.float32))
209
+ self.D = torch.nn.Parameter(torch.ones(4))
210
+ self.dt_bias = torch.nn.Parameter(torch.ones(4))
211
+ self.A.weight_loader = composed_weight_loader(
212
+ default_weight_loader, lambda x: -torch.exp(x.float())
213
+ )
214
+ self.D.weight_loader = default_weight_loader
215
+ self.dt_bias.weight_loader = default_weight_loader
216
+
217
+
218
+def test_layerwise_reload_composed_loader_does_not_drop_params(monkeypatch):
219
+ # Regression test: a composed_weight_loader param (A) used to double-count
220
+ # its elements, finalizing the layer before the trailing param (D) was
221
+ # loaded and leaving it as uninitialized materialized memory.
222
+ layer = _ComposedLoaderLayer()
223
+ model = torch.nn.Sequential(layer)
224
+
225
+ def materialize_with_sentinel(meta_tensor):
226
+ tensor = torch.empty_strided(
227
+ size=tuple(meta_tensor.size()),
228
+ stride=tuple(meta_tensor.stride()),
229
+ dtype=meta_tensor.dtype,
230
+ requires_grad=False,
231
+ )
232
+ tensor.fill_(float("nan"))
233
+ tensor.__class__ = meta_tensor.__class__
234
+ tensor.__dict__ = meta_tensor.__dict__.copy()
235
+ return tensor
236
+
237
+ monkeypatch.setattr(
238
+ reload_meta, "materialize_meta_tensor", materialize_with_sentinel
239
+ )
240
+
241
+ loaded = {
242
+ "A": torch.full((4,), 0.5),
243
+ "dt_bias": torch.full((4,), 3.0),
244
+ "D": torch.full((4,), 7.0),
245
+ }
246
+
247
+ record_metadata_for_reloading(model)
248
+ initialize_layerwise_reload(model)
249
+ # Mimic real load_weights: resolve params once, then load in checkpoint
250
+ # order with D last (the param that was dropped).
251
+ params = dict(layer.named_parameters())
252
+ for name in ("A", "dt_bias", "D"):
253
+ param = params[name]
254
+ param.weight_loader(param, loaded[name])
255
+ finalize_layerwise_reload(model, model_config=None)
256
+
257
+ assert torch.equal(layer.A, -torch.exp(loaded["A"]))
258
+ assert torch.equal(layer.dt_bias, loaded["dt_bias"])
259
+ assert torch.equal(layer.D, loaded["D"])
260
+
261
+
262
def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch):
263
layer = _AliasedBufferLayer()
264
model = torch.nn.Sequential(layer)
vllm/model_executor/model_loader/reload/meta.py
+15
-1
@@ -185,4 +185,18 @@ def get_numel_loaded(
185
"""
186
with CopyCounter() as counter:
187
return_value = weight_loader(*args.args, **args.kwargs)
188
- return counter.copied_numel, return_value
188
+
189
+ # A weight loader fills a single destination parameter, so the number of
190
+ # loaded elements is at most that parameter's size. Some loaders copy into
191
+ # the parameter more than once -- e.g. ``composed_weight_loader`` runs an
192
+ # in-place post-load transform (``param.copy_(fn(param))``) on top of the
193
+ # initial copy -- which would make CopyCounter report twice the parameter
194
+ # size. Over-counting inflates the layer's loaded-element total and can
195
+ # finalize the layer before every parameter is loaded, silently dropping
196
+ # the trailing parameter(s) (e.g. Mamba ``mixer.D``). Cap the count at the
197
+ # destination size to keep the per-layer accounting correct.
198
+ numel = counter.copied_numel
199
+ param = args.arguments.get("param", None)
200
+ if isinstance(param, torch.Tensor):
201
+ numel = min(numel, param.numel())
202
+ return numel, return_value