NVIDIA / NVIDIA/Megatron-LM

Support torch distributed checkpointing for MFSDP v2

Open
#5,534 1 comment 0 reactions 1 assignee Claimed by @ahmadki View on GitHub
MFSDPv2 nemotron PoR
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.