[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