# MegaBlocks: Efficient Sparse Training with Mixture-of-Experts

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

## Key facts

- Cited by 209 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 30 of its 51 references that are in the map.
- First posted on arXiv 29 Nov 2022.

## 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, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [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)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [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)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [The State of Sparsity in Deep Neural Networks](https://mlsanity.com/paper/1902.09574/) (2019)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [Unified Scaling Laws for Routed Language Models](https://mlsanity.com/paper/2202.01169/) (2022)
- [Hash Layers For Large Sparse Models](https://mlsanity.com/paper/2106.04426/) (2021)
- [Efficient Neural Audio Synthesis](https://mlsanity.com/paper/1802.08435/) (2018)
- [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)
- [Sparse GPU Kernels for Deep Learning](https://mlsanity.com/paper/2006.10901/) (2020)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [Fast Sparse ConvNets](https://mlsanity.com/paper/1911.09723/) (2019)
- [Block-Sparse Recurrent Neural Networks](https://mlsanity.com/paper/1711.02782/) (2017)
- [Balanced Sparsity for Efficient DNN Inference on GPU](https://mlsanity.com/paper/1811.00206/) (2018)

## Built on by (most cited first, 30 of 209)

- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [OLMoE: Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2409.02060/) (2024)
- [Challenges and Applications of Large Language Models](https://mlsanity.com/paper/2307.10169/) (2023)
- [A Survey on Efficient Inference for Large Language Models](https://mlsanity.com/paper/2404.14294/) (2024)
- [Arcee's MergeKit: A Toolkit for Merging Large Language Models](https://mlsanity.com/paper/2403.13257/) (2024)
- [OpenMoE: An Early Effort on Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2402.01739/) (2024)
- [Efficient Large Language Models: A Survey](https://mlsanity.com/paper/2312.03863/) (2023)
- [SOLAR 10.7B: Scaling Large Language Models with Simple yet Effective Depth Up-Scaling](https://mlsanity.com/paper/2312.15166/) (2023)
- [Mixture-of-Transformers: A Sparse and Scalable Architecture for Multi-Modal Foundation Models](https://mlsanity.com/paper/2411.04996/) (2024)
- [Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models](https://mlsanity.com/paper/2402.14800/) (2024)
- [Towards Efficient Generative Large Language Model Serving: A Survey from Algorithms to Systems](https://mlsanity.com/paper/2312.15234/) (2023)
- [A Survey of Resource-efficient LLM and Multimodal Foundation Models](https://mlsanity.com/paper/2401.08092/) (2024)
- [Mixtures of Experts Unlock Parameter Scaling for Deep RL](https://mlsanity.com/paper/2402.08609/) (2024)
- [Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA](https://mlsanity.com/paper/2410.20672/) (2024)
- [MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts](https://mlsanity.com/paper/2407.21770/) (2024)
- [JetMoE: Reaching Llama2 Performance with 0.1M Dollars](https://mlsanity.com/paper/2404.07413/) (2024)
- [ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing](https://mlsanity.com/paper/2412.14711/) (2024)
- [Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models](https://mlsanity.com/paper/2501.11873/) (2025)
- [Jamba-1.5: Hybrid Transformer-Mamba Models at Scale](https://mlsanity.com/paper/2408.12570/) (2024)
- [Parameters vs FLOPs: Scaling Laws for Optimal Sparsity for Mixture-of-Experts Language Models](https://mlsanity.com/paper/2501.12370/) (2025)
- [Scaling Laws for Native Multimodal Models](https://mlsanity.com/paper/2504.07951/) (2025)
- [Upcycling Large Language Models into Mixture of Experts](https://mlsanity.com/paper/2410.07524/) (2024)
- [QMoE: Practical Sub-1-Bit Compression of Trillion-Parameter Models](https://mlsanity.com/paper/2310.16795/) (2023)
- [CATS: Contextually-Aware Thresholding for Sparsity in Large Language Models](https://mlsanity.com/paper/2404.08763/) (2024)
- [Mixture Compressor for Mixture-of-Experts LLMs Gains More](https://mlsanity.com/paper/2410.06270/) (2024)
- [Shortcut-connected Expert Parallelism for Accelerating Mixture-of-Experts](https://mlsanity.com/paper/2404.05019/) (2024)
- [Training Sparse Mixture Of Experts Text Embedding Models](https://mlsanity.com/paper/2502.07972/) (2025)
- [How Far Are We From AGI: Are LLMs All We Need?](https://mlsanity.com/paper/2405.10313/) (2024)
- [Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient](https://mlsanity.com/paper/2502.05172/) (2025)
- [HMoE: Heterogeneous Mixture of Experts for Language Modeling](https://mlsanity.com/paper/2408.10681/) (2024)

## Abstract (arXiv)

We present MegaBlocks, a system for efficient Mixture-of-Experts (MoE) training on GPUs. Our system is motivated by the limitations of current frameworks, which restrict the dynamic routing in MoE layers to satisfy the constraints of existing software and hardware. These formulations force a tradeoff between model quality and hardware efficiency, as users must choose between dropping tokens from the computation or wasting computation and memory on padding. To address these limitations, we reformulate MoE computation in terms of block-sparse operations and develop new block-sparse GPU kernels that efficiently handle the dynamism present in MoEs. Our approach never drops tokens and maps efficiently to modern hardware, enabling end-to-end training speedups of up to 40% over MoEs trained with the state-of-the-art Tutel library and 2.4x over DNNs trained with the highly-optimized Megatron-LM framework.