Support torch distributed checkpointing for MFSDP v2
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Goal
Support saving and loading MFSDP v2 checkpoints using PyTorch Distributed Checkpoint (DCP).
## Background
MFSDP v2 introduces a per-module `fully_shard()` path with sharded parameter, gradient, and optimizer state layouts. Checkpointing needs to preserve those layouts well enough to resume training correctly.
The initial target should be the PyTorch Distributed Checkpoint format. Support for Megatron-LM's `torch_dist` checkpoint format should be tracked separately and is likely lower ROI for now.
## Non-goals
- Support for Megatron-LM's `torch_dist` checkpoint format.
- Checkpoint format conversion between DCP and `torch_dist`.
## Related issues and PRs
The items below are for the existing Megatron-FSDP implementation rather than MFSDP v2. They are listed only as references for prior checkpointing issues, fixes, and test coverage.
- #4129
- #4910
- #5265
- #4623
- #4753
- #4843
- #4748
- #3510
- #2805
- #4553
Contributor guide
Assessment
This issue has not been assessed yet.