Fix Megatron-FSDP checkpoint save failure by shjwudp · Pull Request #2138 · NVIDIA/Megatron-LM

shjwudp · 2025-11-05T07:50:23Z

Fix Megatron-FSDP checkpoint save failure by ensuring _replace_param_with_distributed_if_needed is invoked.

Prior to this change, the following error messages would appear in the logs:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/lustre/fsw/coreai_devtech_all/jianbinc/playground/fix_fsdp_dtensor/Megatron-LM/pretrain_gpt.py", line 234, in <module>
[rank0]:     pretrain(
[rank0]:   File "/lustre/fsw/coreai_devtech_all/jianbinc/playground/fix_fsdp_dtensor/Megatron-LM/megatron/training/training.py", line 737, in pretrain
[rank0]:     iteration, num_floating_point_operations_so_far = train(
[rank0]:                                                       ^^^^^^
[rank0]:   File "/lustre/fsw/coreai_devtech_all/jianbinc/playground/fix_fsdp_dtensor/Megatron-LM/megatron/training/training.py", line 2459, in train
[rank0]:     should_exit = checkpoint_and_decide_exit(
[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/lustre/fsw/coreai_devtech_all/jianbinc/playground/fix_fsdp_dtensor/Megatron-LM/megatron/training/training.py", line 1875, in checkpoint_and_decide_exit
[rank0]:     save_checkpoint_and_time(
[rank0]:   File "/lustre/fsw/coreai_devtech_all/jianbinc/playground/fix_fsdp_dtensor/Megatron-LM/megatron/training/training.py", line 1743, in save_checkpoint_and_time
[rank0]:     save_checkpoint(
[rank0]:   File "/lustre/fsw/coreai_devtech_all/jianbinc/playground/fix_fsdp_dtensor/Megatron-LM/megatron/training/checkpointing.py", line 558, in save_checkpoint
[rank0]:     state_dict = preprocess_fsdp_dtensor_state_dict(args, state_dict, model[0])
[rank0]:                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/lustre/fsw/coreai_devtech_all/jianbinc/playground/fix_fsdp_dtensor/Megatron-LM/megatron/training/checkpointing.py", line 854, in preprocess_fsdp_dtensor_state_dict
[rank0]:     model_state_dict, optimizer_state_dict = handle_swiglu_in_state_dict(
[rank0]:                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/lustre/fsw/coreai_devtech_all/jianbinc/playground/fix_fsdp_dtensor/Megatron-LM/megatron/core/transformer/fsdp_dtensor_checkpoint.py", line 275, in handle_swiglu_in_state_dict
[rank0]:     weight_w, weight_v = split_swiglu_linear_fc1(
[rank0]:                          ^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/lustre/fsw/coreai_devtech_all/jianbinc/playground/fix_fsdp_dtensor/Megatron-LM/megatron/core/transformer/fsdp_dtensor_checkpoint.py", line 219, in split_swiglu_linear_fc1
[rank0]:     fsdp_slice = dist_param.megatron_fsdp_slice
[rank0]:                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: AttributeError: 'Parameter' object has no attribute 'megatron_fsdp_slice'

What does this PR do ?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share discuss a design-doc with the team.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
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

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

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.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

…with_distributed_if_needed is called

copy-pr-bot · 2025-11-05T07:50:27Z

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.

cspades

Megatron-LM checkpoints sub-modules that dodges the root module hook?

BoxiangW · 2025-11-05T19:08:34Z


    args.use_distributed_optimizer = ckpt_format != "torch_dcp"
    args.use_dist_ckpt = ckpt_format != "torch"
+    args.use_megatron_fsdp = ckpt_format == "fsdp_dtensor"


Is this line in use?

BoxiangW · 2025-11-05T19:09:32Z

            expected_ckpt_path = ckpt_dir / "mp_rank_00" / "model_optim_rng.pt"
        elif ckpt_format == "torch_dcp":
            expected_ckpt_path = ckpt_dir / ".metadata"
+        elif ckpt_format == "fsdp_dtensor":


Minor issue, this if condition is running the same code as elif ckpt_format == "torch_dcp":, maybe we can combine them

@BoxiangW If you want to apply this structurally across the codebase, ruff can be used to automatically lint and fix these duplicated if else codepaths.

cc @ko3n1g, is there any plan on switching to ruff instead? I know that all other repos are using it already

We have some pretty carefully chosen, mostly non-opinionated ruff rules chosen in PyTorch: https://github.com/pytorch/pytorch/blob/a74fe75c450277eb88a95c764e8b0a664a550a86/pyproject.toml#L130

This specific one we did not enable because of breakpoint debugging / traceback reasons, but on a smaller codebase like this more easy to enforce opinionated lints.

Do you want to create a separate issue for this? It seems would be very helpful for improving Megatron-Core's formatting

More for catching bugs lol, I've found several bugs with ruff that should have not made it past the CI into Megatron.

Do you mean issue for this specific rule or enabling ruff more broadly.

Got it, current lint is more for formatting and import checks.
I meant enabling ruff more broadly actually, I think it would be great to have.

hey @Skylion007, we had to put the plans to switch to ruff on ice for a bit due to the team impact when handling a code-base of this size. We hope that we can pick those plan up again end of the year. Thanks for the tooling you shared!

I enabled most of the rules for most of PyTorch, which is a way larger codebase if you want to discuss how to enable it incrementally @ko3n1g

Skylion007 · 2025-11-05T20:30:21Z

Likely closes: #2035

shjwudp · 2025-11-10T06:54:44Z

/ok to test 98158d9

Fix Megatron-FSDP checkpoint save failure by ensuring _replace_param_…

6f0606d

…with_distributed_if_needed is called

shjwudp requested review from a team as code owners November 5, 2025 07:50

ko3n1g approved these changes Nov 5, 2025

View reviewed changes

cspades approved these changes Nov 5, 2025

View reviewed changes

BoxiangW reviewed Nov 5, 2025

View reviewed changes

code format

2e89e57

shjwudp added Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. Run tests Run functional tests labels Nov 10, 2025

Merge branch 'main' into fix_fsdp_dtensor

98158d9

shjwudp requested review from a team November 10, 2025 05:35

copy-pr-bot Bot temporarily deployed to nemo-ci November 10, 2025 06:54 Inactive

ko3n1g added this to the Core 0.16 milestone Nov 10, 2025

copy-pr-bot Bot temporarily deployed to nemo-ci November 10, 2025 06:55 Inactive

copy-pr-bot Bot temporarily deployed to test November 10, 2025 06:55 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci November 10, 2025 06:57 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci November 10, 2025 06:58 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci November 10, 2025 07:15 Inactive

Uh oh!

Conversation

shjwudp commented Nov 5, 2025 • 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): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

copy-pr-bot Bot commented Nov 5, 2025

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Skylion007 Nov 5, 2025 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Skylion007 Nov 8, 2025 • 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!

Skylion007 commented Nov 5, 2025

Uh oh!

shjwudp commented Nov 10, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

shjwudp commented Nov 5, 2025 •

edited

Loading

(Step 1): Add PR label `Expert Review`

cspades left a comment •

edited

Loading

Skylion007 Nov 5, 2025 •

edited

Loading

Skylion007 Nov 8, 2025 •

edited

Loading