[Dev] Refactor CUDA graph API: decompose cuda_graph_scope into full_iteration impl, inference scope, and per-layer capture modules

[buptzyb]

What does this PR do ?

Main PR #4292

This PR decomposes the overloaded cuda_graph_scope field into three dedicated, semantically distinct concepts and cleans up naming throughout.

Problem

The old API overloaded --cuda-graph-scope with three unrelated concerns:

• --cuda-graph-scope full_iteration → full-iteration training graphs (a capture strategy, not a module)
• --cuda-graph-scope full_iteration_inference → block-owned inference graphs (inference ownership, not a module)
• --cuda-graph-scope attn mlp → per-layer capture regions (the actual intended use)

CudaGraphScope mixed iteration-level control flow with per-layer module selection. The three concepts have nothing in common and cannot be meaningfully combined in one field.

Solution

Four concrete changes:

1. full_iteration becomes its own cuda_graph_impl value.
   --cuda-graph-impl full_iteration replaces --cuda-graph-impl local --cuda-graph-scope full_iteration.

2. --inference-cuda-graph-scope is a new dedicated field.
   InferenceCudaGraphScope (none / layer / block) replaces full_iteration_inference in cuda_graph_scope. The default for --cuda-graph-impl local is layer (preserving prior behaviour).

3. CudaGraphScope is renamed to CudaGraphModule; cuda_graph_scope to cuda_graph_modules.
   With the two non-module values removed, the enum and field names now accurately reflect their purpose: selecting which per-layer modules to capture.

4. Normalization is centralized in cuda_graph_config.py.
   normalize_cuda_graph_modules, normalize_inference_cuda_graph_scope, and validate_deprecated_cuda_graph_modules_migration_inputs are shared between TransformerConfig.__post_init__ and validate_args, removing duplication and making the migration logic testable in isolation.

Backward compatibility

All deprecated inputs are still accepted and silently migrated at startup:

Old input	New equivalent
--enable-cuda-graph	--cuda-graph-impl local
--external-cuda-graph	--cuda-graph-impl transformer_engine
--cuda-graph-scope full_iteration	--cuda-graph-impl full_iteration
--cuda-graph-scope full_iteration_inference	--cuda-graph-impl local --inference-cuda-graph-scope block
--cuda-graph-scope attn mlp	--cuda-graph-modules attn mlp
from megatron.core.transformer.enums import CudaGraphScope	CudaGraphScope alias kept; use CudaGraphModule
TransformerConfig(cuda_graph_scope=...)	cuda_graph_scope field kept, deprecated; use cuda_graph_modules

Conflicting combinations (e.g. passing both a deprecated value and the new flag with an incompatible value) are rejected with a clear assertion.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Victarry]

/claude strict-review

[claude]

Code Review Summary

CRITICAL: 0 | IMPORTANT: 4 | SUGGESTION: 4

Key Findings

IMPORTANT:

1. Deprecated scope migration sets string instead of enum (transformer_config.py:932, arguments.py:248): The setattr(self, attr, value) migration path sets inference_cuda_graph_granularity to a string ('transformer_block') rather than the declared InferenceCudaGraphGranularity enum. This works today because normalize_inference_cuda_graph_granularity() is called afterwards, but creates a fragile order-dependent invariant. Recommend converting to the final type immediately or storing typed values in CUDA_GRAPH_SCOPE_DEPRECATIONS.

2. full_iteration impl allows cuda_graph_impl="none" in deprecated scope validation (cuda_graph_config.py:127): Passing only --cuda-graph-scope full_iteration (without explicit --cuda-graph-impl) now silently enables full-iteration CUDA graphs, whereas in the old API this combination would have been a no-op (since graphs were disabled). This is a subtle backward-incompatible behavior change in an edge case.

3. State leak in megatron_rl_inference_mode (rl_utils.py:1838-1843): inference_cuda_graph_granularity and cuda_graph_scope are overwritten on context entry but never restored on exit. Benign today since cuda_graph_impl is restored to "none", but leaves the config in an inconsistent state.

4. dynamic_engine.py runs warmup passes for full_iteration without capturing graphs (dynamic_engine.py:338-345): The warning does not mention that warmup forward passes still run. Combined with the matching inference/utils.py change that enables num_cuda_graphs for full_iteration, this burns GPU time and memory unnecessarily during inference with full_iteration impl.

Overall Assessment

Risk: Low-Medium. The PR's core design is sound — separating training capture scope from inference granularity into distinct, purpose-specific fields is the right architectural move. The backward compatibility migration is well-structured with centralized normalization in cuda_graph_config.py, proper deprecation warnings, and comprehensive test coverage for migration paths.

The findings above are not blocking — they're fragility/hygiene issues, not correctness bugs in the current code. The most actionable item is the string-vs-enum type inconsistency during migration, which could be fixed with a small change to the deprecation map or the setattr call sites.

The new documentation (cuda_graph.md) and test coverage are excellent additions.

[buptzyb]

/ok to test 145c047

[hxbai]

LGTM.

[buptzyb]

/ok to test d231704