Fix Megatron-FSDP DTensor checkpoint save/load compatibility

[addsubmuldiv]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

This PR adds a dedicated save/load path for Megatron-FSDP DTensor checkpoints.

The existing Megatron checkpoint path uses MCore sharded state dictionaries. Megatron-FSDP exposes DTensor state through state_dict_for_save_checkpoint(), which must be preprocessed and saved with PyTorch Distributed Checkpoint (DCP).

Main changes

• Add swift/megatron/utils/megatron_fsdp_checkpoint.py to identify FSDP DCP checkpoints, run MCore preprocess_fsdp_dtensor_state_dict(), and save/load DTensor state with PyTorch DCP, following Megatron-Core's configurable partial/strict load behavior.
• Route use_megatron_fsdp=true through the DCP path while preserving the existing non-FSDP checkpoint path.
• Save and restore model, distributed optimizer, optimizer scheduler, RNG, iteration, and checkpoint arguments.
• Follow MCore RNG checkpoint semantics: RNG state is keyed by (PP rank, TP rank), with per-DP-rank state gathered when data_parallel_random_init=true.
• Move warmup_jit_function() before checkpoint loading so JIT warmup cannot consume restored RNG state.
• Fill only missing TP/ETP partition attributes before FSDP wrapping, preserving metadata already assigned by MCore.
• Add narrowly scoped MindSpeed 0.16 compatibility for optional gradient scaling, empty local optimizer shards, FSDP DTensor optimizer loading, and NPU optimizer step tensors.

PR scope and limitations

• The FSDP DTensor checkpoint path is synchronous and supports one flat distributed-optimizer state. Unsupported nested or chained optimizer layouts fail explicitly.
• This PR handles Megatron resumable checkpoints and does not change Hugging Face/safetensors export.
• Exact RNG resume is validated for an unchanged parallel topology; topology-changing deterministic resume is not claimed.
• Megatron-FSDP LoRA remains unsupported with Megatron-Core 0.16 because its FSDP initialization registers post-accumulate-gradient hooks on frozen parameters. The upstream behavior changes in later MCore versions and is not backported here.
• On the tested NPU 0.16 stack, gradient_accumulation_fusion=true can fail when BF16 gradients are accumulated into FP32 main_grad; validation therefore uses --gradient_accumulation_fusion false. Further 0.18 stack adaptation is outside this PR.

Validation

• Ran Qwen3-0.6B full SFT with Megatron-Core 0.16, MindSpeed 0.16, two NPUs, DP=2, dropout=0.1, BF16, distributed Adam, and optim_grads_params.
• The uninterrupted run completed three optimizer steps and saved a complete DCP checkpoint at step 2.
• A new process restored model, optimizer, scheduler, iteration, and RNG from step 2, then completed step 3 and saved another checkpoint.
• Uninterrupted and resumed step-3 metrics matched exactly: loss=6.79597759, grad_norm=446.08346558, and learning_rate=0.0.
• Focused checkpoint/optimizer tests, Ruff, Python compilation, and git diff --check passed.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.