Optimize memory usage of partial CUDA graphs by jiemingz · Pull Request #5451 · NVIDIA/Megatron-LM

jiemingz · 2026-06-23T16:36:13Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do ?

⚠️ 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 @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

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 @NVIDIA/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"

When your PR is ready, click Ready for Review.
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.

Signed-off-by: Jieming Zhang <jiemingz@nvidia.com>

Signed-off-by: root <jiemingz@nvidia.com>

copy-pr-bot · 2026-06-23T16:36:16Z

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.

github-actions · 2026-06-23T16:36:34Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

Signed-off-by: root <jiemingz@nvidia.com>

jiemingz · 2026-06-25T20:05:52Z

/claude review

claude

Two bugs in the new make_weakref function (line 340 and the except block) — see inline comment. Otherwise the changes look correct.

Signed-off-by: Jieming Zhang <jiemingz@nvidia.com>

Phlip79

Do we have a unit or functional test to prevent memory regression? If not, should we?

jiemingz · 2026-07-01T04:39:09Z

Do we have a unit or functional test to prevent memory regression? If not, should we?
We do not beyond the usual checks in functional tests, however they are not accurate.

Indeed I am working on a test to track memory usage (as well as tracking the numerical correctness of graphs) but unfortunately measuring memory usage under partial CG isn't straighforward. The usual methods are not currently cleanly supported in pytorch- polling memory usage with cudagraphs pools withtorch.cuda.get_memory_allocated doesnt account for graph mem pools, and directly querying total device memory with get_mem_info is too unstable.
See this comment:
#5451 (comment)
I would suggest merging these optimizations first and revisiting how we catch regressions in a separate PR.

jiemingz · 2026-07-01T04:48:01Z

Do we have a unit or functional test to prevent memory regression? If not, should we?

I created an issue to create these tests:
#5594

jiemingz · 2026-07-01T19:23:10Z

/ok to test 0589635

svcnvidia-nemo-ci · 2026-07-02T03:10:33Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-07-02T14:47:16Z

🔄 Merge queue validation started!

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

jiemingz added 2 commits June 11, 2026 13:41

partial CG memory fixes

33ea587

Signed-off-by: Jieming Zhang <jiemingz@nvidia.com>

cleanup

97a9fa4

Signed-off-by: root <jiemingz@nvidia.com>

jiemingz requested review from a team as code owners June 23, 2026 16:36

jiemingz marked this pull request as draft June 23, 2026 16:36

svcnvidia-nemo-ci marked this pull request as draft June 23, 2026 16:36

deepakn94 changed the title ~~optimize partial cudagraph memory usage~~ Optimize memory usage of partial CUDA graphs Jun 23, 2026

jiemingz marked this pull request as ready for review June 23, 2026 17:15

deepakn94 approved these changes Jun 23, 2026

View reviewed changes

svcnvidia-nemo-ci added the complexity: medium label Jun 23, 2026

jiemingz commented Jun 23, 2026

View reviewed changes

Comment thread megatron/core/transformer/cuda_graphs.py

kvareddy approved these changes Jun 23, 2026

View reviewed changes

jiemingz requested review from fanshiqing and mathemakitten June 23, 2026 17:49

mathemakitten reviewed Jun 23, 2026

View reviewed changes

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

Comment thread megatron/core/transformer/transformer_layer.py

fanshiqing approved these changes Jun 24, 2026

View reviewed changes

Comment thread megatron/core/transformer/moe/token_dispatcher.py

Comment thread megatron/core/transformer/cuda_graphs.py

jiemingz force-pushed the jiemingz/cg_mem branch from 2976222 to b211685 Compare June 25, 2026 19:37

address comments

b2cbe34

Signed-off-by: root <jiemingz@nvidia.com>

jiemingz force-pushed the jiemingz/cg_mem branch from b211685 to b2cbe34 Compare June 25, 2026 19:42

Merge branch 'NVIDIA:main' into jiemingz/cg_mem

d57e702

claude Bot reviewed Jun 25, 2026

View reviewed changes

Comment thread megatron/core/transformer/cuda_graphs.py

claude Bot reviewed Jun 25, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to public June 30, 2026 17:07 Inactive

Phlip79 requested changes Jun 30, 2026

View reviewed changes

Comment thread megatron/core/models/hybrid/hybrid_block.py Outdated

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

move import

0589635

Signed-off-by: Jieming Zhang <jiemingz@nvidia.com>

Phlip79 reviewed Jul 1, 2026

View reviewed changes

Phlip79 approved these changes Jul 1, 2026

View reviewed changes

jiemingz enabled auto-merge July 1, 2026 19:14

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

jiemingz disabled auto-merge July 1, 2026 19:23

copy-pr-bot Bot temporarily deployed to public July 1, 2026 19:37 Inactive

copy-pr-bot Bot temporarily deployed to test July 1, 2026 19:40 Inactive

copy-pr-bot Bot temporarily deployed to public July 1, 2026 19:43 Inactive

copy-pr-bot Bot temporarily deployed to public July 1, 2026 19:48 Inactive

copy-pr-bot Bot temporarily deployed to public July 1, 2026 19:49 Inactive

copy-pr-bot Bot temporarily deployed to public July 1, 2026 19:59 Inactive

yaox12 added this pull request to the merge queue Jul 2, 2026

yaox12 removed this pull request from the merge queue due to a manual request Jul 2, 2026

jiemingz added this pull request to the merge queue Jul 2, 2026

Merged via the queue into NVIDIA:main with commit 5e4fe9b Jul 2, 2026
181 checks passed

svcnvidia-nemo-ci mentioned this pull request Jul 8, 2026

chore: nightly sync main into dev (08_07_2026) #5716

Closed

lhb8125 mentioned this pull request Aug 18, 2026

Megatron-LM Monthly Activity Report — July 2026 (Final, through July 31, 2026) #6607

Open

huthvincent mentioned this pull request Aug 25, 2026

Fix the strong-reference fallback in make_weakref #6856

Open

6 tasks

Uh oh!

Conversation

jiemingz commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Jun 23, 2026

Uh oh!

github-actions Bot commented Jun 23, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jiemingz commented Jun 25, 2026

Uh oh!

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Phlip79 left a comment

Choose a reason for hiding this comment

Uh oh!

jiemingz commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jiemingz commented Jul 1, 2026

Uh oh!

jiemingz commented Jul 1, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Jul 2, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

jiemingz commented Jun 23, 2026 •

edited

Loading

jiemingz commented Jul 1, 2026 •

edited

Loading