Fix fused MLA and MTP checkpointing for Megatron-FSDP. by rapatel · Pull Request #6271 · NVIDIA/Megatron-LM

rapatel · 2026-08-04T22:23:46Z

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

What does this PR do?

Fixes fused MLA and MTP checkpointing for Megatron-FSDP.

Aside from the newly added unit tests, this PR was also tested manually for save/load parity check for the fused MLA + MTP path:

Ran a checkpoint save/load round trip on 2 GPUs through the real code in this PR (preprocess_fsdp_dtensor_state_dict -> the MLA/MTP handlers -> split_fused_fsdp_param), with a real torch.distributed.checkpoint save and load and the new validate_fsdp_dtensor_model_load. All checks passed.
The model reproduces: a fused linear_qkv_down_proj with asymmetric sections, an input_layernorm absorbed as layer_norm_weight, and an MTP layer whose inner layer is mtp_model_layer in the model but transformer_layer on disk, and which itself contains a fused attention, so the MLA-before-MTP handler ordering is covered.

The 3 main checks were:

Unfused checkpoint -> fused model
- 16 of 16 tensors bit-identical to what was saved, 0 mismatches across ranks
- loss 1.2154645919799805 vs 1.2154645919799805, delta 0.000e+00, bit-identical
Fused -> fused round trip (destination model starts from a different seed)
- 16 of 16 on-disk tensors bit-identical
- 13 of 13 weights the forward pass reads bit-identical
- loss 1.1673169136047363 before save and after load, bit-identical
- all gradients bit-identical, e.g. linear_qkv_down_proj.weight gradient norm 0.16416169703006744 before and after
Negative control, identical load with the MLA handler disabled
- CheckpointingException raised, naming the 6 unloadable fused parameters
- with --dist-ckpt-strictness ignore_all, the same load silently keeps those 6 at their initial values, which is the behavior that broke convergence before this change

Additionally, we ran a 64-node (256 GB300 GPUs) DSv3 convergence run with these changes, which ran as expected within the target samples-to-converge range for the target GBS.

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

copy-pr-bot · 2026-08-04T22:23:49Z

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-08-04T22:23:56Z

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.

rapatel · 2026-08-04T22:35:04Z

/ok to test aa30d8f.

cspades

I think it LGTM (there's a lot of code but a good bit is familiar code). Have 1 comment on arguments, and also an ask that you check we at least have loss and grad parity loading a non-trivial checkpoint with the model this PR is supporting, thanks!

cspades · 2026-08-04T22:47:34Z

+    strict_fsdp_dtensor_model_load: Literal[
+        "assume_ok_unexpected",
+        "log_unexpected",
+        "log_all",
+        "raise_unexpected",
+        "raise_all",
+        "return_unexpected",
+        "return_all",
+        "ignore_all",
+    ] = "raise_unexpected"
+    """Determine handling of model weights an FSDP DTensor partial load cannot supply.
+    Check StrictHandling docs for flags meaning. Only applies when strict_fsdp_dtensor_load is
+    False, where the underlying load would otherwise skip such weights silently and leave them
+    at their initialized values. Unlike strict_fsdp_dtensor_load this covers only the model
+    sections, so optimizer and RNG state that is intentionally absent does not trip it. Weights
+    the checkpoint holds but this rank does not need are always fine, as expected under
+    pipeline, tensor and expert parallelism. Ignored when finetune is set."""
+
    dist_ckpt_strictness: Literal[


Can we just pipe the existing dist_ckpt_strictness argument to FSDP?

cspades · 2026-08-04T22:51:18Z

+def _intersect_slice(s1, s2):
+    """Intersection of two step-1 slices, or an empty slice when they do not overlap."""
+    start = max(s1.start, s2.start)
+    stop = min(s1.stop, s2.stop)
+    return slice(0, 0) if start >= stop else slice(start, stop)


There's a similar function in the SwiGLU transforms, I think this can be used for that.

rapatel · 2026-08-05T01:09:01Z

/ok to test 6fee9a5.

rapatel · 2026-08-06T22:24:50Z

@cspades Please see updated description of this PR - it includes the details from the save/load parity testing. Thanks.

Signed-off-by: Ritesh Patel <ripatel@nvidia.com>

…e helpers. Signed-off-by: Ritesh Patel <ripatel@nvidia.com>

rapatel · 2026-08-11T23:04:12Z

/ok to test 1515dc7.

…nit test. Signed-off-by: Ritesh Patel <ripatel@nvidia.com>

gautham-kollu · 2026-08-19T19:42:47Z

/ok to test 6069334

nemo-automation-bot · 2026-08-20T01:19:32Z

🔄 Merge queue validation started!

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

nemo-automation-bot · 2026-08-20T07:47:40Z

🔄 Merge queue validation started!

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

Signed-off-by: Ritesh Patel <ripatel@nvidia.com> Co-authored-by: gautham-kollu <gkollu@nvidia.com> Signed-off-by: Kezhi Kong <kezhik@kezhik-mlt.client.nvidia.com>

Signed-off-by: Ritesh Patel <ripatel@nvidia.com> Co-authored-by: gautham-kollu <gkollu@nvidia.com>

rapatel requested a review from a team as a code owner August 4, 2026 22:23

rapatel requested a review from a team August 4, 2026 22:23

rapatel requested review from a team as code owners August 4, 2026 22:23

svcnvidia-nemo-ci marked this pull request as draft August 4, 2026 22:23

nvidia-megatron-lm-release-bot Bot added the community-request label Aug 4, 2026

rapatel assigned rapatel, cspades and gautham-kollu Aug 4, 2026

rapatel marked this pull request as ready for review August 4, 2026 22:34

svcnvidia-nemo-ci requested a review from a team August 4, 2026 22:35

github-actions Bot added the complexity: medium label Aug 4, 2026

gautham-kollu requested review from ahmadki and asolergi-nv August 4, 2026 22:40

cspades reviewed Aug 4, 2026

View reviewed changes

bvolpato mentioned this pull request Aug 5, 2026

[ckpt] fix: Remap fused MLA and MTP keys for FSDP DTensor loads. NVIDIA-NeMo/Megatron-Bridge#5315

Merged

3 of 5 tasks

rapatel enabled auto-merge August 5, 2026 01:08

chtruong814 removed the community-request label Aug 5, 2026

cspades approved these changes Aug 6, 2026

View reviewed changes

gautham-kollu added the Final Review PR is in the "final review" stage label Aug 10, 2026

rapatel added 2 commits August 11, 2026 16:03

Fix fused MLA and MTP checkpointing for Megatron-FSDP.

f728d12

Signed-off-by: Ritesh Patel <ripatel@nvidia.com>

Reuse dist_ckpt_strictness for the FSDP load check and share the slic…

1515dc7

…e helpers. Signed-off-by: Ritesh Patel <ripatel@nvidia.com>

rapatel force-pushed the ripatel/fsdp_chkpt_fix_main branch from 6fee9a5 to 1515dc7 Compare August 11, 2026 23:03

Cover the fused MLA and MTP fsdp_dtensor split with a Megatron-FSDP u…

a4f7722

…nit test. Signed-off-by: Ritesh Patel <ripatel@nvidia.com>

rapatel force-pushed the ripatel/fsdp_chkpt_fix_main branch from e61006c to a4f7722 Compare August 19, 2026 18:03

gautham-kollu added the Final Review PR is in the "final review" stage label Aug 19, 2026

asolergi-nv approved these changes Aug 19, 2026

View reviewed changes

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

jiemingz approved these changes Aug 19, 2026

View reviewed changes

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

Merge branch 'main' into ripatel/fsdp_chkpt_fix_main

6069334

copy-pr-bot Bot temporarily deployed to public August 19, 2026 19:51 Inactive

copy-pr-bot Bot temporarily deployed to test August 19, 2026 19:53 Inactive

copy-pr-bot Bot temporarily deployed to public August 19, 2026 19:58 Inactive

copy-pr-bot Bot temporarily deployed to public August 19, 2026 20:00 Inactive

copy-pr-bot Bot temporarily deployed to public August 19, 2026 20:04 Inactive

copy-pr-bot Bot temporarily deployed to public August 19, 2026 20:21 Inactive

maanug-nv approved these changes Aug 19, 2026

View reviewed changes

rapatel added this pull request to the merge queue Aug 19, 2026

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

rapatel added this pull request to the merge queue Aug 20, 2026

Merged via the queue into NVIDIA:main with commit 93924d5 Aug 20, 2026
188 checks passed

rapatel deleted the ripatel/fsdp_chkpt_fix_main branch August 20, 2026 09:41

This was referenced Aug 27, 2026

Megatron-LM Monthly Activity Report — August 2026 (Final, through August 31, 2026) #6921

Open

Megatron-LM Monthly Activity Report — August 2026 #7006

Open

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

Fix fused MLA and MTP checkpointing for Megatron-FSDP. (NVIDIA#6271)

82b079d

Signed-off-by: Ritesh Patel <ripatel@nvidia.com> Co-authored-by: gautham-kollu <gkollu@nvidia.com>

Conversation

rapatel commented Aug 4, 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 Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

Uh oh!

rapatel commented Aug 4, 2026

Uh oh!

cspades 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!

cspades Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

rapatel Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

cspades Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

rapatel Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

rapatel commented Aug 5, 2026

Uh oh!

rapatel commented Aug 6, 2026

Uh oh!

rapatel commented Aug 11, 2026

Uh oh!

gautham-kollu commented Aug 19, 2026

Uh oh!

nemo-automation-bot Bot commented Aug 20, 2026

Uh oh!

Uh oh!

nemo-automation-bot Bot commented Aug 20, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

rapatel commented Aug 4, 2026 •

edited

Loading

cspades left a comment •

edited

Loading