Skip to content

Fix inference graph override in RL flow #4323

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tdene merged 2 commits into from
Apr 29, 2026
Merged

Fix inference graph override in RL flow #4323

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6f6599c
Fix inference graph override in RL flow
tdene Apr 15, 2026
dd20976
Fix the fix
tdene Apr 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions megatron/rl/rl_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1973,8 +1973,8 @@ 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 = []
# Set cudagraph scope 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.
Expand Down Expand Up @@ -2031,14 +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 partial capture cudagraph scope for training if this is MoE
# 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:
Expand Down
Loading