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

dimapihtar · 2026-03-18T20:44:33Z

/ok to test d2cfee3

sbak5 · 2026-03-18T21:14:00Z

https://github.com/NVIDIA/Megatron-LM/actions/runs/23266155298/job/67649083947?pr=3899
It's supposed to raise error during writing but this injection may not work with some issue. @dimapihtar

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

dimapihtar · 2026-03-18T23:25:36Z

/ok to test c5c86f5

sbak5 · 2026-03-19T06:29:22Z

megatron/core/dist_checkpointing/strategies/fully_parallel.py


    @debug_time("FullyParallelLoadStrategyWrapper.load", logger)
-    def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> StateDict:
+    def load(


Let's remove any load related changes. We don't have anything yet for loading.

@sbak5 load calss _get_filesystem_reader: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/dist_checkpointing/strategies/torch.py#L807

which uses CachedMetadataFileSystemReader: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/dist_checkpointing/strategies/torch.py#L766 so we need to import it properly in respect to async_strategy.

hmm but I guess we can take it from common state dict
let me see

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

This reverts commit 48c0d95.

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

dimapihtar · 2026-03-19T14:38:00Z

/ok to test ddc37b5

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

dimapihtar · 2026-03-19T14:48:52Z

/ok to test 2c85bea

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

dimapihtar · 2026-03-19T16:22:15Z

/ok to test ffe53e1

This reverts commit c3dabfb.

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

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

copy-pr-bot bot temporarily deployed to test March 18, 2026 20:45 Inactive

fix unit tests

c5c86f5

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

copy-pr-bot bot temporarily deployed to test March 18, 2026 23:26 Inactive

sbak5 reviewed Mar 19, 2026

View reviewed changes

dimapihtar and others added 13 commits March 19, 2026 03:59

avoid async_strategy param at serialization.load()

0f81f06

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

fix unit test

5ceef34

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

fix unit test

57c422c

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

move warning

9f690f0

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

fix unit test

9a36493

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

fix unit test

48c0d95

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

Revert "fix unit test"

67703f2

This reverts commit 48c0d95.

disable async_save

c8ed36c

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

fix unit test

8030764

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

fix typo

a0591c8

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

Merge branch 'main' into sbak/ckpt_migrate

64f4280

disable async_save

1f4979e

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

fix unit test

ddc37b5

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

fix warning

2c85bea

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

copy-pr-bot bot temporarily deployed to test March 19, 2026 14:49 Inactive

fix unit test

ffe53e1

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

copy-pr-bot bot temporarily deployed to test March 19, 2026 16:23 Inactive

dimapihtar and others added 3 commits March 20, 2026 05:03

Revert "update nvrx version"

6833d87

This reverts commit c3dabfb.

Merge branch 'main' into sbak/ckpt_migrate

e26b2d2

update nvrx version

2a92e71

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

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
dimapihtar wants to merge 51 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 18, 2026

Uh oh!

sbak5 commented Mar 18, 2026

Uh oh!

dimapihtar commented Mar 18, 2026

Uh oh!

sbak5 Mar 19, 2026

Uh oh!

dimapihtar Mar 19, 2026

Uh oh!

dimapihtar Mar 19, 2026

Uh oh!

dimapihtar commented Mar 19, 2026

Uh oh!

dimapihtar commented Mar 19, 2026

Uh oh!

dimapihtar commented Mar 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 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 18, 2026

Uh oh!

sbak5 commented Mar 18, 2026

Uh oh!

dimapihtar commented Mar 18, 2026

Uh oh!

sbak5 Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

dimapihtar Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

dimapihtar Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

dimapihtar commented Mar 19, 2026

Uh oh!

dimapihtar commented Mar 19, 2026

Uh oh!

dimapihtar commented Mar 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dimapihtar commented Mar 17, 2026 •

edited

Loading