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.

megatron/core/transformer/cuda_graphs.py

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

megatron/core/transformer/cuda_graphs.py

+        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

megatron/core/transformer/module.py

megatron/core/transformer/cuda_graphs.py

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-02-14T07:18:02Z

/ok to test 9dc24c7

shifangx · 2026-02-14T07:39:35Z

/ok to test 2c8b58e

shifangx · 2026-02-20T02:25:07Z

/ok to test ad6d2d4

shifangx · 2026-02-25T03:54:38Z

Hi, @yashaswikarnati , do you have any more comments?

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

Hi, @NVIDIA/mcore-oncall , can you help to review this pr?

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

/ok to test b714709

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

megatron/core/transformer/cuda_graphs.py Show resolved Hide resolved

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

megatron/core/transformer/module.py Outdated Show resolved Hide resolved

jiemingz reviewed Feb 2, 2026

View reviewed changes

megatron/core/transformer/cuda_graphs.py Outdated Show resolved Hide resolved

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

M4 leftover for TE cuda graph NVIDIA-NeMo/Megatron-Bridge#2372

Open

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

Open

6 tasks

copy-pr-bot bot temporarily deployed to nemo-ci February 14, 2026 07:18 Inactive

copy-pr-bot bot had a problem deploying to nemo-ci February 14, 2026 07:18 Failure

copy-pr-bot bot temporarily deployed to nemo-ci February 14, 2026 07:18 Inactive

m4 leftover for TE cuda graph

cdc9ca2

shifangx force-pushed the shifang/cuda_graph_m4 branch from 9dc24c7 to cdc9ca2 Compare February 14, 2026 07:36

Merge branch 'main' into shifang/cuda_graph_m4

2c8b58e

copy-pr-bot bot temporarily deployed to nemo-ci February 14, 2026 07:39 Inactive

copy-pr-bot bot had a problem deploying to nemo-ci February 14, 2026 07:39 Failure

copy-pr-bot bot temporarily deployed to nemo-ci February 14, 2026 07:39 Inactive

copy-pr-bot bot temporarily deployed to test February 14, 2026 07:40 Inactive

shifangx mentioned this pull request Feb 14, 2026

[main] Add TE CUDA Graph Support for Vision Encoder #3293

Merged

6 tasks

jiemingz approved these changes Feb 17, 2026

View reviewed changes

Merge branch 'main' into shifang/cuda_graph_m4

ad6d2d4

copy-pr-bot bot temporarily deployed to test February 20, 2026 02:25 Inactive

yaox12 added Final Review PR is in the "final review" stage and removed Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. labels Feb 26, 2026

jaredcasper approved these changes Feb 28, 2026

View reviewed changes

Merge branch 'main' into shifang/cuda_graph_m4

b714709

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Mar 19, 2026

shifangx enabled auto-merge March 19, 2026 02:04

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

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!

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 Feb 14, 2026

Uh oh!

shifangx commented Feb 14, 2026

Uh oh!

shifangx commented Feb 20, 2026

Uh oh!

shifangx commented Feb 25, 2026

Uh oh!

shifangx commented Mar 19, 2026

Uh oh!

shifangx commented Mar 19, 2026

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