[ROADMAP] [2026 Q3] Megatron FSDP v2 Roadmap
- 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
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