NVIDIA / NVIDIA/Megatron-LM

[ROADMAP] [2026 Q3] Megatron FSDP v2 Roadmap

Open
#6,812 5 comments 0 reactions 0 assignees View on GitHub
call for contribution
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

This roadmap outlines the key features, integrations, performance improvements, and reliability work planned for Megatron FSDP v2 (MFSDP v2). The work focuses on maturing the experimental per-module `fully_shard()` implementation and expanding its support across the Megatron ecosystem.

This is a tentative roadmap and is subject to change. MFSDP v2 remains experimental, and its APIs and configuration may evolve.

For broader Megatron Core priorities, see the [Q3 2026 Megatron Core Roadmap #5676](https://github.com/NVIDIA/Megatron-LM/issues/5676). For information about released functionality, see the [Megatron Core release notes](https://github.com/NVIDIA/Megatron-LM/releases).

---

## MFSDP v2 Roadmap

### Integrations

- **Megatron Core training integration** \- Added an experimental path for training Megatron Core models with MFSDP v2 ([\#5614](https://github.com/NVIDIA/Megatron-LM/issues/5614)).
- **Functional `fully_shard()` foundation** \- Established the initial per-module sharding path and mixed-precision optimizer adapter; this works independently of Megatron Core and supports Transformer Engine, standard PyTorch, and other modules ([\#5659](https://github.com/NVIDIA/Megatron-LM/issues/5659), [\#5411](https://github.com/NVIDIA/Megatron-LM/pull/5411)).
- **Megatron Bridge integration** \- Expose MFSDP v2 through experimental Bridge configuration ([\#6325](https://github.com/NVIDIA/Megatron-LM/issues/6325)).

### Parallelism

- **Sharding modes** \- Added no-shard, ZeRO-1, and ZeRO-2, zero-3 configurations ([\#6137](https://github.com/NVIDIA/Megatron-LM/pull/6137)).
- **HSDP and HFSDP** \- for tradeoffs between memory usage and communication time. ([\#5616](https://github.com/NVIDIA/Megatron-LM/issues/5616)).
- **EP composability** \- Support MFSDP v2 with MCore/TE’s expert parallelism ( with per-expert parameters) ([\#5656](https://github.com/NVIDIA/Megatron-LM/issues/5656),).

### Optimizers and Precision

- **Muon optimizer support** \- Establish a correctness baseline, add efficient redistribution for tensor-atomic momentum updates, benchmark the implementation, and extend it to TP and EP configurations ([\#5533](https://github.com/NVIDIA/Megatron-LM/issues/5533)).
- **MXFP8** \- Support mxfp8 via block-atomic sharding, which preserves MXFP8 data and scale locality across sharding, materialization, gradient reduction, and parameter updates ([\#5615](https://github.com/NVIDIA/Megatron-LM/issues/5615), [\#6587](https://github.com/NVIDIA/Megatron-LM/issues/6587)).

### Checkpointing

- **PyTorch Distributed Checkpoint support** \- Support saving and resuming MFSDP v2 parameter, gradient, and optimizer layouts using PyTorch Distributed Checkpoint ([\#5534](https://github.com/NVIDIA/Megatron-LM/issues/5534)).

### Performance and Memory

- **Communication/computation overlap** \- Continue improving all-gather and reduce-scatter overlap, prefetch behavior, double buffering, and allocator efficiency ([\#5647](https://github.com/NVIDIA/Megatron-LM/issues/5647)).
- **Opt-in single communication stream.** This trades longer communication time for less memory and more determinism. ([\#6471](https://github.com/NVIDIA/Megatron-LM/issues/6471)).
- **Symmetric-memory staging** \- Build on MFSDP v2 symmetric-memory support and its CUDA Graph compatibility ([\#5655](https://github.com/NVIDIA/Megatron-LM/issues/5655), [\#6154](https://github.com/NVIDIA/Megatron-LM/issues/6154)).

### CUDA Graphs

- **Full-iteration CUDA Graphs** \- Build on standalone full-iteration capture and extend coverage through Megatron Core integration ([\#5717](https://github.com/NVIDIA/Megatron-LM/issues/5717), [\#6152](https://github.com/NVIDIA/Megatron-LM/issues/6152)).

---

## How to Provide Feedback

We welcome community input on prioritization. Please:

1. React to items you would like prioritized.
2. Comment on this issue with use cases, constraints, and hardware, model, and parallelism configurations.
3. Open focused feature requests with the `enhancement` label.
4. Contribute pull requests for roadmap items where possible.

---

## Credits

This roadmap reflects the collective efforts of NVIDIA, external contributors, and the Megatron Core community.

Contributor guide

Open the contributing guide

Research direction

This is a planning roadmap rather than an implementation task; no file, test, or entry point is named. Start by reviewing the linked MFSDP v2 issues and the broader roadmap, then choose a focused item with defined scope. Done criteria are not specified here and must be established in that item.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, documentation, machine-learning
Issue type
Documentation
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.