Skip to content

Fix Megatron-FSDP checkpoint save failure #2138

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shjwudp merged 3 commits into from
Nov 12, 2025
Merged

Fix Megatron-FSDP checkpoint save failure #2138

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6f0606d
Fix Megatron-FSDP checkpoint save failure by ensuring _replace_param_…
shjwudp Nov 5, 2025
2e89e57
code format
shjwudp Nov 5, 2025
98158d9
Merge branch 'main' into fix_fsdp_dtensor
shjwudp Nov 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -898,9 +898,10 @@ def forward_hook(_module, inputs, output):

# Register pre state_dict hook to ensure that the module parameters are
# distributed before saving the state_dict.
self._state_dict_pre_hook = self.module.register_state_dict_pre_hook(
lambda *args, **kwargs: self._replace_param_with_distributed_if_needed()
)
for name, module in self.named_modules():
module.register_state_dict_pre_hook(
lambda *args, **kwargs: self._replace_param_with_distributed_if_needed()
)

@contextmanager
def no_sync(self):
Expand Down
18 changes: 17 additions & 1 deletion tests/unit_tests/test_checkpointing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
import torch
import torch.distributed.checkpoint

from megatron.core.distributed import DistributedDataParallelConfig
from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel
from megatron.core.num_microbatches_calculator import (
init_num_microbatches_calculator,
unset_num_microbatches_calculator,
Expand Down Expand Up @@ -111,6 +113,8 @@ def create_args():
args.retro_add_retriever = False
args.ckpt_convert_update_legacy_dist_opt_format = False
args.ckpt_step = None
args.swiglu = True
args.num_experts = 1

yield args

Expand Down Expand Up @@ -191,7 +195,7 @@ def test_load_base_checkpoint(
assert ckpt_type == expected_ckpt_type


@pytest.mark.parametrize("ckpt_format", ["torch", "torch_dcp"])
@pytest.mark.parametrize("ckpt_format", ["torch", "torch_dcp", "fsdp_dtensor"])
def test_save_checkpoint(init_model_parallel, create_args, tmp_path_dist_ckpt, ckpt_format):
"""Test save_checkpoint."""
args = create_args
Expand All @@ -202,11 +206,21 @@ def test_save_checkpoint(init_model_parallel, create_args, tmp_path_dist_ckpt, c

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

@BoxiangW BoxiangW Nov 5, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this line in use?


iteration = 123
config = TransformerConfig(num_layers=1, kv_channels=1)
model = MockModel(config)
optimizer = MockState({"optimizer": "optimizer_state"})
if ckpt_format == "fsdp_dtensor":
model = FullyShardedDataParallel(
config=config,
ddp_config=DistributedDataParallelConfig(
use_distributed_optimizer=True, use_megatron_fsdp=True
),
module=model,
)
optimizer = MockState({"state": {}})
opt_param_scheduler = MockState({"opt_param_scheduler": "scheduler_state"})
num_floating_point_operations_so_far = 456

Expand All @@ -228,6 +242,8 @@ def test_save_checkpoint(init_model_parallel, create_args, tmp_path_dist_ckpt, c
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":

@BoxiangW BoxiangW Nov 5, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@Skylion007 Skylion007 Nov 5, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@BoxiangW BoxiangW Nov 5, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@Skylion007 Skylion007 Nov 5, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@BoxiangW BoxiangW Nov 5, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@Skylion007 Skylion007 Nov 5, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@BoxiangW BoxiangW Nov 5, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@ko3n1g ko3n1g Nov 5, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!

@Skylion007 Skylion007 Nov 8, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

expected_ckpt_path = ckpt_dir / ".metadata"

assert os.path.exists(expected_ckpt_path)

Expand Down