Scattered Mixture-of-Experts Implementation (2403.08245v2)

Abstract: We present ScatterMoE, an implementation of Sparse Mixture-of-Experts (SMoE) on GPUs. ScatterMoE builds upon existing implementations, and overcoming some of the limitations to improve inference and training speed, and memory footprint. This implementation achieves this by avoiding padding and making excessive copies of the input. We introduce ParallelLinear, the main component we use to build our implementation and the various kernels used to speed up the operation. We benchmark our implementation against Megablocks, and show that it enables a higher throughput and lower memory footprint. We also show how ParallelLinear enables extension of the Mixture-of-Experts concept by demonstrating with an implementation of Mixture of Attention.

• The paper introduces ScatterMoE with its ParallelLinear component that fuses grouped matrix operations to significantly reduce memory usage and boost efficiency.
• Benchmarking reveals up to 66.2% memory savings during training, 53.6% during inference, and a 38.1% throughput increase over previous implementations.
• The approach extends to Mixture-of-Attention, demonstrating a scalable, adaptable framework for optimizing SMoE models and related neural network architectures on GPUs.

Scattered Mixture-of-Experts Implementation Enhances GPU Efficiency

Introduction

The implementation of Sparse Mixture-of-Experts (SMoE) on GPUs has been a challenging endeavor due to the inherent complexity in efficiently leveraging GPU parallelism. The paper presents ScatterMoE, an innovative approach to SMoE implementation that significantly improves inference and training speed, along with reducing the memory footprint on GPUs. By addressing the limitations found in existing implementations such as padding and excessive copying of input data, ScatterMoE introduces a more memory and computationally efficient method for executing SMoE models.

Key Contributions

ParallelLinear Component

A cornerstone of ScatterMoE is the introduction of the ParallelLinear component. This linear module is capable of executing grouped matrix operations on scattered groups of data without necessitating padding or making unnecessary copies of the input tensor. ParallelLinear is fundamentally designed to minimize memory usage while maximizing computational efficiency by fusing the operations of grouping, matrix multiplication, and scattering into a single GPU-optimized step. This approach not only reduces the intermediate memory requirement drastically but also leverages the GPU's parallel processing capabilities to its full extent.

Efficiency and Extensibility

The implementation showcases impressive benchmarks when compared against Megablocks, showcasing higher throughput and reduced memory usage. In particular, ScatterMoE's implementation allows for a significant reduction in memory allocation during both the forward and backward passes of the network, saving up to 66.2% of memory during training and 53.6% during inference, in comparison to Megablocks. Furthermore, the adaptable architecture of ScatterMoE facilitates easy extensions of SMoE methodologies to other linear transformation-based modules, such as attention layers, without the need for additional complex operations.

Mixture-of-Attention Implementation

Expanding upon the concept of SMoE, the paper also presents an implementation of Mixture-of-Attention (MoA) leveraging the ParallelLinear component. This approach maintains the efficiency benefits of SMoE while extending its applicability to attention mechanisms. By avoiding redundant grouping and scattering operations inherent in existing implementations, ScatterMoE's MoA implementation demonstrates superior performance, especially in high granularity settings where the number of experts per attention head is increased.

Benchmarking and Performance

Through rigorous benchmarking within a 1.5B parameter configuration and comparison with Megablocks and a naive HuggingFace implementation, ScatterMoE evidences its superiority in throughput and memory efficiency. Notably, it achieved up to 38.1% higher throughput in training settings against Megablocks and demonstrated favorable scaling with increasing granularity, which is crucial for optimizing the performance of highly granular SMoE models.

Future Implications

The development of ScatterMoE represents a significant step towards more efficient and scalable implementation of SMoE models on GPUs. Its ability to minimize memory footprint while maximizing throughput without compromising on flexibility or extensibility sets a new benchmark for future developments in the field. The paper speculates that the principles and methodologies introduced by ScatterMoE will facilitate broader adoption and more innovative uses of SMoE models in various domains, potentially leading to the development of even larger and more complex neural network architectures optimized for GPU environments.

Conclusion

In conclusion, ScatterMoE offers a compelling advancement in the implementation of Sparse Mixture-of-Experts on GPUs, addressing critical efficiency and scalability challenges. By introducing the ParallelLinear component and showcasing its application in Mixture-of-Attention, this work lays the groundwork for future research and development in optimizing neural network architectures for parallel computing environments.