fix(inference): restore NemotronH mixer.D after vLLM 0.22 layerwise reload (#2714)

Author: hallerite

src/prime_rl/inference/patches.py

@@ -18,7 +18,6 @@ def transformers_v5_compat():
     _patch_qwen35_lora()
     _patch_lora_key_prefix()
     monkey_patch_deep_gemm_silu_mul_quant_int64()
-    monkey_patch_vllm_layerwise_reload_alias_buffers()
     monkey_patch_vllm_padded_input_scrub()
     monkey_patch_return_routed_experts_with_nixl_connector()
 
@@ -67,38 +66,6 @@ def _post_init(config: VllmConfig):
     logger.warning("Enabled vLLM routed-experts capture with NIXL connector patch.")
 
 
-def monkey_patch_vllm_layerwise_reload_alias_buffers():
-    # vLLM's layerwise reload materializes each buffer as an independent tensor
-    # and then copies it back into the original kernel storage. When a buffer
-    # aliases a parameter (e.g. NemotronH Mamba's mixer.conv_weights, a view of
-    # mixer.conv1d.weight), the buffer copy stamps garbage into the parameter's
-    # storage *after* the parameter has been correctly reloaded. Skip the copy
-    # for any buffer that shares storage with a parameter; _place_kernel_tensors
-    # re-registers the original view, which trivially reflects the parameter.
-    # Remove this patch once https://github.com/vllm-project/vllm/pull/42481 is
-    # included in the vLLM release we pin/use.
-    from vllm.logger import init_logger
-    from vllm.model_executor.model_loader.reload import layerwise as reload_layerwise
-
-    logger = init_logger(__name__)
-
-    def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: reload_layerwise.LayerReloadingInfo):
-        assert info.kernel_tensors is not None
-        parameters, buffers = info.kernel_tensors
-        param_storage_ptrs = {p.untyped_storage().data_ptr() for p in layer.parameters(recurse=True)}
-        for name, param in parameters.items():
-            param.data.copy_(getattr(layer, name))
-        for name, buffer in buffers.items():
-            if buffer.untyped_storage().data_ptr() in param_storage_ptrs:
-                continue
-            buffer.data.copy_(getattr(layer, name))
-
-        reload_layerwise._place_kernel_tensors(layer, info)
-
-    reload_layerwise._copy_and_restore_kernel_tensors = _copy_and_restore_kernel_tensors
-    logger.warning("Enabled vLLM layerwise reload alias-buffer patch.")
-
-
 @triton.jit
 def _silu_mul_per_token_group_quant_fp8_colmajor_int64_kernel(
     y_ptr,

src/prime_rl/inference/vllm/worker/nccl.py

@@ -24,6 +24,37 @@
 
 logger = init_logger("vllm.inference.vllm.worker_nccl")
 
+# NemotronH mixer.D is dropped by vLLM 0.22's layerwise online-reload path (left uninitialized).
+_RELOAD_CORRUPTED_SUFFIXES = (".mixer.D",)
+
+
+def _restore_reload_corrupted_params(model: Module, received: dict[str, torch.Tensor]) -> None:
+    """Work around a vLLM 0.22 layerwise-reload bug for NemotronH.
+
+    The online reload drops the weight load for every Mamba layer's ``mixer.D`` (the SSD skip
+    connection), leaving it as uninitialized ``empty_strided`` memory -- it reads back as garbage
+    (NaN/inf) and the logits go NaN after a weight update. The received broadcast value is correct,
+    so restore D from it via the param's own ``weight_loader``. Remove once the upstream bug is fixed.
+    """
+
+    def _layer_key(name: str) -> str:
+        index = name.find("layers.")
+        return name[index:] if index >= 0 else name
+
+    received_by_key = {_layer_key(name): tensor for name, tensor in received.items()}
+    for name, param in model.named_parameters():
+        if not name.endswith(_RELOAD_CORRUPTED_SUFFIXES):
+            continue
+        tensor = received_by_key.get(_layer_key(name))
+        if tensor is None:
+            continue
+        tensor = tensor.to(device=param.device)
+        weight_loader = getattr(param, "weight_loader", None)
+        if weight_loader is not None:
+            weight_loader(param, tensor)
+        elif tensor.shape == param.shape:
+            param.data.copy_(tensor.to(param.dtype))
+
 
 def receive_integer(communicator: PyNcclCommunicator) -> int:
     """Receive an integer from the trainer master rank using NCCL communicator."""
@@ -148,9 +179,20 @@ def update_weights_from_path(self, weight_dir: str) -> None:
             update_mla_absorbed_weights(model)
             return
 
+        # vLLM 0.22's layerwise reload drops NemotronH mixer.D's weight load (see
+        # _restore_reload_corrupted_params). Capture the correct received value to restore after.
+        received_reload_fix: dict[str, torch.Tensor] = {}
+
+        def _capture_reload_fix(weights):
+            for name, tensor in weights:
+                if name.endswith(_RELOAD_CORRUPTED_SUFFIXES):
+                    received_reload_fix[name] = tensor.detach().to("cpu", copy=True)
+                yield name, tensor
+
         load_weights_checkpoint_layerwise(
             model,
-            state_iter,
+            _capture_reload_fix(state_iter),
             self.model_runner.model_config,
             self.vllm_config,
         )
+        _restore_reload_corrupted_params(model, received_reload_fix)