MegaBlocks: Efficient Sparse Training with Mixture-of-Experts

By Trevor Gale, D. Narayanan, Cliff Young, Matei A. Zaharia. 2022

234 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Large Language Models, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Large Language Models

In Efficiency & Compression

Citations in the corpus

Mixtral of Experts (2024, background): “For example, Megablocks [13] casts the feed-forward network (FFN) operations of the MoE layer as large sparse matrix multiplications, significantly enhancing the execution speed and naturally handling cases where different experts get a variable number of tokens assigned to them.”
A Survey on Mixture of Experts in Large Language Models (2024, unknown): “However, OLMoE [142] ﬁnds that expert-choice gating exhibits lower accuracy compared to token-choice dropless MoE [112].”
QMoE: Practical Sub-1-Bit Compression of Trillion-Parameter Models (2023, background): “Since MoEs can feature rather different computational profiles from standard dense models, there is also significant research on optimizing inference and training systems (Barham et al., 2022; Gale et al., 2023; Hwang et al., 2023).”
Training Sparse Mixture Of Experts Text Embedding Models (2025, unknown): “Recent advances in MoE training, such as upcycling from pretrained transformers (Komatsuzaki et al., 2023) and efficient block-sparse implementations (Gale et al., 2022), have made MoE training even more efficient.”
Mixture-of-Experts in Remote Sensing: A Survey (2026, unknown): “Works such as FastMoE [45], DeepSpeed-MoE [118], Tutel [55], HetuMoE [108], FasterMoE [46] and MegaBlocks [32] focus on efficient distributed training and inference, including all-to-all communication scheduling, CUDA kernel optimizations, expert placement, and parallelism strategies.”