You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
The focus for Megatron Core MoE is to provide comprehensive support for latest MoE architectures, advanced parallelism strategies, and performance optimizations for Blackwell. This is a tentative roadmap and subject to change.
⚠️ This Roadmap is based on the dev branch; please see the details in its README.
✅ Optimizer CPU offload, with a configurable fraction of parameter updates executed on CPU and overlapped transfers
✅ Optimizer-state and master-weight offload with GPU-side updates, using chunked staging to reduce peak optimizer memory
Communication
✅ DeepEP support for H100 and B200
✅ HybridEP support for GB200
✅ DeepEP v2 token-dispatcher integration
✅ NCCL EP base dispatcher support
✅ 1F1B EP A2A Overlap - Hiding Expert Parallel Communication with 1F1B Pipeline Schedule
✅ DP/PP/TP/EP Communication Overlapping
Computation
✅ Advanced fusions for Router, Permutation, MLA/RoPE, FP8 casting, and memory-efficient linear cross-entropy on Blackwell
✅ cuDNN fused Attention and FlashAttn integration
✅ GroupedGEMM and Gradient Accumulation Fusion
✅ Production-ready CUDA Graph support for MoE, including partial capture with EP overlap and THD training, with 🚀 chunk-wise and fixed-GA full-iteration THD capture in progress
Optimizer Support
✅ Distributed Optimizer for MoE with optimizer-state sharding across dense and expert data-parallel groups (ZeRO-1 equivalent)
✅ Muon with LayerWise Distributed Optimizer, with non-matrix parameters handled by distributed Adam or Lion
✅ Precision-aware distributed Adam with configurable main-parameter, gradient, and optimizer-state dtypes, including BF16 optimizer states for FP8 training; 🚀 precision-aware Muon for BF16, blockwise FP8, and MXFP8 training, with checkpoint and evaluation validation in progress
✅ Fine-grained precision selection for supported model components via per-module TE recipes; 🚀 BF16/MXFP8 options for specialized DeepSeek-V4 hybrid-attention components, including the DSA indexer and grouped output projection, are in progress
Extend precision-aware optimizer support from distributed Adam to Muon, and validate BF16, blockwise FP8, and MXFP8 training together with checkpoint and evaluation correctness
Validate HybridModel execution and 1F1B overlap composition, including logical-layer scheduling, dense/MoE composition, pipeline boundaries, deferred weight-gradient computation, repeated-MTP sharing, and Distributed Optimizer
Qualify long-context packed-THD training with static and dynamic CP, including correctness, checkpoint resume, memory, and performance
E2E Training Recipes and Delivery
Validate model and optimizer checkpoint continuity, including repeated-MTP sharing, Muon state, and low-precision parameters
Qualify reproducible full-model BF16 and MXFP8 recipes on GB200 and GB300, including correctness, memory margin, performance evidence, and short convergence
Compose KDA/MLA hybrid layers and AttnRes with the MCore MoE path
Performance and CUDA Graph
Validate end-to-end CUDA Graph execution for the Kimi-K3 KDA path in composition with the MCore MoE stack, building on graph-capable backend foundations (cuDNN Frontend #556)
Training Validation and Recipes
Establish a reproducible Kimi-K3 training recipe and complete initial correctness and convergence validation
General Optimizations
Packed Sequence and Dynamic Context Parallel (DCP)
Validate representative variable-length packed/THD workloads end to end with static CP and DCP
Mature DP-balanced scheduling for production-scale packed variable-length workloads
Complete DCP efficiency and load-imbalance metrics
Publish representative workload and configuration guidance
CUDA Graph
Enable and validate end-to-end CUDA Graph execution for KDA-based models
The roadmap's feature-development baseline remains dev. This section summarizes capability-level readiness on main: checked groups are already landed, while unchecked groups are remaining promotion gaps. Inclusion does not constitute a Q3 delivery commitment or imply active ownership. PR-level implementation details are intentionally omitted from the public roadmap.
Landed Baseline
Establish data and sequence-handling foundations on main, including THD/DCP plumbing, sequence packing, variable-length data handling, and CP-aware CUDA Graph compatibility
Establish distributed-execution foundations on main, including schedule planning, expert-parallel communication, Megatron FSDP integration and overlap, checkpointing, and core sharding modes
Establish precision and heterogeneous-model foundations on main, including low-precision parameter movement, heterogeneous model-parallel axes, and grouped-module initialization
Promotion Readiness Gaps
Complete model and HybridModel promotion readiness for DeepSeek-V4, Qwen3.5, Kimi-K3, and shared HybridModel integration
Complete expert-parallel backend and overlap readiness, including dispatcher, checkpoint, recompute, and training composition
Complete THD, sequence-packing, DCP, and CUDA Graph readiness for variable-length and long-context workloads
Complete remaining Megatron FSDP v2 readiness across parallelism modes, pipeline schedules, optimizer checkpointing, and grouped-expert weights
Close optimizer, precision, normalization, initialization, and gradient-handling parity gaps
Ongoing Long-term Features
E2E Performance optimization for DeepSeek-V4, Qwen3.5 and other fine-grained MoEs
Extreme-sparsity MoE execution and scaling, including large expert counts, small-token dispatch, and router/dispatcher co-design
Migration from GPTModel to HybridModel
CUDA Graph readiness across key MoE training workloads
Long-context MoE training scalability and productionization
RL readiness for MoE models, including policy-weight refit, rollout-training consistency, and end-to-end recipe validation
NVFP4 MoE training maturity, including numerical stability, convergence, and production recipe readiness
Broader MegaMoE and MegaKernel research beyond the Q3 integration scope
Production-ready, topology-aware MoE load balancing through runtime scheduling and expert placement
Description
The focus for Megatron Core MoE is to provide comprehensive support for latest MoE architectures, advanced parallelism strategies, and performance optimizations for Blackwell. This is a tentative roadmap and subject to change.
Megatron MoE Supported Features
Model Support
Core MoE Functionality
Advanced Parallelism
Data and Sequence Handling
Optimizations
Optimizer Support
Precision Support
Developer Experience
(2026 Q3) Megatron MoE Roadmap
Model Supports
DeepSeek-V4
See #4468 for detailed tracker.
Functionality
Performance Optimization
DSv4HybridAttention#5722, TE #3160)E2E Training Recipes and Validation
Long-Term Topics
Qwen3.5
Performance Optimization
Long Context Training
Qwen3.5-VL and Multimodal Training
GLM-5.2
See #6392 for detailed tracker.
Functionality and Model/Optimizer Parity
Performance, Memory, and CUDA Graph
Long Context and Parallelism
E2E Training Recipes and Delivery
Kimi-K3
Functionality and Model Integration
Performance and CUDA Graph
Training Validation and Recipes
General Optimizations
Packed Sequence and Dynamic Context Parallel (DCP)
CUDA Graph
Token-Imbalance Scheduling and Load Balancing
Communication Optimization
Megatron FSDP
GEMM Optimization
Optimizer
Mainline
The roadmap's feature-development baseline remains
dev. This section summarizes capability-level readiness onmain: checked groups are already landed, while unchecked groups are remaining promotion gaps. Inclusion does not constitute a Q3 delivery commitment or imply active ownership. PR-level implementation details are intentionally omitted from the public roadmap.Landed Baseline
Promotion Readiness Gaps
Ongoing Long-term Features
(2026 Q2) Highlights
Model & Architecture
Performance & Optimizers
Parallelism, Dispatch, and CUDA Graph
Call for Community Contributions
This roadmap reflects the collective efforts of NVIDIA and our collaborators.
Credits: MCore MoE Team and @sbhavani
Labels:
roadmap,moe,call-for-contribution