[fix] Use MSC for checking checkpoint existence by pavelgein · Pull Request #4251 · NVIDIA/Megatron-LM

pavelgein · 2026-04-10T12:08:00Z

What does this PR do ?

This PR fixed behaviour, when MSC path is used

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

copy-pr-bot · 2026-04-10T12:08:05Z

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-04-10T12:08:10Z

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.

pavelgein · 2026-04-13T04:18:40Z

@mcore-oncall

pavelgein · 2026-04-14T08:48:59Z

@mcore-oncall hi, could you please take a look?

pavelgein · 2026-04-16T15:32:14Z

hi @chtruong814 sorry to bother you, maybe you can help me get this merged?

pavelgein · 2026-04-27T16:23:20Z

@mcore-oncall hi, could you please take a look?

asolergi-nv

Thanks for the fix — verify_checkpoint was indeed missing the MSC branch that the rest of this file already has (core.py:63, validation.py:224 in verify_checkpoint_and_load_strategy, serialization.py:379, strategies/common.py:94,156), so loading from an msc://… URL would falsely report the checkpoint as missing.

A couple of small suggestions inline to keep the new block consistent with the surrounding conventions in this file.

Signed-off-by: Pavel Gein <pavel.gein@gmail.com>

Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

pavelgein · 2026-04-30T16:56:09Z

@asolergi-nv thanks for the review.

There is actually a test in the file tests/unit_tests/dist_checkpointing/test_msc.py, and it uses a MSC, but the problem with this test is that it uses only local paths, so it passes without fix.

pavelgein · 2026-05-13T08:16:59Z

Hi @asolergi-nv could you, please, help to land this PR? I don't know anyone else to ping...

ericharper · 2026-05-14T17:18:22Z

/ok to test 62e1b40

pavelgein · 2026-05-15T03:11:14Z

Tests failures seems unrelated

asolergi-nv · 2026-05-15T15:16:00Z

/ok to test 207cd98

svcnvidia-nemo-ci · 2026-05-16T03:49:53Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-05-16T04:30:16Z

🔄 Merge queue validation started!

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

* origin/main: (138 commits) Refactor CUDA graph API: decompose cuda_graph_scope into full_iteration impl, inference scope, and per-layer capture modules (NVIDIA#4292) Add high-priority A2A stream and HybridEP preprocessing SMs (NVIDIA#4694) add is_torch_min_version in fsdp src (NVIDIA#4812) [Main][feat] Support A2A Overlap for Megatron-FSDP (NVIDIA#3797) Reorder mtp_post_process after attention backward in 1F1B schedule plan (NVIDIA#4695) [fix] Use MSC for checking checkpoint existence (NVIDIA#4251) Combine GEMM + SwiGLU fused MLP PRs (3890, 4071, 4095, 4219, 4311, 4324) → main (NVIDIA#4636) Strengthen test_checkpoint to verify distributed checkpoint behavior (NVIDIA#4711) Disable MSC by default; opt in via --enable-msc (NVIDIA#4629) additional tests for nvrx (NVIDIA#4522) Update copy-pr-bot.yaml [skip ci] ci: tolerate git-gc race in /home/runner chown after checkout (NVIDIA#4808) Inference: Optimize Prefill Engine Steps for Nemotron (NVIDIA#4764) chore: Update nightly tests golden values (NVIDIA#4805) ci: Update workflow to use same commit for building docker image and running tests (NVIDIA#4787) Update owners (NVIDIA#4794) Bump nvidia-modelopt>=0.44.0 (NVIDIA#4803) fix tokenizers in respect to newer transformers (NVIDIA#4608) Use Protocols to type-check linear_proj submodules of Attention (NVIDIA#3434) Fix recompute checkpointing + training CGs (NVIDIA#3919) ... # Conflicts: # megatron/core/transformer/moe/moe_utils.py

Signed-off-by: Pavel Gein <pavel.gein@gmail.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

Signed-off-by: Pavel Gein <pavel.gein@gmail.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

Signed-off-by: Pavel Gein <pavel.gein@gmail.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

Signed-off-by: Pavel Gein <pavel.gein@gmail.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

Signed-off-by: Pavel Gein <pavel.gein@gmail.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>

pavelgein requested review from a team as code owners April 10, 2026 12:08

svcnvidia-nemo-ci marked this pull request as draft April 10, 2026 12:08

github-actions Bot added the community-request label Apr 10, 2026

pavelgein marked this pull request as ready for review April 10, 2026 12:09

svcnvidia-nemo-ci requested a review from a team April 10, 2026 12:09

chtruong814 added the needs-follow-up Issue needs follow-up label Apr 12, 2026

chtruong814 removed the needs-follow-up Issue needs follow-up label Apr 13, 2026

pavelgein mentioned this pull request Apr 13, 2026

[feature] Support MultiStorageClient (MSC) for FSDP checkpoints NVIDIA-NeMo/Megatron-Bridge#3261

Closed

pavelgein force-pushed the msc_support_in_load branch from 3600b8c to 9f57ac6 Compare April 14, 2026 07:18

chtruong814 added the needs-follow-up Issue needs follow-up label Apr 18, 2026

svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed needs-follow-up Issue needs follow-up labels Apr 21, 2026

svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Apr 27, 2026

asolergi-nv reviewed Apr 30, 2026

View reviewed changes

Comment thread megatron/core/dist_checkpointing/validation.py Outdated

Comment thread megatron/core/dist_checkpointing/validation.py Outdated

Comment thread megatron/core/dist_checkpointing/validation.py

svcnvidia-nemo-ci added waiting-on-customer Waiting on the original author to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Apr 30, 2026

pavelgein and others added 2 commits April 30, 2026 21:48

[fix] Use MSC for checking checkpoint existence

4798c70

Signed-off-by: Pavel Gein <pavel.gein@gmail.com>

style

62e1b40

Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

pavelgein force-pushed the msc_support_in_load branch from 450a621 to 62e1b40 Compare April 30, 2026 16:49

pavelgein requested a review from asolergi-nv May 5, 2026 09:43

svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed waiting-on-customer Waiting on the original author to respond labels May 7, 2026

jaredcasper approved these changes May 13, 2026

View reviewed changes

svcnvidia-nemo-ci added waiting-on-customer Waiting on the original author to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels May 13, 2026

ericharper approved these changes May 14, 2026

View reviewed changes

ericharper enabled auto-merge May 14, 2026 17:14

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

copy-pr-bot Bot temporarily deployed to test May 14, 2026 17:19 Inactive

svcnvidia-nemo-ci removed the waiting-on-customer Waiting on the original author to respond label May 15, 2026

Merge branch 'main' into msc_support_in_load

207cd98

copy-pr-bot Bot temporarily deployed to public May 15, 2026 15:16 Inactive

copy-pr-bot Bot temporarily deployed to test May 15, 2026 15:17 Inactive

ericharper added this pull request to the merge queue May 15, 2026

Merged via the queue into NVIDIA:main with commit f2070ee May 16, 2026
63 of 65 checks passed

svcnvidia-nemo-ci mentioned this pull request May 16, 2026

chore: nightly sync main into dev (16_05_2026) #4831

Closed

janEbert pushed a commit to janEbert/Megatron-LM that referenced this pull request Jun 2, 2026

[fix] Use MSC for checking checkpoint existence (NVIDIA#4251)

ab15e8a

Signed-off-by: Pavel Gein <pavel.gein@gmail.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

maanug-nv mentioned this pull request Jun 10, 2026

Reduce boilerplate around MultiStorageClient feature checks #5269

Merged

6 tasks

terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026

[fix] Use MSC for checking checkpoint existence (NVIDIA#4251)

d1562ff

Signed-off-by: Pavel Gein <pavel.gein@gmail.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

Uh oh!

Conversation

pavelgein commented Apr 10, 2026

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 Apr 10, 2026

Uh oh!

github-actions Bot commented Apr 10, 2026

Uh oh!

pavelgein commented Apr 13, 2026

Uh oh!

pavelgein commented Apr 14, 2026

Uh oh!

pavelgein commented Apr 16, 2026

Uh oh!

pavelgein commented Apr 27, 2026

Uh oh!

asolergi-nv left a comment • 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!

Uh oh!

Uh oh!

Uh oh!

pavelgein commented Apr 30, 2026

Uh oh!

pavelgein commented May 13, 2026

Uh oh!

ericharper commented May 14, 2026

Uh oh!

pavelgein commented May 15, 2026

Uh oh!

asolergi-nv commented May 15, 2026

Uh oh!

svcnvidia-nemo-ci commented May 16, 2026

Uh oh!

svcnvidia-nemo-ci commented May 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

asolergi-nv left a comment •

edited

Loading