From 6f6599cdfd03e8f73d627033a5daa0fecaa0b3d8 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 15 Apr 2026 12:31:45 -0500 Subject: [PATCH 1/2] Fix inference graph override in RL flow --- megatron/rl/rl_utils.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 8a564315dc3..90d3250fb26 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -1973,8 +1973,11 @@ def megatron_rl_inference_mode( logger.debug(f"[{dist.get_rank()}] Entering inference mode") - # Change cudagraph scope for inference (empty list = full-layer capture) - model[0].config.cuda_graph_scope = [] + # Save original scope so we can restore it after inference + original_cuda_graph_scope = model[0].config.cuda_graph_scope + + # Set cudagraph scope and impl for inference + model[0].config.cuda_graph_scope = args.cuda_graph_scope model[0].config.cuda_graph_impl = "local" # If we get a lower precision wrapper, we go one object deeper. @@ -2031,14 +2034,8 @@ def megatron_rl_inference_mode( # Reset drop_and_pad leaked from inference decode set_decode_expert_padding(unwrap_model(model[0]), set_to=False) - # Restore partial capture cudagraph scope for training if this is MoE - if args.num_experts is not None: - model[0].config.cuda_graph_scope = [ - CudaGraphScope.mamba, - CudaGraphScope.attn, - CudaGraphScope.moe_router, - CudaGraphScope.moe_preprocess, - ] + # Restore the original cudagraph scope for training + model[0].config.cuda_graph_scope = original_cuda_graph_scope # Switch MoE layers to partial CUDA graph capture for training if args.rl_training_cuda_graphs and args.num_experts is not None: From dd20976ff88511021f6d61119e682a6653f1b3f7 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 15 Apr 2026 14:12:38 -0500 Subject: [PATCH 2/2] Fix the fix --- megatron/rl/rl_utils.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 90d3250fb26..4d018217cec 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -1973,10 +1973,7 @@ def megatron_rl_inference_mode( logger.debug(f"[{dist.get_rank()}] Entering inference mode") - # Save original scope so we can restore it after inference - original_cuda_graph_scope = model[0].config.cuda_graph_scope - - # Set cudagraph scope and impl for inference + # Set cudagraph scope for inference. model[0].config.cuda_graph_scope = args.cuda_graph_scope model[0].config.cuda_graph_impl = "local" @@ -2034,8 +2031,19 @@ def megatron_rl_inference_mode( # Reset drop_and_pad leaked from inference decode set_decode_expert_padding(unwrap_model(model[0]), set_to=False) - # Restore the original cudagraph scope for training - model[0].config.cuda_graph_scope = original_cuda_graph_scope + # Restore cudagraph scope for training. + # MoE partial capture requires specific scopes that aren't user-facing. + if args.num_experts is not None: + model[0].config.cuda_graph_scope = [ + CudaGraphScope.mamba, + CudaGraphScope.attn, + CudaGraphScope.moe_router, + CudaGraphScope.moe_preprocess, + ] + else: + model[0].config.cuda_graph_scope = [ + s for s in args.cuda_graph_scope if s != CudaGraphScope.full_iteration_inference + ] # Switch MoE layers to partial CUDA graph capture for training if args.rl_training_cuda_graphs and args.num_experts is not None: