M4 leftover for TE cuda graph by shifangx · Pull Request #3137 · NVIDIA/Megatron-LM

shifangx · 2026-01-29T07:37:48Z

What does this PR do ?

fix this issue #3135
Global pg groups, such as _TENSOR_MODEL_PARALLEL_GROUP , in parallel_state.py will be deprecated after M4.
So we need pass pg_collection while init TECudaGraphHelper, and use pg group in pg_collection later.

⚠️ 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

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

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!

(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 · 2026-01-29T07:37:52Z

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.

yashaswikarnati · 2026-01-29T18:51:06Z

+        self.dp_cp_group = self.pg_collection.dp_cp
+        self.pp_group = self.pg_collection.pp
+        from megatron.core.pipeline_parallel.p2p_communication import P2PCommunicator
+        self.p2p_communicator = P2PCommunicator(pp_group=self.pp_group, config=self.config)


what is this self.config? looks like p2p communicator needs ModelParallelConfig, not TransformerConfig

The type of this config is megatron.bridge.models.qwen_vl.qwen3_vl_provider.Qwen3VLMoEModelProvider.

isnt this cuda graphs file agnostic to a particular model. I'm wondering how ould this be related to qwen3vlrprovider, also given this is megatron LM and not bridge

This pr is not for particular model, it is for M4.
In order to support M4, we need use pg groups in pg_collection everywhere, including in this cuda_graph file.

Yeah I get that. My question/confusion was what's the type of config we pass here to P2Pcommunicator. Is it TransformerConfig or ModelParallelConfig

The inheritance relationship is Qwen3VLModelProvider -> GPTModelProvider -> TransformerConfig -> ModelParallelConfig.

shifangx · 2026-02-14T07:04:14Z

cuda graph is very helpful for QWen3-VL training performance, and many customers want to try training QWen3-VL with both cuda graph and M4.
This pr is needed for this case.

Hi, @jiemingz, I have addressed reviewer's comments. Can you help to review my update?

If @jiemingz is out of office for Chinese New Year, is there any other expert can help to review this pr? @NVIDIA/core-adlr @NVIDIA/mcore-oncall

shifangx · 2026-02-14T07:10:01Z

/ok to test 81d8781

shifangx · 2026-03-19T02:05:42Z

/ok to test b714709

shifangx · 2026-04-15T18:02:43Z

/ok to test 50efc83

shifangx · 2026-04-16T15:44:34Z

/ok to test bbe34d4 bbe34d4

shifangx · 2026-04-16T16:26:42Z

/ok to test 0d2cd53

shifangx · 2026-04-16T16:35:59Z

/ok to test 13e9033

shifangx · 2026-04-16T16:39:26Z

/ok to test 556cb65

copy-pr-bot · 2026-04-16T16:39:30Z

/ok to test 556cb65

@shifangx, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

shifangx · 2026-04-16T16:42:00Z

/ok to test 806d161

svcnvidia-nemo-ci · 2026-04-16T18:56:25Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/24528325068

svcnvidia-nemo-ci · 2026-04-16T19:58:25Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/24531041761

svcnvidia-nemo-ci · 2026-04-16T20:00:11Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/24531115334

svcnvidia-nemo-ci · 2026-04-16T20:26:30Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/24532272940

shifangx requested review from a team as code owners January 29, 2026 07:37

ko3n1g requested a review from a team January 29, 2026 07:38

shifangx changed the title ~~fix issue with ccuda graph and m4~~ fix issue with cuda graph and m4 Jan 29, 2026

shifangx force-pushed the shifang/cuda_graph_m4 branch 3 times, most recently from bfce1c9 to 98d5fcf Compare January 29, 2026 13:47

yaoyu-33 added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Jan 29, 2026

yaoyu-33 approved these changes Jan 29, 2026

View reviewed changes

yashaswikarnati reviewed Jan 29, 2026

View reviewed changes

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

yashaswikarnati reviewed Jan 29, 2026

View reviewed changes

jiemingz added the complexity: low label Jan 29, 2026

jiemingz self-requested a review January 29, 2026 19:06

shifangx force-pushed the shifang/cuda_graph_m4 branch 2 times, most recently from 14c15b0 to 844848b Compare January 30, 2026 13:30

jiemingz reviewed Feb 2, 2026

View reviewed changes

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

jiemingz reviewed Feb 2, 2026

View reviewed changes

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

shifangx mentioned this pull request Feb 7, 2026

Need TECudaGraphHelper support M4 #3135

Open

shifangx changed the title ~~fix issue with cuda graph and m4~~ m4 leftover for cuda graph Feb 13, 2026

shifangx force-pushed the shifang/cuda_graph_m4 branch from 5210c51 to 9e676af Compare February 13, 2026 14:07

shifangx changed the title ~~m4 leftover for cuda graph~~ m4 leftover for TE cuda graph Feb 13, 2026

shifangx force-pushed the shifang/cuda_graph_m4 branch 2 times, most recently from 3ba536f to d470fca Compare February 13, 2026 14:31

shifangx mentioned this pull request Feb 13, 2026

[QWen3VL] Adding CUDA Graph Support for LLM Sub-module NVIDIA-NeMo/Megatron-Bridge#2372

Closed

5 tasks

shifangx changed the title ~~m4 leftover for TE cuda graph~~ M4 leftover for TE cuda graph Feb 13, 2026

shifangx mentioned this pull request Feb 14, 2026

[dev] Add TE CUDA Graph Support for Vision Encoder for dev branch #3294

Closed

6 tasks

copy-pr-bot Bot temporarily deployed to test March 19, 2026 02:06 Inactive

Phlip79 removed the request for review from a team April 15, 2026 16:18

copy-pr-bot Bot temporarily deployed to test April 15, 2026 18:03 Inactive

shifangx disabled auto-merge April 15, 2026 18:17

shifangx force-pushed the shifang/cuda_graph_m4 branch from 50efc83 to 5e04c06 Compare April 16, 2026 15:43

shifangx mentioned this pull request Apr 16, 2026

m4 leftover for TE cuda graph NVIDIA-NeMo/Megatron-Bridge#3358

Merged

5 tasks

shifangx enabled auto-merge April 16, 2026 15:59

shifangx force-pushed the shifang/cuda_graph_m4 branch from 0d2cd53 to d801e5c Compare April 16, 2026 16:35

m4 leftover for TE cuda graph

556cb65

shifangx force-pushed the shifang/cuda_graph_m4 branch from 13e9033 to 556cb65 Compare April 16, 2026 16:38

Merge branch 'main' into shifang/cuda_graph_m4

806d161

copy-pr-bot Bot temporarily deployed to test April 16, 2026 16:43 Inactive

shifangx added this pull request to the merge queue Apr 16, 2026

Merged via the queue into NVIDIA:main with commit 01eb7e8 Apr 16, 2026
64 of 65 checks passed

shifangx deleted the shifang/cuda_graph_m4 branch April 16, 2026 21:28

Phlip79 pushed a commit to Phlip79/Megatron-LM that referenced this pull request Apr 17, 2026

M4 leftover for TE cuda graph (NVIDIA#3137)

f04e15e

Phlip79 mentioned this pull request Apr 17, 2026

M4 leftover for TE cuda graph (dev) #4369

Merged

Conversation

shifangx commented Jan 29, 2026 • 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 Jan 29, 2026

Uh oh!

Uh oh!

yashaswikarnati Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

shifangx Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

yashaswikarnati Feb 25, 2026

Choose a reason for hiding this comment

Uh oh!

shifangx Mar 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yashaswikarnati Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

shifangx Apr 16, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

shifangx commented Feb 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shifangx commented Feb 14, 2026

Uh oh!

shifangx commented Mar 19, 2026

Uh oh!

shifangx commented Apr 15, 2026

Uh oh!

shifangx commented Apr 16, 2026

Uh oh!

shifangx commented Apr 16, 2026

Uh oh!

shifangx commented Apr 16, 2026

Uh oh!

shifangx commented Apr 16, 2026

Uh oh!

copy-pr-bot Bot commented Apr 16, 2026

Uh oh!

shifangx commented Apr 16, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 16, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 16, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 16, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

shifangx commented Jan 29, 2026 •

edited

Loading

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

shifangx Mar 19, 2026 •

edited

Loading

shifangx commented Feb 14, 2026 •

edited

Loading