[Main] Partial CUDA Graph support for EP Overlap by Wohox · Pull Request #2184 · NVIDIA/Megatron-LM

Wohox · 2025-11-10T08:57:15Z

Based on #1920

What does this PR do ?

EP Overlap brings extra cpu overhead and may cause GPU bubble during execution, partial CUDA graph helps release cpu pressure within the selected scope. This PR supports partial CUDA graph for EP Overlap, the supported scopes are attn, moe_router, moe_preprocess (moe and mlp are not supported).

Usage
To enable this feature, refer to the following example:

--overlap-moe-expert-parallel-comm \
--delay-wgrad-compute \
--cuda-graph-scope attn moe_router moe_preprocess \

(--delay-wgrad-compute is optional)

Correctness
The loss value can be bitwise aligned when enable & disable partial CUDA graph.

This PR is for main branch and the PR for dev branch is in PR2168

Dependencies
enable CUDA graph with delay-wgrad-compute relies on TE PR:

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
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

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

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.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

copy-pr-bot · 2025-11-10T08:57:19Z

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.

kvareddy · 2025-11-13T09:47:50Z

@fanshiqing @jiemingz can you please take a look at this MR?

jiemingz · 2025-11-13T22:15:34Z

it looks like this include changes from #1920 , is that supposed to be merged first and this be rebased?

Wohox · 2025-11-17T05:38:10Z

it looks like this include changes from #1920 , is that supposed to be merged first and this be rebased?

@jiemingz Yes, but rebase can happen later, since this MR requires #1920.

Wohox · 2025-12-02T06:14:46Z

/ok to test 125fa43

lhb8125

LGTM, thanks!

jiemingz

LGTM!

Wohox · 2026-01-15T02:26:05Z

/ok to test 9eababb

Wohox · 2026-01-15T02:33:16Z

@ericharper Can you help review this PR, still needs approval from NeMo and GPT group, thanks~

Phlip79 · 2026-01-15T19:21:27Z

/ok to test ecbaff4

Phlip79 · 2026-01-15T20:05:32Z

/ok to test c3219e6

Phlip79 · 2026-01-15T22:57:05Z

/ok to test e01fcab

Wohox · 2026-01-16T02:42:44Z

/ok to test e180d4d

Wohox requested review from a team as code owners November 10, 2025 08:57

Wohox changed the title ~~Pingtian/support cuda graph for ep overlap main~~ [Draft][Main] Partial CUDA Graph support for EP Overlap Nov 10, 2025

Wohox mentioned this pull request Nov 10, 2025

[Dev] Partial CUDA Graph support for EP Overlap #2168

Merged

6 tasks

Wohox requested review from a team as code owners November 13, 2025 09:07

Wohox changed the title ~~[Draft][Main] Partial CUDA Graph support for EP Overlap~~ [Main] Partial CUDA Graph support for EP Overlap Nov 13, 2025

jiemingz self-assigned this Nov 13, 2025

Wohox added module: moe Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. labels Nov 17, 2025

copy-pr-bot Bot temporarily deployed to nemo-ci December 2, 2025 06:14 Inactive

ko3n1g added this to the Core 0.16 milestone Dec 2, 2025

copy-pr-bot Bot temporarily deployed to nemo-ci December 2, 2025 06:15 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci December 2, 2025 06:15 Failure

copy-pr-bot Bot temporarily deployed to test December 2, 2025 06:15 Inactive

ericharper added the complexity: high label Dec 8, 2025

Wohox and others added 2 commits January 5, 2026 00:33

[Dev] Partial CUDA Graph support for EP Overlap (#2168)

2fe52d4

fix for main branch

0509ed1

format

c08dcee

lhb8125 reviewed Jan 13, 2026

View reviewed changes

Comment thread megatron/core/transformer/transformer_layer.py

jiemingz reviewed Jan 13, 2026

View reviewed changes

Comment thread megatron/core/transformer/transformer_layer.py Outdated

jiemingz reviewed Jan 13, 2026

View reviewed changes

Comment thread megatron/core/transformer/module.py Outdated

lhb8125 approved these changes Jan 14, 2026

View reviewed changes

Wohox and others added 2 commits January 13, 2026 18:35

refactor

677bd59

Merge branch 'main' into pingtian/support_cuda_graph_for_ep_overlap_main

293ebb4

jiemingz approved these changes Jan 14, 2026

View reviewed changes

kvareddy approved these changes Jan 14, 2026

View reviewed changes

yanring approved these changes Jan 15, 2026

View reviewed changes

Merge branch 'main' into pingtian/support_cuda_graph_for_ep_overlap_main

9eababb

copy-pr-bot Bot temporarily deployed to nemo-ci January 15, 2026 02:26 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci January 15, 2026 02:26 Failure

copy-pr-bot Bot temporarily deployed to nemo-ci January 15, 2026 02:26 Inactive

copy-pr-bot Bot temporarily deployed to test January 15, 2026 02:27 Inactive

ericharper approved these changes Jan 15, 2026

View reviewed changes

ericharper enabled auto-merge January 15, 2026 18:22

Merge branch 'main' into pingtian/support_cuda_graph_for_ep_overlap_main

ecbaff4

Merge branch 'main' into pingtian/support_cuda_graph_for_ep_overlap_main

c3219e6

Merge branch 'main' into pingtian/support_cuda_graph_for_ep_overlap_main

e01fcab

Wohox and others added 2 commits January 16, 2026 10:18

Merge branch 'main' into pingtian/support_cuda_graph_for_ep_overlap_main

f39e82a

fix ut

e180d4d

sbhavani mentioned this pull request Mar 23, 2026

[ROADMAP][2026 Q1] Megatron Core Roadmap #4003

Open

Conversation

Wohox commented Nov 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

copy-pr-bot Bot commented Nov 10, 2025

Uh oh!

kvareddy commented Nov 13, 2025

Uh oh!

jiemingz commented Nov 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Wohox commented Nov 17, 2025

Uh oh!

Wohox commented Dec 2, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lhb8125 left a comment

Choose a reason for hiding this comment

Uh oh!

jiemingz left a comment

Choose a reason for hiding this comment

Uh oh!

Wohox commented Jan 15, 2026

Uh oh!

Wohox commented Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Phlip79 commented Jan 15, 2026

Uh oh!

Phlip79 commented Jan 15, 2026

Uh oh!

Phlip79 commented Jan 15, 2026

Uh oh!

Wohox commented Jan 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

Wohox commented Nov 10, 2025 •

edited

Loading

(Step 1): Add PR label `Expert Review`

jiemingz commented Nov 13, 2025 •

edited

Loading

Wohox commented Jan 15, 2026 •

edited

Loading