Adding NVRx as a dependency and keeping the current code base optionally by dimapihtar · Pull Request #3899 · NVIDIA/Megatron-LM

dimapihtar · 2026-03-17T11:47:16Z

What does this PR do ?

Adds NVRx async save routines and deprecates mcore's.
Introduces async_strategy param to make it configurable from the user perspective.
NVRx strategy is set as default async strategy. MCore's async strategy will be used if NVRx is not installed.
Still possible to use MCore's async strategy by setting --async-strategy mcore.
Keeps mcore's async strategy for backward compatibility but will be fully removed and subsituted with related nvrx solution in 1-2 releases.

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

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 the @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.

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.

Fix dequantize check to use `"dequantize" in type(ten).__dict__` instead of `hasattr(ten, "dequantize")`. The latter returns True for all torch.Tensor objects since dequantize is defined on the base class, causing RuntimeError on non-quantized tensors. The new check only matches TE subclasses (e.g. Float8Tensor, MXFP8Tensor) that define their own dequantize override.

This reverts commit a760b5e.

copy-pr-bot · 2026-03-17T11:47:20Z

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.

Signed-off-by: dimapihtar <dpihtar@gmail.com>

sbak5 · 2026-03-17T19:56:16Z

/claude review

Signed-off-by: dimapihtar <dpihtar@gmail.com>

dimapihtar · 2026-03-18T14:31:20Z

/ok to test cc6fd79

Signed-off-by: dimapihtar <dpihtar@gmail.com>

dimapihtar · 2026-03-30T21:47:34Z

/ok to test 7d7feb7

- Add `_get_async_calls_queue()` helper for lazy initialization of the async calls queue, defaulting async_strategy to "nvrx" if not set in args - Guard `maybe_finalize_async_save` and `is_empty_async_queue` against an uninitialized (`None`) queue to avoid crashes when no save has been scheduled - Use `getattr` with default for `async_strategy` and `use_persistent_ckpt_worker` to handle missing args gracefully - Revert removal of `async_save=True` from `test_basic_save_load_scenarios` and `test_local_ckpt_sharded_state_dict` parametrized tests, restoring async save coverage Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

sbak5 · 2026-03-31T00:01:19Z

/ok to test 68032da

svcnvidia-nemo-ci · 2026-03-31T05:55:48Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-03-31T06:10:26Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-03-31T07:10:48Z

🔄 Merge queue validation started!

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

… optionally (#3899)" This reverts commit 704c7ee.

…e optionally (#3899)" This reverts commit 925e9f8.

…code base optionally (#3899)"" This reverts commit 25c7636.

… optionally (#3899)" This reverts commit 704c7ee.

…e optionally (#3899)" This reverts commit 0057095.

sbak5 and others added 4 commits March 12, 2026 17:15

Lint applied

1587314

Import resiliency-ext async checkpointing

a1aad02

Revert "Handle quantized CUDA tensors in async checkpoint writer"

65bb819

This reverts commit a760b5e.

dimapihtar added 3 commits March 17, 2026 08:14

keep both nvrx & mcore async save strategies

814f2fb

Signed-off-by: dimapihtar <dpihtar@gmail.com>

rename variable

9fc5757

Signed-off-by: dimapihtar <dpihtar@gmail.com>

refactor get_async_strategy

c2559d9

Signed-off-by: dimapihtar <dpihtar@gmail.com>

sbak5 changed the title ~~Sbak/ckpt migrate~~ Adding NVRx as a dependency and keeping the current code base optionally Mar 17, 2026

dimapihtar and others added 7 commits March 18, 2026 06:16

pass async strategy to save/load strategy

1b9ad87

Signed-off-by: dimapihtar <dpihtar@gmail.com>

Merge branch 'main' into sbak/ckpt_migrate

55b6a00

fix imports

4fa312f

Signed-off-by: dimapihtar <dpihtar@gmail.com>

properly pass async_strategy to async_save

cb04e5a

Signed-off-by: dimapihtar <dpihtar@gmail.com>

properly pass async_strategy

247027f

Signed-off-by: dimapihtar <dpihtar@gmail.com>

properly pass async_strategy load

c3673fd

Signed-off-by: dimapihtar <dpihtar@gmail.com>

remove extra code

cc6fd79

Signed-off-by: dimapihtar <dpihtar@gmail.com>

dimapihtar marked this pull request as ready for review March 18, 2026 14:31

dimapihtar requested review from a team as code owners March 18, 2026 14:31

dimapihtar added the complexity: high label Mar 18, 2026

svcnvidia-nemo-ci requested a review from a team March 18, 2026 14:31

svcnvidia-nemo-ci added this to the Core 0.16 milestone Mar 18, 2026

dimapihtar added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Mar 18, 2026

svcnvidia-nemo-ci added complexity: medium and removed complexity: high labels Mar 18, 2026

dimapihtar requested a review from yaoyu-33 March 30, 2026 19:22

yaoyu-33 approved these changes Mar 30, 2026

View reviewed changes

raise error

5cef931

Signed-off-by: dimapihtar <dpihtar@gmail.com>

svcnvidia-nemo-ci removed the Final Review PR is in the "final review" stage label Mar 30, 2026

dimapihtar added 3 commits March 30, 2026 14:44

revert uv.lock

fd06594

Signed-off-by: dimapihtar <dpihtar@gmail.com>

Merge branch 'main' into sbak/ckpt_migrate

938880d

upgrade uv.lock

7d7feb7

Signed-off-by: dimapihtar <dpihtar@gmail.com>

copy-pr-bot Bot temporarily deployed to test March 30, 2026 21:48 Inactive

copy-pr-bot Bot temporarily deployed to test March 31, 2026 00:02 Inactive

sbak5 enabled auto-merge March 31, 2026 03:50

jaredcasper approved these changes Mar 31, 2026

View reviewed changes

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

sbak5 added this pull request to the merge queue Mar 31, 2026

jaredcasper removed this pull request from the merge queue due to a manual request Mar 31, 2026

ko3n1g added this pull request to the merge queue Mar 31, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Mar 31, 2026

ko3n1g added this pull request to the merge queue Mar 31, 2026

Merged via the queue into NVIDIA:main with commit 704c7ee Mar 31, 2026
133 of 135 checks passed

ko3n1g added a commit that referenced this pull request Apr 4, 2026

Revert "Adding NVRx as a dependency and keeping the current code base…

925e9f8

… optionally (#3899)" This reverts commit 704c7ee.

ko3n1g added a commit that referenced this pull request Apr 4, 2026

Reapply "Adding NVRx as a dependency and keeping the current code bas…

25c7636

…e optionally (#3899)" This reverts commit 925e9f8.

ko3n1g added a commit that referenced this pull request Apr 4, 2026

Revert "Reapply "Adding NVRx as a dependency and keeping the current …

fd88c5b

…code base optionally (#3899)"" This reverts commit 25c7636.

ko3n1g added a commit that referenced this pull request Apr 4, 2026

Revert "Adding NVRx as a dependency and keeping the current code base…

0057095

… optionally (#3899)" This reverts commit 704c7ee.

ko3n1g added a commit that referenced this pull request Apr 4, 2026

Reapply "Adding NVRx as a dependency and keeping the current code bas…

e51170c

…e optionally (#3899)" This reverts commit 0057095.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Adding NVRx as a dependency and keeping the current code base optionally#3899

Adding NVRx as a dependency and keeping the current code base optionally#3899
ko3n1g merged 69 commits intoNVIDIA:mainfrom
sbak5:sbak/ckpt_migrate

dimapihtar commented Mar 17, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Mar 17, 2026

Uh oh!

sbak5 commented Mar 17, 2026

Uh oh!

dimapihtar commented Mar 18, 2026

Uh oh!

dimapihtar commented Mar 30, 2026

Uh oh!

sbak5 commented Mar 31, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 31, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Mar 31, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Mar 31, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Conversation

dimapihtar commented Mar 17, 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: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Mar 17, 2026

Uh oh!

sbak5 commented Mar 17, 2026

Uh oh!

dimapihtar commented Mar 18, 2026

Uh oh!

dimapihtar commented Mar 30, 2026

Uh oh!

sbak5 commented Mar 31, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 31, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Mar 31, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Mar 31, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

dimapihtar commented Mar 17, 2026 •

edited

Loading