MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (2211.15841v1)

Abstract: 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.

• The paper introduces MegaBlocks, which uses a block-sparse matrix approach to significantly enhance the efficiency of MoE training on GPUs.
• It eliminates the need for token dropping and padding by dynamically routing tokens across expert networks, optimizing hardware utilization.
• Benchmark tests reveal up to 40% faster end-to-end MoE training and 1.8–2.4× speed-ups compared to state-of-the-art frameworks.

Introduction to Efficient MoE Training

The presented research introduces a system known as MegaBlocks, which is developed to improve the training efficiency of Mixture-of-Experts (MoE) models on Graphics Processing Units (GPUs). MoE models offer a structured way to capture sparsity in deep neural networks (DNNs) for reduced computation, which leads to faster training times and can potentially enhance model quality. However, conventional MoE training faces a challenge in dealing with the dynamic routing of data through various computational paths, which either leads to wasteful computation or the need to drop data tokens during training. MegaBlocks addresses these drawbacks by employing a block-sparse matrix approach that ensures no data tokens are dropped and makes better use of the hardware, ultimately increasing training speeds significantly.

Overcoming Limitations of Current Frameworks

Current methodologies for MoE training enforce constraints to manage dynamic routing within layers. These constraints lead to a tradeoff between model performance and hardware efficiency, wherein users must decide to either drop tokens or fill up computational paths with padding. The MegaBlocks system bypasses this issue by redefining MoE computation to utilize block-sparse matrices, allowing it to work effectively without dropping any tokens and mapping efficiently onto the architecture of modern GPUs. This approach provides significant improvements in training speed compared to other state-of-the-art libraries like Tutel and even outperforms highly-optimized frameworks for standard DNNs such as Megatron-LM.

Block-Sparse Matrix Approach

MegaBlocks uses high-performance GPU kernels specially designed for block-sparse matrix products capable of adapting to the imbalanced distribution of tokens to experts, a typical scenario in MoE computation. Essentially, computation is performed in parallel for sets of tokens assigned across various "expert" neural networks without the need for padding or dropping tokens. By utilizing novel encoding strategies and creative management of transposed matrix data, MegaBlocks enhances the speed at which these operations are executed.

Practical Implications and Benchmarks

In benchmark tests, MegaBlocks has shown to accelerate end-to-end training of MoE models by up to 40% compared to the Tutel library. Moreover, when compared against DNNs optimized with Megatron-LM, MegaBlocks delivers training speed-ups ranging from 1.8 to 2.4 times faster. These benchmarks are a significant indication of MegaBlocks' potential to make the training of vast, sparse models like MoEs more practical and cost-effective.

In conclusion, the MegaBlocks system offers a powerful advancement in training MoE models on GPUs, providing accelerated training speeds and eliminating the trade-offs between model quality and hardware efficiency inherent in prior frameworks. This block-sparse matrix approach not only retains the integrity of the data throughout the training process but also maximizes the computational power of modern hardware accelerators.