# MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core

arXiv: https://arxiv.org/abs/2504.14960 · HTML: https://mlsanity.com/paper/2504.14960/

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Placed in Efficient MoE training and inference systems, one of 12 directions in Mixture-of-experts.
- Builds on 24 of its 33 references that are in the map.
- First posted on arXiv 21 Apr 2025.

## Placed in

- Mixture-of-experts / [Efficient MoE training and inference systems](https://mlsanity.com/moe/efficient-moe-training-and-inference-systems/)

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models](https://mlsanity.com/paper/2401.06066/) (2024)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [MegaBlocks: Efficient Sparse Training with Mixture-of-Experts](https://mlsanity.com/paper/2211.15841/) (2022)
- [Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints](https://mlsanity.com/paper/2212.05055/) (2022)
- [Tutel: Adaptive Mixture-of-Experts at Scale](https://mlsanity.com/paper/2206.03382/) (2022)
- [Efficient Large Scale Language Modeling with Mixtures of Experts](https://mlsanity.com/paper/2112.10684/) (2021)
- [Scaling Laws for Fine-Grained Mixture of Experts](https://mlsanity.com/paper/2402.07871/) (2024)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [DAPPLE: A Pipelined Data Parallel Approach for Training Large Models](https://mlsanity.com/paper/2007.01045/) (2020)
- [Upcycling Large Language Models into Mixture of Experts](https://mlsanity.com/paper/2410.07524/) (2024)
- [A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training](https://mlsanity.com/paper/2303.06318/) (2023)
- [MoESys: A Distributed and Efficient Mixture-of-Experts Training and Inference System for Internet Services](https://mlsanity.com/paper/2205.10034/) (2022)

## Built on by (most cited first, 17 of 17)

- [Scalable Training of Mixture-of-Experts Models with Megatron Core](https://mlsanity.com/paper/2603.07685/) (2026)
- [UltraEP: Unleash MoE Training and Inference on Rack-Scale Nodes with Near-Optimal Load Balancing](https://mlsanity.com/paper/2606.04101/) (2026)
- [LAER-MoE: Load-Adaptive Expert Re-layout for Efficient Mixture-of-Experts Training](https://mlsanity.com/paper/2602.11686/) (2026)
- [ReLibra: Routing-Replay-Guided Load Balancing for MoE Training in Reinforcement Learning](https://mlsanity.com/paper/2605.08639/) (2026)
- [LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget](https://mlsanity.com/paper/2607.14952/) (2026)
- [WeChat-YATT: A Scalable, Simple, Efficient, and Production Ready Training Library](https://mlsanity.com/paper/2508.07970/) (2025)
- [BOOST: BOttleneck-Optimized Scalable Training Framework for Low-Rank Large Language Models](https://mlsanity.com/paper/2512.12131/) (2025)
- [G-Core: A Simple, Scalable and Balanced RLHF Trainer](https://mlsanity.com/paper/2507.22789/) (2025)
- [DualSparse-MoE: Coordinating Tensor/Neuron-Level Sparsity with Expert Partition and Reconstruction](https://mlsanity.com/paper/2508.18376/) (2025)
- [MixServe: An Automatic Distributed Serving System for MoE Models with Hybrid Parallelism Based on Fused Communication Algorithm](https://mlsanity.com/paper/2601.08800/) (2026)
- [MoEless: Efficient MoE LLM Serving via Serverless Computing](https://mlsanity.com/paper/2603.06350/) (2026)
- [MinT: Managed Infrastructure for Training and Serving Millions of LLMs](https://mlsanity.com/paper/2605.13779/) (2026)
- [Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference](https://mlsanity.com/paper/2604.26294/) (2026)
- [MoEMB: Scaling Universal Multimodal Embeddings with Efficient Mixture-of-Experts Models](https://mlsanity.com/paper/2609.08663/) (2026)
- [Memory-Efficient Expert Routing for Distributed MoE Training](https://mlsanity.com/paper/2610.07333/) (2026)
- [Expert Coupling in MoE Pretraining: Reducing All-to-All Overhead with Correlated Placement and Token Shuffling](https://mlsanity.com/paper/2610.09372/) (2026)
- [ATTUNE: Agentic Throughput Tuning for LLM Training](https://mlsanity.com/paper/2610.17528/) (2026)

## Abstract (arXiv)

Mixture of Experts (MoE) models enhance neural network scalability by dynamically selecting relevant experts per input token, enabling larger model sizes while maintaining manageable computation costs. However, efficient training of large-scale MoE models across thousands of GPUs presents significant challenges due to limitations in existing parallelism strategies. We introduce an end-to-end training framework for large-scale MoE models that utilizes five-dimensional hybrid parallelism: Tensor Parallelism, Expert Parallelism, Context Parallelism, Data Parallelism, and Pipeline Parallelism. Central to our approach is MoE Parallel Folding, a novel strategy that decouples the parallelization of attention and MoE layers in Transformer models, allowing each layer type to adopt optimal parallel configurations. Additionally, we develop a flexible token-level dispatcher that supports both token-dropping and token-dropless MoE training across all five dimensions of parallelism. This dispatcher accommodates dynamic tensor shapes and coordinates different parallelism schemes for Attention and MoE layers, facilitating complex parallelism implementations. Our experiments demonstrate significant improvements in training efficiency and scalability. We achieve up to 49.3% Model Flops Utilization (MFU) for the Mixtral 8x22B model and 39.0% MFU for the Qwen2-57B-A14B model on H100 GPUs, outperforming existing methods. The framework scales efficiently up to 1,024 GPUs and maintains high performance with sequence lengths up to 128K tokens, validating its effectiveness for large-scale MoE model training. The code is available in Megatron-Core.