Skip to content

[ROADMAP][2026 Q3] Megatron Core MoE Roadmap #6757

Description

@buptzyb

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

  • 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) with per-weight model-parallel sharding and grouped-expert prefetch

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


(2026 Q3) Megatron MoE Roadmap

Model Supports

DeepSeek-V4

See #4468 for detailed tracker.

Functionality

Performance Optimization

E2E Training Recipes and Validation

Long-Term Topics

  • Anticipatory Routing
  • Flexible Activation Checkpointing
  • MegaMoE

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

  • 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

Functionality and Model Integration

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

Token-Imbalance Scheduling and Load Balancing

Communication Optimization

Megatron FSDP

  • Characterize and improve MoE memory behavior under Megatron FSDP v2 across TP/EP and no-shard/ZeRO modes

GEMM Optimization

  • Reduce GroupedGEMM and grouped-MLP host and launch overhead, including ragged and small-token expert workloads (TE #2897, TE #3099)
  • Complete low-precision grouped-expert kernel and fusion paths (TE #3315, TE #3359, TE #3261)

Optimizer

Mainline

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

(2026 Q2) Highlights

Model & Architecture

Performance & Optimizers

Parallelism, Dispatch, and CUDA Graph


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

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions