[ROADMAP][2026 Q3] Megatron Core MoE Roadmap

## 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](https://github.com/NVIDIA/Megatron-LM/tree/dev); 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
- https://github.com/NVIDIA/Megatron-LM/issues/4815
- https://github.com/NVIDIA/Megatron-LM/issues/1729

---

# Megatron MoE Supported Features

### Model Support

- **DeepSeek**
- ✅ DeepSeek-V2
- ✅ DeepSeek-V3, including MTP
- ✅ DeepSeek-V3.2
- ✅ DeepSeek-V4
- **Qwen**
- ✅ Qwen2-57B-A14B
- ✅ Qwen3-235B-A22B
- ✅ Qwen3.5
- 🚀 Qwen3.5-VL
- **Kimi**
- ✅ Kimi-K2
- 🚀 Kimi-K3
- **GLM**
- 🚀 GLM-5.2

### 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
- ✅ **Shared-expert support**

### Advanced Parallelism

- Basic Parallelism
- ✅ **Expert Parallel (EP)** with 3D parallelism integration
- ✅ **Context Parallel (CP) and Dynamic Context Parallel (DCP)** for long-sequence training and imbalanced sequence workloads
- ✅ **Pipeline Parallelism (PP)** with **VPP and flexible stage configuration** for workload balancing
- Parallelism Composition and Mapping
- ✅ **Full parallelism combo**: EP + DP + TP + PP + CP + SP + VPP support
- ✅ **Parallel Folding** - Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training
- Data- and Model-Parallel State Sharding
- ✅ **Megatron FSDP / HSDP** with **expert-parallel support**
- 🚀 **Megatron FSDP v2 expanded MoE and ZeRO-mode coverage**
- 🚀 **Generalized Tensor Parallelism (GTP) dev integration**, bringing per-weight model-parallel sharding and grouped-expert prefetch from `main` into the `dev` MoE stack (#6882)

### Data and Sequence Handling

- ✅ **Variable-length dataset ingestion** for packed-sequence training
- ✅ **DP×CP-aware sequence packing and batch scheduling**, including DP-balanced and DCP-aware scheduling
- ✅ **End-to-end packed-sequence / THD execution and metadata propagation across TP/PP/CP**, including MTP

### Optimizations

- Memory
- ✅ **Memory Efficient token permutation**
- ✅ **Pipeline-aware fine-grained activation offloading**
- ✅ **Fine-grained Recomputations** (mla, moe, mlp, moe_act, norm)
- ✅ **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

### Precision Support

- ✅ **GroupedGEMM** including FP8/MXFP8 support
- ✅ **FP8 training** full support
- ✅ [**FP8 primary weights and parameter all-gather**](https://github.com/NVIDIA/Megatron-LM/blob/dev/megatron/core/transformer/moe/README.md#fp8-training) across Distributed Optimizer and Megatron-FSDP
- ✅ **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
- 🚀 **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
- ✅ **Single-GPU memory and kernel profiling for distributed MoE layouts** with Fake Process Group
- ✅ [**DeepSeek-V3 GB200 performance reproduction guide**](https://github.com/NVIDIA/Megatron-LM/blob/dev/docs/discussions/deepseek-v3-gb200-optimization/deepseek-v3-gb200-reproduce-guide.md)
- ✅ [**Megatron Core MoE and advanced-feature guides**](https://github.com/NVIDIA/Megatron-LM/tree/dev/docs/user-guide/features), covering MoE architecture, parallelism, CUDA Graph, memory/offload, optimizer/FSDP, pipeline layouts, and data loading
- ✅ [**Public MoE training recipes**](https://github.com/NVIDIA/Megatron-LM/tree/dev/examples/moe_recipes)

---

# (2026 Q3) Megatron MoE Roadmap

## Model Supports

### DeepSeek-V4
See #4468 for detailed tracker.

#### Functionality

- Convergence
- [ ] **E2E convergence verification**
- [ ] **Deterministic and checkpoint-restart validation** (#5785)
- Long Context
- [ ] **64K-to-1M context training curriculum validation**
- Muon Support
- [ ] **Extend precision-aware optimizer support from distributed Adam to Muon**, and validate BF16, blockwise FP8, and MXFP8 training together with checkpoint and evaluation correctness
- [ ] **Complete decoupled-optimizer checkpoint and LayerWise parameter-gather fixes** (#6015)
- [x] **Correctness and numerical-stability fixes for hybrid attention and DSA** (#5960, #6172, #6349, #6372)

#### Performance Optimization

- mHC
- [x] **mHC selective recompute with CUDA Graph under EP overlap** (#5841)
- Hybrid Attention
- [ ] **Fused Attention Kernel for Dense MQA**
- [ ] **Complete BF16/MXFP8 compact DSA indexer integration and validation** (#5992, [cuDNN Frontend #370](https://github.com/NVIDIA/cudnn-frontend/pull/370))
- [ ] **Complete BF16/MXFP8 batched GEMM** (#5722, [TE #3160](https://github.com/NVIDIA/TransformerEngine/pull/3160))
- [x] **Dispatch THD CSA compressor pooling to fused cuDNN Frontend kernels** (#5984)
- ClampedSwiGLU
- [ ] **Additional fused kernel in cuDNN Frontend**
- Long Context and Memory
- [ ] **Complete load-balanced CP indexer support**, building on delivered CP communication overlap (#6058, #5691)
- [x] **Optimize MTP contiguous packed-CP rolls** (#6246)
- [ ] **Reduce indexer and hybrid-attention memory**, including inverse-RoPE fusion (#6185)
#### E2E Training Recipes and Validation

- [x] **Refresh public 4K DeepSeek-V4 Flash and Pro training recipes** (#6335)
- [ ] **Complete 4K dense training convergence validation** (recipe delivered in #5266)
- [ ] **Complete 16K dense training convergence validation**
- [ ] **Complete 64K sparse training convergence validation**
- [ ] **Complete 1M sparse training convergence validation**

#### Long-Term Topics

- [ ] Anticipatory Routing
- [ ] Flexible Activation Checkpointing

### Qwen3.5

#### Performance Optimization

- Gated Delta Rule (GDN) Optimization
- [x] **Refactor shared GDN/GDN2 implementation foundations** (#6088)
- [ ] **GDN kernel fusion on Blackwell**
- [ ] **Gated RMSNorm implementation in TE**
- Memory Optimization
- [ ] **GDN fine-grained activation offloading**
- [ ] **Enable and validate partial/full-iteration CUDA Graph paths for the full Qwen3.5 model**, including THD and DCP combinations

#### Long Context Training

- [x] **Add mixed-attention CP layout handling for SBHD and THD workloads** (#6387)
- [ ] **Validate THD, CP, DCP, and CUDA Graph combinations for Qwen3.5**

#### Qwen3.5-VL and Multimodal Training

- [ ] **Complete the Qwen3.5-VL HybridModel migration** (#6315)
- [ ] **Complete Qwen3.5-VL pipeline-parallel support and validation** (#6081)
- [ ] **Complete full-model Qwen3.5-VL MDP validation**, closing the remaining encoder/decoder CP and full-iteration CUDA Graph gaps and qualifying composition with DP/EP overlap, THD/DCP, and recompute
- [ ] **Complete and standardize the variable-length multimodal data path**, including the Energon dataloader and vision-sequence packing (#5927)
- [ ] **Evaluate packing and global-batch scheduling strategies to reduce padding and data imbalance for dynamic-resolution multimodal workloads**
- [ ] **Prepare and validate the Qwen3.5-VL recipe with Megatron FSDP v2, gated on readiness of the required multimodal FSDP v2 path** (#6139)
- [ ] **Complete multimodal CI and model-level validation** (#6083, #6173)

### GLM-5.2

See #6392 for detailed tracker.

#### Functionality and Model/Optimizer Parity

- [x] **Land the ordinary DSA baseline and backend selection** (#6020)
- [ ] **Add latent-attention normalization and indexer-projection precision controls** (#6204, #6279)
- [ ] **Establish per-head Muon and packed-MLA optimizer parity**
- [ ] **Complete repeated-MTP IndexShare/KVShare sharing, end-to-end TV training, and future-row alignment** (#6472, #6473, #6741)

#### Performance, Memory, and CUDA Graph

- [x] **Skip invalid slots in sparse-indexer backward** (#6166)
- [ ] **Correct DSA FLOPs accounting and resolve the CP1 indexer-scoring path** (#6753, #6206)
- [x] **Fuse standard RoPE packing for MLA and DSA** (#6343)
- [ ] **Make DSA sparse-attention backward fixed-shape for empty-row cases** (#6910)
- [x] **Enable low-precision AbsorbedMLA up-projection recompute** (#6178)
- [ ] **Close full-model memory safety with selective recompute and activation offload**
- [x] **Reduce DP-balanced THD rerouting memory** (#6378)
- [ ] **Validate DSA, repeated-MTP, and whole-MoE paged-stash CUDA Graph composition** (#6022)

#### Long Context and Parallelism

- [ ] **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

### Kimi-K3

See https://github.com/NVIDIA/Megatron-LM/issues/6872 for detailed tracker.

#### Functionality and Model Integration

- [x] **Establish the baseline Quantile Balancing path for MoE routing** (#5349)
- [ ] **Complete Kimi-K3 global-batch Quantile Balancing in Megatron Core and its supporting dense-routing-map integration, including fused CUDA Graph readiness** (#6637, #6614, [Transformer Engine #3395](https://github.com/NVIDIA/TransformerEngine/pull/3395), [Transformer Engine #3426](https://github.com/NVIDIA/TransformerEngine/pull/3426))
- [ ] **Add optional output normalization for latent MoE** (#6448, #6449)
- [ ] **Integrate latent-MoE RMSNorm into Megatron Core training** (https://github.com/NVIDIA/Megatron-LM/pull/6804)
- [ ] **Complete end-to-end Kimi-K3 SiTU-GLU integration across dense, routed-expert, and shared-expert FFNs**, building on merged Transformer Engine and cuDNN Frontend foundations (#6673, [Transformer Engine #3402](https://github.com/NVIDIA/TransformerEngine/pull/3402), [cuDNN Frontend #645](https://github.com/NVIDIA/cudnn-frontend/pull/645), [cuDNN Frontend #670](https://github.com/NVIDIA/cudnn-frontend/pull/670))
- [x] **Add baseline KDA and Gated MLA hybrid-model support in MCore** (#6556)
- [ ] **Complete KDA low-rank projections and extended MLA integration and validation** (#6877)
- [ ] **Attention Residual support in MCore** [#6840](https://github.com/NVIDIA/Megatron-LM/pull/6840)

#### Performance and CUDA Graph

- [ ] **Migrate GDN kernel fusion works to KDA**
- [ ] **MLA latent CP** (https://github.com/NVIDIA/Megatron-LM/pull/6829)
- [ ] **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](https://github.com/NVIDIA/cudnn-frontend/pull/556))

#### Training Validation and Recipes

- [ ] **Establish a reproducible Kimi-K3 training recipe and complete initial correctness and convergence validation**

#### P1 items

- [ ] **Per-head Muon**
- [ ] **Multimodal Support (MoonViT-V2)**
- [ ] **MXFP4 Support**

## 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**
- [ ] **Reduce DCP communication-memory overhead and support arbitrary CP sizes**, including non-power-of-two configurations (#6801, [TE #3420](https://github.com/NVIDIA/TransformerEngine/pull/3420))
- [ ] **Develop workload-model-driven DCP configuration optimization and publish representative guidance** for variable-length THD workloads across supported attention backends

### CUDA Graph

- [ ] **Enable and validate end-to-end CUDA Graph execution for KDA-based models**
- [ ] **Enable and validate end-to-end CUDA Graph execution for DSA-based models** (#6392)
- [ ] **Padding-aware GEMM for CUDA Graph-padded variable-length workloads**, avoiding tail-padding compute while preserving replay-stable buffer shapes
- [ ] **Chunk-wise whole-block CUDA Graph support for THD** (#5258)
- [ ] **DCP partial CUDA Graph capture and shared-pool slot reuse, including characterization and reduction of model-dependent graph-memory overhead** (#5618, #6479, [TE #3353](https://github.com/NVIDIA/TransformerEngine/pull/3353))
- [ ] **Megatron FSDP v2 integration with partial CUDA Graph capture**, including end-to-end MoE composition validation
- [ ] **Paged stash for whole-MoE CUDA Graph capture** (#6022)
- [ ] **Full-iteration THD CUDA Graph under fixed gradient accumulation** (#5807)

### MegaMoE

- [ ] **Advance MegaMoE integration and correctness/determinism qualification across representative MoE workloads**, including MOK and CuTe DSL paths (#6572, [TE #3350](https://github.com/NVIDIA/TransformerEngine/pull/3350))
- [ ] **Benchmark MegaMoE variants against one another and against established EP-overlap baselines across representative workloads**, covering performance, memory, scalability, and composition constraints

### Token-Imbalance Scheduling and Load Balancing

- [ ] **Build a common benchmark for runtime MoE token imbalance**, comparing the current baseline with ECHO (#2368), UltraEP, MoonEP (#6134), and related approaches
- [ ] **Design and evaluate a common runtime-aware MoEScheduler framework**, adapt ECHO to the common interface, and establish unified correctness, performance, and integration validation across UltraEP, MoonEP, and ECHO backends

### Communication Optimization

- [x] **Full activation recompute with EP A2A overlap** (#5869)
- [x] **Per-layer activation recompute with EP A2A overlap** (#6311)
- [ ] **HybridModel EP-overlap integration, checkpoint compatibility, and performance/memory parity with the GPTModel path for the DeepSeek-V3 proxy workload**
- [ ] **HybridEP robustness for zero-copy, small-token, ultra-sparse, and paged-stash workloads**

### Megatron FSDP

- [ ] **Characterize and improve Megatron FSDP v2 MoE memory and execution behavior** across TP/EP, PP/VPP and 1F1B overlap, no-shard/ZeRO modes, heterogeneous dense/expert sharding, memory-safe initialization from scratch, and LayerWise Muon bucket overlap

### GEMM Optimization

- [ ] **Validate device-initiated GroupedGEMM with paged-stash composition across representative MoE workloads**, including correctness, performance, and in-flight buffer memory behavior
- [ ] **Enable and validate 2D MXFP8 quantization for Megatron Core GEMM paths on Blackwell**, building on merged Transformer Engine support ([TE #2634](https://github.com/NVIDIA/TransformerEngine/pull/2634))
- [ ] **Reduce GroupedGEMM and grouped-MLP host and launch overhead**, including ragged and small-token expert workloads ([TE #2897](https://github.com/NVIDIA/TransformerEngine/issues/2897), [TE #3099](https://github.com/NVIDIA/TransformerEngine/pull/3099))
- [ ] **Complete low-precision grouped-expert kernel and fusion paths** ([TE #3315](https://github.com/NVIDIA/TransformerEngine/pull/3315), [TE #3359](https://github.com/NVIDIA/TransformerEngine/pull/3359), [TE #3261](https://github.com/NVIDIA/TransformerEngine/pull/3261))

### Optimizer

- [x] **Add optimizer-state and master-weight offload with GPU-side updates**, using chunked staging to reduce peak optimizer memory for Distributed Adam and compact-layout BF16 LayerWise Muon (#6244)
- [ ] **Complete parameter-structure-aware Muon updates** for fused attention projections and gated-MLP weights (#6326, #6688)
- [ ] **Complete explicit optimizer assignment** for MoE routers and other parameters excluded from Muon, building on corrected LayerWise buffer routing (#6504, #6642)

## Release Qualification and Branch Alignment

This section separates `dev` release qualification from bidirectional branch alignment. Branch-alignment items are capability-level: inclusion does not constitute a Q3 delivery commitment or imply active ownership, and PR-level implementation details are intentionally omitted.

### Dev Release Qualification

- [ ] **Publish a qualified dev-branch release tag**, gated by convergence and performance validation of one canonical configuration each for DeepSeek-V3 4K, DeepSeek-V4 Flash 64K, and Qwen3.5-397B

### Dev-to-Main Promotion

**Already on main**

- [x] **Establish data and sequence-handling foundations on main**, including THD/DCP plumbing, sequence packing, variable-length data handling, and CP-aware CUDA Graph compatibility
- [x] **Establish distributed-execution foundations on main**, including schedule planning, expert-parallel communication, Megatron FSDP integration and overlap, checkpointing, and core sharding modes
- [x] **Establish precision and heterogeneous-model foundations on main**, including low-precision parameter movement, heterogeneous model-parallel axes, and grouped-module initialization

**Remaining promotion 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**

### Main-to-Dev Synchronization

- [ ] **Integrate the main-branch GTP stack into dev and validate it with the dev MoE stack**

## Ongoing Long-term Features

- **E2E Performance optimization** for DeepSeek-V4, Qwen3.5 and other fine-grained MoEs
- **Future-generation model architecture exploration and training-system readiness**; for example:
- **Qwen3.8-Flash-Next architectural components:** QSA, gated residuals, and N-gram modeling
- **GLM-5.3-Flash architectural components:** IndexPool
- **Sparse/linear-attention hybrids:** GDN + QSA in Qwen3.8-Flash-Next and KDA + DSA in GLM-5.3-Flash
- **Higher-sparsity MoE architectures and scalable execution:** larger expert counts, small-token dispatch, and router/dispatcher co-design
- **Looped Transformers and other iterative-depth architectures**
- **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
- **Production-ready, topology-aware MoE load balancing through runtime scheduling and expert placement**

# (2026 Q2) Highlights

### Model & Architecture

- [x] **DeepSeek-V4 architecture baseline**, including hybrid attention, Hash MoE, ClampedSwiGLU, mHC, MTP, and HybridModel (#4468, #5762)
- [x] **DeepSeek-V4 THD and context-parallel support** (#4816, #5011, #5087)
- [x] **DeepSeek-V4-Flash 4K training recipe** (#5266)
- [x] **DeepSeek-V4 HF <-> MCore checkpoint conversion** ([Megatron Bridge #3562](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/3562))
- [x] **Qwen3.5 architecture, packed sequence, and checkpoint conversion** (#4749, #4494, #4748)
- [x] **Qwen3.5 GDN context parallel, fused mRoPE, OOTB E2E THD/FSDP example, and heterogeneous encoder-decoder parallelism foundations** (#4913, #3282, #5962, #4751, #5148, #5150, #5260)

### Performance & Optimizers

- [x] **mHC, hybrid-attention, and ClampedSwiGLU fusion foundations** (#3828, #4624, #4894, #5130, [TE #2938](https://github.com/NVIDIA/TransformerEngine/pull/2938))
- [x] **Correct mHC boundaries in EP-overlap schedules** (#5471)
- [x] **LayerWise optimizer, Newton–Schulz integration, and FP8-aware Muon parameter gather** (#4509, #5320, #5388, #5470)
- [x] **Qwen3.5 pre-GDR, chunk-wise CP pre-GDR, and selective recompute** (#5361, #5638, #5296)

### Parallelism, Dispatch, and CUDA Graph

- [x] **Variable-length THD, real/padded-length propagation, and DCP foundations** (#4832, #5541, #5544, #4226)
- [x] **DeepEP v2 and NCCL EP dispatcher integration** (#4793, #5129)
- [x] **FSDP communication with EP A2A overlap** (#3796)
- [x] **CUDA Graph interface refactor and partial THD capture** (#4293, #4359)

---

## 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`