[megatron] fix: forward mhc_multistream to MTP and skip activation reclaim for MTP checkpoints

verl/models/mcore/mtp_patch.py

@@ -22,11 +22,7 @@
 import torch
 from megatron.core import parallel_state, tensor_parallel
 from megatron.core.models.gpt.gpt_model import GPTModel
-from megatron.core.transformer.multi_token_prediction import (
-    MTPLossAutoScaler,
-    MTPLossLoggingHelper,
-    roll_tensor,
-)
+from megatron.core.transformer.multi_token_prediction import MTPLossAutoScaler, MTPLossLoggingHelper, roll_tensor
 
 try:
     from megatron.core.transformer.multi_token_prediction import process_mtp_loss as _process_mtp_loss
@@ -93,6 +89,7 @@ def _megatron_gptmodel_postprocess(
     output_processor=None,
     output_processor_context=None,
     is_spec_decode=None,
+    mhc_multistream=None,
 ):
     """Postprocesses decoder hidden states to generate logits or compute loss.
 
@@ -111,6 +108,8 @@ def _megatron_gptmodel_postprocess(
             self.mtp._forward_has_padding_mask = "padding_mask" in signature(self.mtp.forward).parameters
         if self.mtp._forward_has_padding_mask:
             mtp_kwargs["padding_mask"] = padding_mask
+        if mhc_multistream is not None:
+            mtp_kwargs["mhc_multistream"] = mhc_multistream
         hidden_states = self.mtp(
             input_ids=input_ids,
             position_ids=position_ids,

verl/models/mcore/patch.py

@@ -642,14 +642,29 @@ def patch_backward(ctx, *args):
             for inp in detached_inputs
         )
         cur_stream = torch.cuda.current_stream()
-        # Release original input and grad tensors
-        for t in detached_inputs:
-            if isinstance(t, torch.Tensor) and t.requires_grad:
-                t.record_stream(cur_stream)
-                t.untyped_storage().resize_(0)
-                if t.grad is not None:
-                    t.grad.record_stream(cur_stream)
-                    t.grad.untyped_storage().resize_(0)
+        # Release saved input/grad tensors (MoE residual-memory leak fix, commit 04df110c).
+        # Skip MTP layer checkpoints: their saved ``hidden_states`` aliases the decoder
+        # output (a ``torch.chunk`` view), and MTP backward runs before the decoder
+        # backward, so ``resize_(0)`` would truncate storage the decoder still needs
+        # -> async CUDA illegal-memory-access.
+        is_mtp_checkpoint = False
+        run_fn = getattr(ctx, "run_function", None)
+        for cell in getattr(run_fn, "__closure__", None) or ():
+            try:
+                obj = cell.cell_contents
+            except ValueError:
+                continue
+            if obj.__class__.__name__ == "MultiTokenPredictionLayer":
+                is_mtp_checkpoint = True
+                break
+        if not is_mtp_checkpoint:
+            for t in detached_inputs:
+                if isinstance(t, torch.Tensor) and t.requires_grad:
+                    t.record_stream(cur_stream)
+                    t.untyped_storage().resize_(0)
+                    if t.grad is not None:
+                        t.grad.record_stream(cur_stream)
+                        t.grad.untyped_storage().resize_(0)
         # ctx.saved_tensors = None
         return (None, None) + grads
 

verl/utils/vllm/vllm_quant_utils.py

@@ -30,18 +30,23 @@
 
 import importlib.metadata
 import logging
+from contextlib import nullcontext
 from dataclasses import dataclass, field
 from typing import Any
 
 import torch
 from packaging import version
 
 try:
-    from vllm.model_executor.layers.fused_moe.layer import FusedMoE
     from vllm.model_executor.layers.linear import LinearBase
 except ImportError as e:
     raise ImportError("FP8 quantization not available") from e
 
+try:
+    from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+except ImportError:
+    FusedMoE = None
+
 from verl.utils.kernel.fp8_kernel import scaled_fp8_blockwise
 from verl.utils.vllm.vllm_fp4_utils import (
     is_deepseek_v4_model,
@@ -91,10 +96,12 @@ def is_fp8_model(vllm_config):
 # ``nn.Module`` class. ``FusedMoE`` is now a factory *function* that builds a
 # ``MoERunner``, and the fused expert weight tensors (``w13_weight`` /
 # ``w2_weight``) moved onto a ``RoutedExperts`` submodule owned by the runner
-# (i.e. ``experts`` -> ``experts.routed_experts``). Resolve the concrete module
-# classes once so ``isinstance`` checks keep working across vLLM versions --
-# calling ``isinstance(x, FusedMoE)`` when ``FusedMoE`` is a function raises
-# ``TypeError: isinstance() arg 2 must be a type``.
+# (i.e. ``experts`` -> ``experts.routed_experts``). vLLM 0.26.1 removed the
+# ``FusedMoE`` name entirely. Resolve the concrete module classes once so
+# ``isinstance`` checks keep working across vLLM versions -- calling
+# ``isinstance(x, FusedMoE)`` when ``FusedMoE`` is a function raises
+# ``TypeError: isinstance() arg 2 must be a type``, and ``FusedMoE`` may be
+# ``None`` when the name no longer exists.
 if isinstance(FusedMoE, type):
     # vLLM < 0.24: ``FusedMoE`` is itself the expert-weight-holding module.
     _MOE_STOP_CLASSES = (FusedMoE,)
@@ -380,9 +387,17 @@ def load_quanted_weights(weights, model_runner, is_drafter=False):
         if hasattr(param, "subclass_type"):
             param.orig_type = param.__class__
             param.__class__ = param.subclass_type
-    # Finally load the weights into vllm
+    # Finally load the weights into vllm.
+    # MTP completeness check assumes a single full-checkpoint load; in RL refit
+    # weights arrive bucketed, so disable it (like vLLM's own NCCL/IPC engines).
+    # ``nullcontext`` covers older vLLM without the guard.
+    try:
+        from vllm.model_executor.model_loader.mtp_validation import disable_mtp_completeness_check
+    except ImportError:
+        disable_mtp_completeness_check = nullcontext
     try:
-        loaded_params = model.load_weights(weights_quantized)
+        with disable_mtp_completeness_check():
+            loaded_params = model.load_weights(weights_quantized)
     finally:
         # Undo the type change above to the original type
         for name, param in model.named_parameters():