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.

• ⚠️ This Roadmap is based on the dev branch; please see the details in its README.
• 🎉 Our technical report for Megatron Core MoE is published in https://arxiv.org/abs/2603.07685
• Previous Roadmaps
  • [ROADMAP][Updated on April 07] Megatron Core MoE Roadmap #1729

Megatron MoE Supported Features

Model Support

• ✅ DeepSeek
  • ✅ DeepSeek-V2
  • ✅ DeepSeek-V3, including MTP
  • ✅ DeepSeek-V3.2
• ✅ Qwen
  • ✅ Qwen2-57B-A14B
  • ✅ Qwen3-235B-A22B
  • ✅ Qwen3.5
• ✅ Kimi-K2
• ✅ Mixtral
• 🚀 DeepSeek-V4 (On-going)

Core MoE Functionality

• ✅ Token dropless MoE - Advanced routing without token dropping
• ✅ Top-K Router with flexible K selection
• ✅ Load balancing losses for expert load balancing optimization

Advanced Parallelism

• ✅ Expert Parallel (EP) with 3D parallelism integration
• ✅ Full parallelism combo: EP + DP + TP + PP + SP support
• ✅ Context Parallel (CP) for long sequence MoE training
• ✅ Parallel Folding - Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training
• ✅ Distributed Optimizer for MoE (ZeRO-1 equivalent)
• ✅ (🚀 New!) Megatron FSDP / HSDP with full expert parallel support

Optimizations

• Memory
  • ✅ Memory Efficient token permutation
  • ✅ Pipeline-aware fine-grained activation offloading [Dev] feat(moe): Fine-grained activation offloading #1912
  • ✅ Fine-grained Recomputations (mla, moe, mlp, moe_act, norm)
• Communication
  • ✅ DeepEP support for H100 and B200
  • ✅ (🚀 New!) HybridEP for GB200
  • ✅ 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, etc.
  • ✅ cuDNN fused Attention and FlashAttn integration
  • ✅ GroupedGEMM and Gradient Accumulation Fusion
  • ✅ Production-ready cudaGraph support for MoE
• Optimizer
  • ✅ Muon and Layer-wise distributed optimizer

Precision Support

• ✅ GroupedGEMM including FP8/MXFP8 support
• ✅ FP8 weights with BF16 optimizer states
• ✅ FP8 training full support
• 🚀 NVFP4 Training

Developer Experience

• ✅ [MoE Model Zoo](https://github.com/yanring/Megatron-MoE-ModelZoo) with pre-training best practices
• ✅ MCore2HF Converter for ecosystem compatibility in megatron-bridge
• ✅ Distributed Checkpointing Support
• ✅ Runtime Upcycling Support for efficient model scaling
• ✅ Layer-wise logging for detailed monitoring

(2026 Q2) Megatron MoE Roadmap

Model Supports

DeepSeek-V4

See #4468 for detailed tracker.

Functionality

• Model Architecture
  • mHC support [dev] feat(mHC): Add basic pytorch implementation of manifold hyper connection(mHC). #2943
  • sqrtsoftplus MoE score function [Dev][MoE] Add a new score function to the router #4193
  • DSA support DeepSeek V3.2 support #2440
  • Hybrid Attention with CSA and HCA [dev] [DeepSeek-v4] Part 1: Hybrid Attention with CSA and HCA #4458
  • Hash MoE, ClampedSwiGLU, mHC contract [dev] [DeepSeek-v4] Part 2: Hash MoE and SwiGLU clamp #4481
  • MTP support with mHC [dev] [DeepSeek-v4] Part 3: MTP support with mHC and new mHC contract #4518
• Muon support
  • Coefficients of Newton-Schulz bump emerging-optimizer package for deepseek v4 coefficients type #4523
  • Distributed Layerwise Optimizer Integrate LayerWiseDistributedOptimizer with DDP buffer infrastructure #4509
  • FP8 primary weight for Muon feat(optim): support --fp8-param-gather for muon + mxfp8 in LayerWise #4987
  • Precision-aware optimizer [dev] Muon fp8 param gather for decoupled layout #5470
• Packed Sequence
  • DSv4HybridAttention THD support [dev] [DeepSeek-v4] Packed Sequence (THD) support for DSv4 Hybrid Attention #5011
  • Dispatcher support for THD [dev] add support for deepep/hybridep dispatcher under thd format training #4816
• Long Context
  • DSv4HybridAttention CP support [dev] [DeepSeek-v4] Context Parallel support #5087
• Convergence
  • HF <-> MCore Checkpoint conversion DeepSeek V4 Bridge NVIDIA-NeMo/Megatron-Bridge#3562
  • E2E convergence verification

Performance Optimization

• mHC
  • mHC fusion kernels [dev] mHC kernel fusion  #3828
  • (perf) mHC kernel benchmark and optimization [dev]: faster implementation of mHC fused kernels #4624
  • mHC support with EP Overlapping [feat] Hybrid model ep overlapping main #4798 [feat] HybridStack grouped syntax and EP-overlap (2/4 of #4798) #4942 [feat] FSDP support for HybridStack EP-overlap (3/4 of #4798) #4943 [fix] Training wire-up for hybrid EP-overlap (4/4 of #4798) #4944
• Hybrid Attention
  • Fused Attention Kernel for Dense MQA
  • Fused Attention Kernel for CSA and HCA [dev] [DeepSeek-v4] Part 4: Fusion Kernels for DSv4 Hybrid Attention #4894
  • FP8 Indexer Computation
• ClampedSwiGLU
  • TE unfused path TE#2938
  • Fused kernel in Cudnn-frontend
• batchedGEMM fusion and FP8 support

E2E Training Recipes #5261

• 4K Dense Training
• 16K Dense Training
• 64K Sparse Training
• 1M Sparse Training

Long-Term Topics

• Anticipatory Routing
• Flexible Activation Checkpointing
• MegaMoE

Qwen3.5

Functionality

• Qwen3-Next / Qwen3.5 functionality supported [dev] [3/5] Qwen3.5 support: SharedExpertMLP meta init #4749 [dev] [4/5] Qwen3.5 support: Interleaved MRoPE layout #4750 [dev] [follow-up] Qwen3.5 support: MoE aux loss padding_mask #4776
• Packed sequence support [Dev] fix(mtp): use padded cu_seqlens in MTP roll for THD with CP #4494
• HF <-> MCore Checkpoint conversion [dev] [2/5] Qwen3.5 support: FSDP DTensor Bridge checkpoint compatibility #4748

Performance Optimization

• Gated Delta Rule (GDN) Optimization
  • GDN kernel fusion on Blackwell (working)
  • GDN FlashQLA integration [dev] moe(perf): Add FlashQLA backend for Hopper/SM90 gated delta net (GDN). #5564
  • Pre-gated delta rule kernel fusion (conv1d + memory ops) [dev] moe(perf): Pre-GDR kernel fusion #5361
• Perf Optimization
  • Full-iteration CUDA graph to reduce CPU overhead of MoE
  • Gated RMSNorm implementation in TE
  • kernel fusion for MRoPE feat(fusions): fused mRoPE for Qwen3.5-VL #5294
• Memory Optimization
  • GDN fine-grained activation offloading (P0.5)
  • GDN selective recompute (in_proj, conv1d, gated_delta_rule) feat(ssm): whole-module 'gdn' selective recompute for GatedDeltaNet #5296

Long Context Training

• GDN context parallel in A2A/Ulysses mode Fuse per-sequence AlltoAll into a unified one in GDN forward #4913
• GDN context parallel in AllGather mode — enables Qwen3.5-397B with >64K context (FLA PR, MCore POC PR) [dev] moe(feat): support chunkwise context parallelism for GDN. #3282
• Dynamic CP support for Multi-modal Training

Multi-Modal Training for Qwen3.5

• OOTB E2E FSDP Example for THD Training Support with Qwen3.5 [dev] [5/5] Qwen3.5 support: Qwen3.5-VL training example #4751
• Dynamic CP Support
• Heterogeneous Parallelism for Encoder-Decoder Add named layouts to HyperCommGrid for heterogeneous parallelism #5148 Apply MIMO SP/CP sharding with explicit groups and enable THD in non-colocated path #5150 Add MIMO hetero topology + distributed bootstrap (examples/mimo training-loop folder) #5260

General Optimizations

Packed Sequence Support

• Add support of flex dispatcher for THD format [dev] add support for deepep/hybridep dispatcher under thd format training #4816
• New varlen dataset compatible with HF datasets varlendataset for thd e2e and benchmark #4832

CUDA Graph

• CUDA Graph Interface Refactor [Dev] Refactor CUDA graph API: decompose cuda_graph_scope into full_iteration impl, inference scope, and per-layer capture modules #4293
• THD + CUDA Graph [dev] partial cuda graph support for thd format training. #4359 Cudagraphs: MCore local path training with variable-sized sequences #5046 [dev] Add chunk-wise (whole-block) CUDA graph support for THD training #5258

Long Context & Context Parallel

• Dynamic CP Support Minor improvements for Dynamic-cp #4226 Add Dynamic Context Parallelism support (port from dev) #5252
• Dynamic CP + CUDA Graph

Communication Optimization

• FSDP + A2A Overlap [Dev][feat] Support A2A Overlap for Megatron-FSDP #3796
• DeepEP v2 Token Dispatcher Integration [Dev] Add DeepEP v2 flex dispatcher backend #4793
• NCCL-EP Support NCCL EP support #5129

Ongoing Long-term Features

• E2E Performance optimization for DeepSeek-V4, Qwen-3.5 and other fine-grained MoEs
• Extreme sparse MoE training exploration
• Migration from GPTModel to HybridModel
• Sync-Free and Full-Iter cudaGraph MoE Training
  • MoE ECHO for dropless MoE load balancing MoE ECHO: Unlocking Sync-Free, Full CUDA-Graph Support for Dropless MoE via Elastic Cloning #2368
• THD and Long Context
  • THD Format E2E Support - End-to-end THD format support [Dev] Add E2E support for THD format #2924
  • Dynamic Context Parallel for Imbalanced Long-Sequence Training
• Megatron FSDP Performance Optimization for MoE Training
• MoE Training with MegaKernels
• Kernel fusions and optimizations for MoE models from TE - MoE training optimization TransformerEngine#2438

(2026 Q1) Highlights

Performance & Kernel Optimizations

• Absorbed MLA - MLA computation optimization for DSA Add absorbed-mla #3193
• Optimization of the standalone permute path Optimization of the standalone permute path deepseek-ai/DeepEP#625
• HybridEP fusion with permutation - Add Permute/Unpermute Fusion Code Path deepseek-ai/DeepEP#588, [Dev] Add permute/unpermute fusion with dispatch/combine in Hybrid-EP #4073
• mHC kernel fusion [dev] mHC kernel fusion  #3828
• GEMM + SwiGLU fused MLP Support GEMM + Swiglu fused MLP #3890

Long Context & Context Parallel

• Hybrid CP Part 2 - Enhanced hybrid data x context parallelism [Dev] feat: Dynamic CP (part 2) #2000
• THD Format E2E Support - End-to-end THD format support [Dev] Add E2E support for THD format #2924

Model & Architecture

• Manifold Hyper Connection (mHC) [dev] feat(mHC): Add basic pytorch implementation of manifold hyper connection(mHC). #2943
• GDN THD Support - Packed sequence support for gated delta net [dev] feat(moe): Support packed sequence for gated delta net (GDN) #2644
• GDN Refinement - Refine gated delta net implementation [dev] perf(moe): Refine gated delta net implementation #3040

Advanced Functionality

• Router replay support for RL training (in progress) feat: add routing replay for Mcore #2693
• Megatron FSDP performance optimization for MoE training

CUDA Graph Enhancements

• Paged Stashing - Dynamic tensor support for dropless MoE with CUDA graph [Dev] Paged Stashing #2690
• CUDA Graph + Offloading - Support CUDA Graph capture with offloading modules [Dev][feat] Support CUDA Graph capture offloading modules #3219
• Optimizer CUDA Graph - Enable CUDA graph for ADAM optimizer Enable CUDA graph for ADAM optimizer #3429

Call for Community Contributions

• Model implementations - Additional MoE model variants
• Performance testing - Performance tests across different platforms and workloads
• Documentation and tutorials - Best practices and optimization guides
• Bug fixes

This roadmap reflects the collective efforts of NVIDIA and our collaborators.

Credits: MCore MoE Team and @sbhavani

Labels: roadmap, moe, call-for-contribution