Baike.dev
Log in
Back to tool/Back to issues
#6757·Megatron-LM

[ROADMAP][2026 Q3] Megatron Core MoE Roadmap

Author: buptzybCreated Aug 21, 2026Updated Sep 18, 2026
Labelscall for contribution

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

  • 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)
  • Correctness and numerical-stability fixes for hybrid attention and DSA (#5960, #6172, #6349, #6372)

Performance Optimization

  • mHC
    • 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)
    • Complete BF16/MXFP8 batched GEMM (#5722, TE #3160)
    • 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)
    • Optimize MTP contiguous packed-CP rolls (#6246)
    • Reduce indexer and hybrid-attention memory, including inverse-RoPE fusion (#6185)

E2E Training Recipes and Validation

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

  • 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

  • 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

  • Skip invalid slots in sparse-indexer backward (#6166)
  • Correct DSA FLOPs accounting and resolve the CP1 indexer-scoring path (#6753, #6206)
  • Fuse standard RoPE packing for MLA and DSA (#6343)
  • Make DSA sparse-attention backward fixed-shape for empty-row cases (#6910)
  • Enable low-precision AbsorbedMLA up-projection recompute (#6178)
  • Close full-model memory safety with selective recompute and activation offload
  • 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

  • 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, Transformer Engine #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, cuDNN Frontend #645, cuDNN Frontend #670)
  • 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

Performance and CUDA Graph

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

  • Full activation recompute with EP A2A overlap (#5869)
  • 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)
  • 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

  • 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

  • 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

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

Ma

Source: NVIDIA/Megatron-LM

View original on GitHubView discussion on GitHub