Scattered Mixture-of-Experts Implementation

2024/03/13 by Shawn Tan, Tan, Shawn, Yikang Shen +5 · 1 citation
Earth and Planetary Sciences · #Environmental Monitoring and Data Management

paper · pdf · doi:10.48550/arxiv.2403.08245

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.

Cited by

• SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

Related

• MegaBlocks: Efficient Sparse Training with Mixture-of-Experts
• MoBiLE: Efficient Mixture-of-Experts Inference on Consumer GPU with Mixture of Big Little Experts
• Janus: Disaggregating Attention and Experts for Scalable MoE Inference
• HeterMoE: Efficient Training of Mixture-of-Experts Models on Heterogeneous GPUs
• Mixture of Diverse Size Experts