From c9e5bf813530fb9ce06024e075da0f520b0718c8 Mon Sep 17 00:00:00 2001 From: hallerite Date: Wed, 10 Jun 2026 15:42:05 +0200 Subject: [PATCH] [Bugfix] Fix layerwise reload dropping params after a composed weight loader (#44814) Signed-off-by: hallerite Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Kyle Sayers --- .../model_loader/test_reload.py | 81 +++++++++++++++++++ .../model_loader/reload/meta.py | 16 +++- 2 files changed, 96 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 0a290a00a8..b3ed0c11bb 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -25,6 +25,10 @@ from vllm.model_executor.model_loader.reload.meta import ( ) from vllm.model_executor.model_loader.reload.types import LayerReloadingInfo from vllm.model_executor.model_loader.reload.utils import get_layer_tensors +from vllm.model_executor.model_loader.weight_utils import ( + composed_weight_loader, + default_weight_loader, +) from vllm.platforms import current_platform @@ -178,6 +182,83 @@ def test_get_numel_loaded(): assert ret == "value" +def test_get_numel_loaded_caps_at_param_size(): + # composed_weight_loader copies into the param twice (the load and the + # in-place post-load transform), but only param.numel() distinct elements + # are loaded. get_numel_loaded must not double-count, otherwise a layer's + # loaded-element total can be reached early and trailing params get dropped. + param = torch.empty(10) + loaded_weight = torch.ones(10) + loader = composed_weight_loader(default_weight_loader, lambda x: x + 1) + + args = inspect.signature(loader).bind(param, loaded_weight) + num_loaded, _ = get_numel_loaded(loader, args) + assert num_loaded == 10 + + +class _ComposedLoaderLayer(torch.nn.Module): + """Mimics a Mamba2 mixer's equal-numel direct params (A, D, dt_bias). + + ``A`` uses ``composed_weight_loader`` (an extra in-place transform copy), + matching ``MambaMixer2`` where ``A`` is loaded as ``-exp(A_log)``. + """ + + def __init__(self): + super().__init__() + self.A = torch.nn.Parameter(torch.empty(4, dtype=torch.float32)) + self.D = torch.nn.Parameter(torch.ones(4)) + self.dt_bias = torch.nn.Parameter(torch.ones(4)) + self.A.weight_loader = composed_weight_loader( + default_weight_loader, lambda x: -torch.exp(x.float()) + ) + self.D.weight_loader = default_weight_loader + self.dt_bias.weight_loader = default_weight_loader + + +def test_layerwise_reload_composed_loader_does_not_drop_params(monkeypatch): + # Regression test: a composed_weight_loader param (A) used to double-count + # its elements, finalizing the layer before the trailing param (D) was + # loaded and leaving it as uninitialized materialized memory. + layer = _ComposedLoaderLayer() + model = torch.nn.Sequential(layer) + + def materialize_with_sentinel(meta_tensor): + tensor = torch.empty_strided( + size=tuple(meta_tensor.size()), + stride=tuple(meta_tensor.stride()), + dtype=meta_tensor.dtype, + requires_grad=False, + ) + tensor.fill_(float("nan")) + tensor.__class__ = meta_tensor.__class__ + tensor.__dict__ = meta_tensor.__dict__.copy() + return tensor + + monkeypatch.setattr( + reload_meta, "materialize_meta_tensor", materialize_with_sentinel + ) + + loaded = { + "A": torch.full((4,), 0.5), + "dt_bias": torch.full((4,), 3.0), + "D": torch.full((4,), 7.0), + } + + record_metadata_for_reloading(model) + initialize_layerwise_reload(model) + # Mimic real load_weights: resolve params once, then load in checkpoint + # order with D last (the param that was dropped). + params = dict(layer.named_parameters()) + for name in ("A", "dt_bias", "D"): + param = params[name] + param.weight_loader(param, loaded[name]) + finalize_layerwise_reload(model, model_config=None) + + assert torch.equal(layer.A, -torch.exp(loaded["A"])) + assert torch.equal(layer.dt_bias, loaded["dt_bias"]) + assert torch.equal(layer.D, loaded["D"]) + + def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch): layer = _AliasedBufferLayer() model = torch.nn.Sequential(layer) diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index 283a98de28..824d5c8b0f 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -185,4 +185,18 @@ def get_numel_loaded( """ with CopyCounter() as counter: return_value = weight_loader(*args.args, **args.kwargs) - return counter.copied_numel, return_value + + # A weight loader fills a single destination parameter, so the number of + # loaded elements is at most that parameter's size. Some loaders copy into + # the parameter more than once -- e.g. ``composed_weight_loader`` runs an + # in-place post-load transform (``param.copy_(fn(param))``) on top of the + # initial copy -- which would make CopyCounter report twice the parameter + # size. Over-counting inflates the layer's loaded-element total and can + # finalize the layer before every parameter is loaded, silently dropping + # the trailing parameter(s) (e.g. Mamba ``mixer.D``). Cap the count at the + # destination size to keep the per-layer accounting correct. + numel = counter.copied_numel + param = args.arguments.get("param", None) + if isinstance(param, torch.Tensor): + numel = min(numel, param.numel()) + return numel, return_value