[Bugfix][Model Loader] Make model.load_weights safe to invoke on already-initialized model

tests/model_executor/model_loader/test_reload.py

@@ -10,9 +10,11 @@
 
 import vllm.model_executor.model_loader.reload.meta as reload_meta
 from vllm.model_executor.layers.linear import QKVParallelLinear
+from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.reload.layerwise import (
     finalize_layerwise_reload,
     initialize_layerwise_reload,
+    make_load_weights_safe_for_reload,
     record_metadata_for_reloading,
 )
 from vllm.model_executor.model_loader.reload.meta import (
@@ -237,6 +239,201 @@ def materialize_with_sentinel(meta_tensor):
     )
 
 
+class _LayoutSwapMethod(QuantizeMethodBase):
+    """Mimics `UnquantizedFusedMoEMethod` for the FlashInfer CUTLASS backend.
+
+    `process_weights_after_loading` swaps the two halves of `layer.weight`
+    in place, modelling `swap_w13_to_w31`. The per-shard weight loader writes
+    the first half on `shard_id="w1"` and the second half on `shard_id="w3"`
+    — same convention as `FusedMoE._load_w13`. The combination triggers
+    https://github.com/vllm-project/vllm/issues/42821 when `load_weights`
+    is called a second time without re-routing through the layerwise reload
+    pipeline.
+    """
+
+    def create_weights(self, layer, *args, **kwargs):  # pragma: no cover
+        return
+
+    def apply(self, layer, *args, **kwargs):  # pragma: no cover
+        return layer.weight
+
+    def process_weights_after_loading(self, layer):
+        # Swap halves of layer.weight in place (analog of `swap_w13_to_w31`).
+        # After loading checkpoint-format `[w1; w3]`, the swap produces
+        # `[w3; w1]` which is the kernel-expected layout.
+        n = layer.weight.shape[0] // 2
+        swapped = torch.cat(
+            [layer.weight.data[n:].clone(), layer.weight.data[:n].clone()], dim=0
+        )
+        layer.weight.data.copy_(swapped)
+
+
+class _LayoutSwapLayer(torch.nn.Module):
+    """Layer with a sharded checkpoint weight loader + destructive process step."""
+
+    def __init__(self, half_size: int = 2):
+        super().__init__()
+        self._half_size = half_size
+        self.weight = torch.nn.Parameter(
+            torch.zeros(2 * half_size, dtype=torch.float32)
+        )
+
+        def shard_loader(param, loaded_weight, shard_id):
+            if shard_id == "w1":
+                param.data[:half_size].copy_(loaded_weight)
+            else:
+                assert shard_id == "w3"
+                param.data[half_size:].copy_(loaded_weight)
+
+        self.weight.weight_loader = shard_loader
+        self.quant_method = _LayoutSwapMethod()
+
+
+class _LayoutSwapModel(torch.nn.Module):
+    """Tiny model with a single `_LayoutSwapLayer` for regression testing."""
+
+    def __init__(self):
+        super().__init__()
+        self.layer = _LayoutSwapLayer()
+
+    def load_weights(self, weights):
+        loaded = set()
+        for name, value, shard_id in weights:
+            assert name == "layer.weight"
+            self.layer.weight.weight_loader(
+                self.layer.weight, value, shard_id=shard_id
+            )
+            loaded.add(name)
+        return loaded
+
+
+def test_make_load_weights_safe_for_reload_is_idempotent():
+    """Re-wrapping `model.load_weights` is a no-op.
+
+    Guards against accumulating layers of `initialize_layerwise_reload`
+    indirection if a loader's `load_model` is invoked more than once on
+    the same model instance.
+    """
+    model = _LayoutSwapModel()
+    record_metadata_for_reloading(model)
+
+    make_load_weights_safe_for_reload(model, model_config=None)
+    wrapped_once = model.load_weights
+    assert getattr(wrapped_once, "_vllm_safe_reload_wrapped", False)
+
+    make_load_weights_safe_for_reload(model, model_config=None)
+    assert model.load_weights is wrapped_once
+
+
+def test_load_weights_idempotent_under_destructive_process_step():
+    """Regression test for https://github.com/vllm-project/vllm/issues/42821.
+
+    Calling `model.load_weights` a second time with the same checkpoint must
+    not silently corrupt parameters whose `process_weights_after_loading`
+    rewrites their layout. Without the wrapper, the second invocation writes
+    checkpoint-format bytes into the swapped-layout buffer and the layer
+    drifts away from its post-init state on every reload.
+    """
+    model = _LayoutSwapModel()
+    record_metadata_for_reloading(model)
+
+    # Initial load: checkpoint provides [w1=(1,2), w3=(3,4)] which yields
+    # the checkpoint-format buffer [1, 2, 3, 4]. The layout swap then
+    # produces the kernel-format buffer [3, 4, 1, 2].
+    initial_weights = [
+        ("layer.weight", torch.tensor([1.0, 2.0]), "w1"),
+        ("layer.weight", torch.tensor([3.0, 4.0]), "w3"),
+    ]
+    model.load_weights(iter(initial_weights))
+    model.layer.quant_method.process_weights_after_loading(model.layer)
+
+    post_init_state = model.layer.weight.data.clone()
+    assert torch.equal(post_init_state, torch.tensor([3.0, 4.0, 1.0, 2.0]))
+
+    # Without the wrapper, a second `load_weights` writes raw [1, 2, 3, 4]
+    # into the swapped-layout buffer, leaving it inconsistent with the
+    # kernel's expected [3, 4, 1, 2] layout.
+    make_load_weights_safe_for_reload(model, model_config=None)
+    model.load_weights(iter(initial_weights))
+
+    assert torch.equal(model.layer.weight.data, post_init_state), (
+        f"Reload corrupted parameter layout: got {model.layer.weight.data}, "
+        f"expected {post_init_state}"
+    )
+
+    # Idempotency across many reloads.
+    for _ in range(3):
+        model.load_weights(iter(initial_weights))
+        assert torch.equal(model.layer.weight.data, post_init_state)
+
+
+def test_safe_reload_wrapper_preserves_kernel_storage_address():
+    """The wrapper preserves the parameter's storage `data_ptr` across reload.
+
+    This is critical for captured CUDA graphs in RL weight-update loops.
+    """
+    model = _LayoutSwapModel()
+    record_metadata_for_reloading(model)
+
+    initial_weights = [
+        ("layer.weight", torch.tensor([1.0, 2.0]), "w1"),
+        ("layer.weight", torch.tensor([3.0, 4.0]), "w3"),
+    ]
+    model.load_weights(iter(initial_weights))
+    model.layer.quant_method.process_weights_after_loading(model.layer)
+    storage_before = model.layer.weight.untyped_storage().data_ptr()
+
+    make_load_weights_safe_for_reload(model, model_config=None)
+    model.load_weights(iter(initial_weights))
+    storage_after = model.layer.weight.untyped_storage().data_ptr()
+
+    assert storage_before == storage_after, (
+        "Wrapper must preserve parameter storage address across reload."
+    )
+
+
+def test_safe_reload_wrapper_finalizes_on_loader_exception():
+    """If the inner `load_weights` raises, the wrapper still runs finalize.
+
+    The `finally` branch must call `finalize_layerwise_reload` so that
+    per-layer `info` is reset; otherwise the next `load_weights` call would
+    short-circuit `initialize_layerwise_reload` (which skips layers whose
+    `info.can_load()` is already True), causing wedged reload state.
+    """
+    model = _LayoutSwapModel()
+    record_metadata_for_reloading(model)
+
+    initial_weights = [
+        ("layer.weight", torch.tensor([1.0, 2.0]), "w1"),
+        ("layer.weight", torch.tensor([3.0, 4.0]), "w3"),
+    ]
+    model.load_weights(iter(initial_weights))
+    model.layer.quant_method.process_weights_after_loading(model.layer)
+    post_init_state = model.layer.weight.data.clone()
+
+    make_load_weights_safe_for_reload(model, model_config=None)
+
+    class _ExplodingIter:
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            raise RuntimeError("simulated checkpoint read failure")
+
+    with pytest.raises(RuntimeError, match="simulated checkpoint read failure"):
+        model.load_weights(_ExplodingIter())
+
+    # The next successful reload must produce the same post-init state,
+    # i.e. the wrapper recovers cleanly from the exception (`info` was
+    # reset by the `finally`-clause finalize so the second reload runs
+    # the full pipeline rather than short-circuiting).
+    model.load_weights(iter(initial_weights))
+    assert torch.equal(model.layer.weight.data, post_init_state), (
+        f"Wrapper failed to recover after exception: "
+        f"got {model.layer.weight.data}, expected {post_init_state}"
+    )
+
+
 @pytest.mark.parametrize(
     "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]
 )

vllm/model_executor/model_loader/base_loader.py

@@ -9,7 +9,10 @@
 from vllm.config import ModelConfig, VllmConfig
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
-from vllm.model_executor.model_loader.reload import finalize_layerwise_processing
+from vllm.model_executor.model_loader.reload import (
+    finalize_layerwise_processing,
+    make_load_weights_safe_for_reload,
+)
 from vllm.model_executor.model_loader.utils import (
     initialize_model,
     process_weights_after_loading,
@@ -79,6 +82,17 @@ def load_model(
 
             process_weights_after_loading(model, model_config, target_device)
 
+            # Make subsequent direct invocations of `model.load_weights`
+            # (e.g. from external RL frameworks performing in-place weight
+            # updates over `collective_rpc`) idempotent on a live model.
+            # Without this, MoE backends that rewrite the parameter layout
+            # in `process_weights_after_loading` (FlashInfer CUTLASS /
+            # TRT-LLM) silently corrupt subsequent forward output because
+            # the persisted per-expert `weight_loader` writes raw
+            # checkpoint-format bytes into the kernel-layout buffer. See
+            # https://github.com/vllm-project/vllm/issues/42821.
+            make_load_weights_safe_for_reload(model, model_config)
+
         return model.eval()
 
 

vllm/model_executor/model_loader/gguf_loader.py

@@ -15,6 +15,9 @@
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
+from vllm.model_executor.model_loader.reload import (
+    make_load_weights_safe_for_reload,
+)
 from vllm.model_executor.model_loader.utils import (
     initialize_model,
     process_weights_after_loading,
@@ -450,4 +453,10 @@ def load_model(
             self.load_weights(model, model_config)
 
             process_weights_after_loading(model, model_config, target_device)
+
+            # See `BaseModelLoader.load_model` for rationale.
+            # Required for parity with the default loader so that direct
+            # `model.load_weights` calls after init are safe on GGUF models
+            # too. https://github.com/vllm-project/vllm/issues/42821
+            make_load_weights_safe_for_reload(model, model_config)
         return model

vllm/model_executor/model_loader/reload/__init__.py

@@ -22,6 +22,7 @@
     "initialize_layerwise_reload",
     "finalize_layerwise_processing",
     "finalize_layerwise_reload",
+    "make_load_weights_safe_for_reload",
     "set_torchao_reload_attrs",
     "support_quantized_model_reload_from_hp_weights",
 ]
@@ -30,6 +31,7 @@
     finalize_layerwise_processing,
     finalize_layerwise_reload,
     initialize_layerwise_reload,
+    make_load_weights_safe_for_reload,
     record_metadata_for_reloading,
 )
 from .torchao_decorator import (

vllm/model_executor/model_loader/reload/layerwise.py

@@ -37,6 +37,7 @@
     "initialize_layerwise_reload",
     "finalize_layerwise_processing",
     "finalize_layerwise_reload",
+    "make_load_weights_safe_for_reload",
 ]
 
 
@@ -397,3 +398,88 @@ def _place_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo):
         layer.register_parameter(name, param)
     for name, buffer in buffers.items():
         layer.register_buffer(name, buffer)
+
+
+# Attribute used to mark a wrapped `model.load_weights` so we can detect and
+# avoid double-wrapping. We also expose the original callable for callers that
+# want to opt out of the layerwise reload (e.g. tests that pre-wrap manually).
+_SAFE_RELOAD_WRAPPED_ATTR = "_vllm_safe_reload_wrapped"
+_SAFE_RELOAD_ORIGINAL_ATTR = "_vllm_original_load_weights"
+
+
+def make_load_weights_safe_for_reload(
+    model: torch.nn.Module, model_config: ModelConfig | None
+) -> None:
+    """Make ``model.load_weights`` safe to invoke on an already-initialized model.
+
+    The first call to ``model.load_weights`` is made by the model loader before
+    ``process_weights_after_loading`` runs and is therefore correct by
+    construction: weight loaders write checkpoint-format bytes into freshly
+    created checkpoint-format buffers, and the kernel-layout transform is
+    applied exactly once afterwards.
+
+    Subsequent direct calls (e.g. from RL frameworks doing in-place weight
+    updates over ``collective_rpc``) are not idempotent on backends that
+    rewrite the parameter layout in ``process_weights_after_loading`` — for
+    example, the FlashInfer CUTLASS and FlashInfer TRT-LLM unquantized MoE
+    backends apply ``swap_w13_to_w31`` (and, for TRT-LLM, an additional block
+    permutation) to ``layer.w13_weight`` once at engine init. The per-expert
+    ``weight_loader`` attribute is preserved across that transform via
+    ``replace_parameter``, so a second ``load_weights`` call routes raw
+    ``[w1; w3]`` checkpoint bytes into a buffer that the kernel reads as
+    ``[w3; w1]`` (or block-permuted). The forward output silently collapses
+    into multilingual subword soup. See
+    https://github.com/vllm-project/vllm/issues/42821.
+
+    This helper installs a wrapper around ``model.load_weights`` that, on
+    every invocation after the initial load, runs the call through the
+    layerwise reload pipeline (``initialize_layerwise_reload`` /
+    ``finalize_layerwise_reload``) — the same pipeline that
+    :meth:`GPUModelRunner.reload_weights` already uses. That pipeline
+    restores parameters to their checkpoint-format storage (via captured
+    meta tensors), replays the weight loaders into a fresh buffer, re-runs
+    each layer's ``process_weights_after_loading`` (re-applying the
+    kernel-layout transform against the freshly loaded weights), and finally
+    copies the result back into the original kernel-layout parameter storage
+    so captured CUDA graphs remain valid.
+
+    The wrapper is a no-op if ``model.load_weights`` is already wrapped (so
+    callers nesting through ``reload_weights`` do not pay extra cost; the
+    inner ``initialize_layerwise_reload`` short-circuits for layers already
+    in the loadable state, and the inner ``finalize_layerwise_reload`` is a
+    no-op once the outer one has already reset per-layer info).
+
+    Args:
+        model: The fully-loaded model whose ``load_weights`` should be wrapped.
+            ``record_metadata_for_reloading`` must already have been called on
+            this model (which is the case after ``initialize_model``).
+        model_config: ``ModelConfig`` to forward to ``finalize_layerwise_reload``.
+            Required so attention layers can re-run their
+            ``process_weights_after_loading(dtype)`` finalize step. Pass
+            ``None`` only when the model has no attention layers (test paths).
+    """
+    original_load_weights = getattr(model, "load_weights", None)
+    if original_load_weights is None:
+        return
+
+    if getattr(original_load_weights, _SAFE_RELOAD_WRAPPED_ATTR, False):
+        return
+
+    @wraps(original_load_weights)
+    def safe_reload_load_weights(*args, **kwargs):
+        # Nesting note: when this wrapper is invoked from within
+        # `GPUModelRunner.reload_weights` (which already calls
+        # `initialize_layerwise_reload` itself), the inner call here is a
+        # no-op because `initialize_layerwise_reload` skips layers whose
+        # `info.can_load()` is already True. Symmetrically, the outer
+        # finalize seen by `reload_weights` becomes a no-op because the
+        # inner finalize below has already reset every layer's info.
+        initialize_layerwise_reload(model)
+        try:
+            return original_load_weights(*args, **kwargs)
+        finally:
+            finalize_layerwise_reload(model, model_config)
+
+    setattr(safe_reload_load_weights, _SAFE_RELOAD_WRAPPED_ATTR, True)
+    setattr(safe_reload_load_weights, _SAFE_RELOAD_ORIGINAL_ATTR, original_load_weights)
+    model.load_weights = safe_reload_load_weights  # type: ignore[assignment]