Scattered Mixture-of-Experts Implementation
This work addresses performance bottlenecks in SMoE for researchers and practitioners using large-scale models, but it is incremental as it builds upon existing implementations.
The paper tackles the problem of inefficient GPU implementation for Sparse Mixture-of-Experts (SMoE) by introducing ScatterMoE, which improves inference and training speed and reduces memory footprint, achieving higher throughput and lower memory usage compared to Megablocks.
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.