Support torch distributed checkpointing for MFSDP v2 #5534

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.

• Megatron-FSDP checkpoint resumption functional test currently does not support FSDP DTensor checkpointing. #4129
• Megatron-FSDP optimizer checkpoint hangs in save with optimizer_cpu_offload=True #4910
• Cannot save optimizer when swiglu and FSDP #5265
• Fix Megatron-FSDP optimizer CPU offload and checkpointing #4623
• [main] [2/5] Qwen3.5 support: FSDP DTensor Bridge checkpoint compatibility #4753
• Fix GDN DTensor splitting for FSDP checkpointing #4843
• [dev] [2/5] Qwen3.5 support: FSDP DTensor Bridge checkpoint compatibility #4748
• Fix Megatron-FSDP fully_shard() optimizer state DCP checkpointing, and fix DTensor deepcopy bug from PyTorch 26.01. #3510
• How to convert checkpoint from fsdp_dtensor to torch_dist or huggingface in megatron-fsdp mode? #2805
• Megatron-Fsdp checkpoint fails on Qwen3.5 with tensor and expert paralelism #4553