diff --git a/verl/models/mcore/mtp_patch.py b/verl/models/mcore/mtp_patch.py index 649e2ad9da6..a8b989b1ac4 100644 --- a/verl/models/mcore/mtp_patch.py +++ b/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, diff --git a/verl/models/mcore/patch.py b/verl/models/mcore/patch.py index 7e9167872e0..4b6b5a326a0 100644 --- a/verl/models/mcore/patch.py +++ b/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 diff --git a/verl/utils/vllm/vllm_quant_utils.py b/verl/utils/vllm/vllm_quant_utils.py index 04524847c9d..aaade858b5d 100644 --- a/verl/utils/vllm/vllm_quant_utils.py +++ b/verl/utils/vllm/vllm_quant_utils.py @@ -30,6 +30,7 @@ import importlib.metadata import logging +from contextlib import nullcontext from dataclasses import dataclass, field from typing import Any @@ -37,11 +38,15 @@ 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():