Remove torch cuda empty cache by shurkat-nvidia · Pull Request #5678 · NVIDIA/Megatron-LM

shurkat-nvidia · 2026-07-06T17:51:12Z

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

What does this PR do?

The PR removes an expensive torch.cuda.empty_cache() call within save_checkpoint_and_time(). This call took multiple seconds on every checkpoint and also increased the size of the generate_state_dict() function by multiple seconds. To avoid OOM errors, the code also modifies the Nemotron3 Super script to 1) use SHM to avoid expensive rebuild_cuda_tensor() calls during the IPC handoff and 2) garbage collect the torch allocator at 80% GPU allocation with a lower overhead than a raw empty_cache() call, and to use larger blocks to prevent fragmentation in the first place.

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

(N/A) I have added relevant unit tests
(N/A) I have added relevant functional tests
(N/A) I have added proper typing to my code Typing guidelines
(N/A) 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.

copy-pr-bot · 2026-07-06T17:51:15Z

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.

asolergi-nv · 2026-07-21T14:54:45Z

/ok to test debb98a

cuichenx · 2026-08-10T18:23:07Z

This is a replay of the change reverted in #5366, is CI passing with this PR?

Phlip79 · 2026-08-10T18:54:34Z

This is a replay of the change reverted in #5366, is CI passing with this PR?

This test was recently updated in #6239. @shurkat-nvidia can you please check to see if the changes to the config are needed anymore?

torch.cuda.empty_cache() was added as a workaround for out-of-memory errors on some models. With async-shm checkpointing, some of the memory pressure is reduced. Remaining memory fragmentation and pressure can be reduced by properly configuring the torch allocator. Signed-off-by: Skand Hurkat <shurkat@nvidia.com>

- Update checkpointing to use nvrx with shm. - Provide PYTORCH_ALLOC_CONF environment variable to initiate garbage collection. Use larger blocks to avoid fragmentation due to splitting. Signed-off-by: Skand Hurkat <shurkat@nvidia.com>

shurkat-nvidia · 2026-08-10T21:06:48Z

This is a replay of the change reverted in #5366, is CI passing with this PR?

@cuichenx, @balasaajay ran the test manually a few weeks ago, it appeared to have addressed the OOM issue.

This test was recently updated in #6239. @shurkat-nvidia can you please check to see if the changes to the config are needed anymore?

@Phlip79, I rebased the branch to the latest main. Still need changes to checkpoint args and to the PYTORCH_ALLOC_CONF environment variable.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Skand Hurkat <shurkat@nvidia.com>

deepakn94 · 2026-08-10T21:33:20Z

/ok to test 0f9153b

svcnvidia-nemo-ci · 2026-08-13T16:15:16Z

🔄 Merge queue validation started!

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

Signed-off-by: Skand Hurkat <shurkat@nvidia.com> Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

shurkat-nvidia force-pushed the remove-torch-cuda-empty-cache branch 2 times, most recently from ddbde1f to da16039 Compare July 6, 2026 18:54

shurkat-nvidia marked this pull request as ready for review July 15, 2026 18:05

shurkat-nvidia requested a review from a team as a code owner July 15, 2026 18:05

shurkat-nvidia requested a review from a team July 15, 2026 18:05

asolergi-nv added the Run functional tests label Jul 21, 2026

copy-pr-bot Bot temporarily deployed to public July 21, 2026 14:55 Inactive

copy-pr-bot Bot temporarily deployed to test July 21, 2026 14:56 Inactive

copy-pr-bot Bot temporarily deployed to public July 21, 2026 14:58 Inactive

copy-pr-bot Bot temporarily deployed to public July 21, 2026 15:06 Inactive

Phlip79 added nemotron complexity: low labels Aug 7, 2026

cuichenx approved these changes Aug 10, 2026

View reviewed changes

deepakn94 mentioned this pull request Aug 10, 2026

Gate checkpoint-time GPU memory reclaim behind --ckpt-release-gpu-memory shurkat-nvidia/Megatron-LM#1

Closed

2 of 4 tasks

shurkat-nvidia added 2 commits August 10, 2026 13:51

Update N3-Super script: checkpointing, allocation

c7830f1

- Update checkpointing to use nvrx with shm. - Provide PYTORCH_ALLOC_CONF environment variable to initiate garbage collection. Use larger blocks to avoid fragmentation due to splitting. Signed-off-by: Skand Hurkat <shurkat@nvidia.com>

shurkat-nvidia force-pushed the remove-torch-cuda-empty-cache branch from debb98a to c7830f1 Compare August 10, 2026 21:01

deepakn94 reviewed Aug 10, 2026

View reviewed changes

Comment thread tests/functional_tests/test_cases/nemotron/nemotron3_super_release_gb200/model_config.yaml Outdated

shurkat-nvidia and others added 2 commits August 10, 2026 14:25

Update N3-Super PYTORCH_ALLOC_CONF to only set GC threshold

416dd0a

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Skand Hurkat <shurkat@nvidia.com>

Update N3-Super PYTORCH_CUDA_ALLOC_CONF to add GC threshold

0f9153b

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Skand Hurkat <shurkat@nvidia.com>

deepakn94 approved these changes Aug 10, 2026

View reviewed changes

svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Aug 10, 2026

copy-pr-bot Bot temporarily deployed to public August 10, 2026 21:33 Inactive

copy-pr-bot Bot temporarily deployed to test August 10, 2026 21:34 Inactive

copy-pr-bot Bot temporarily deployed to public August 10, 2026 21:37 Inactive

copy-pr-bot Bot temporarily deployed to public August 10, 2026 21:38 Inactive

copy-pr-bot Bot temporarily deployed to public August 10, 2026 21:49 Inactive

deepakn94 added this pull request to the merge queue Aug 13, 2026

Merged via the queue into NVIDIA:main with commit 2119719 Aug 13, 2026
185 checks passed

deepakn94 mentioned this pull request Aug 16, 2026

Instrument Megatron-LM with OpenTelemetry traces and metrics #6438

Merged

4 of 5 tasks

lhb8125 mentioned this pull request Aug 17, 2026

2026年7月 Megatron-LM 工作汇总 lhb8125/Megatron-LM#56

Closed

shurkat-nvidia deleted the remove-torch-cuda-empty-cache branch August 17, 2026 18:01

lhb8125 mentioned this pull request Aug 18, 2026

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

Open

ksivaman pushed a commit to ksivaman/Megatron-LM that referenced this pull request Sep 1, 2026

Remove torch cuda empty cache (NVIDIA#5678)

74af842

Signed-off-by: Skand Hurkat <shurkat@nvidia.com> Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

Navigation Menu

Repository navigation

Remove torch cuda empty cache - #5678

Remove torch cuda empty cache#5678
deepakn94 merged 4 commits into
NVIDIA:mainfrom
shurkat-nvidia:remove-torch-cuda-empty-cache

shurkat-nvidia commented Jul 6, 2026

Uh oh!

copy-pr-bot Bot commented Jul 6, 2026

Uh oh!

asolergi-nv commented Jul 21, 2026

Uh oh!

cuichenx commented Aug 10, 2026

Uh oh!

Phlip79 commented Aug 10, 2026

Uh oh!

shurkat-nvidia commented Aug 10, 2026 •

edited

Loading

Uh oh!

Uh oh!

deepakn94 commented Aug 10, 2026

Uh oh!

svcnvidia-nemo-ci commented Aug 13, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

shurkat-nvidia commented Jul 6, 2026

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 Jul 6, 2026

Uh oh!

asolergi-nv commented Jul 21, 2026

Uh oh!

cuichenx commented Aug 10, 2026

Uh oh!

Phlip79 commented Aug 10, 2026

Uh oh!

shurkat-nvidia commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

deepakn94 commented Aug 10, 2026

Uh oh!

svcnvidia-nemo-ci commented Aug 13, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

shurkat-nvidia commented Aug 10, 2026 •

edited

Loading