A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training
Siddharth Singha, Olatunji Ruwaseb, Ammar Ahmad Awanc, Samyam Rajbhandarib, Yuxiong Hed, Abhinav S Bhatelee
- aDepartment of Computer Science, University of Maryland, College Park, Maryland, USA
- bMicrosoft, Inc., Redmond, Washington, USA
- cMicrosoft, Inc., Redmond, Washington, United States
- dMicrosoft, Inc., Redmond, Washington, United States of America
- eDepartment of Computer Science, University of Maryland, College Park, Maryland, United States of America
Papers open inside SciRah from the original publisher or repository; SciRah never hosts them.
Abstract
Mixture-of-Experts (MoE) is a neural network architecture that adds sparsely activated expert blocks to a base model, increasing the number of parameters without impacting computational costs. However, current distributed deep learning frameworks are limited in their ability to train high-quality MoE models with large base models. In this work, we present DeepSpeed-TED, a novel, three-dimensional, hybrid parallel algorithm that combines data, tensor, and expert parallelism to enable the training of MoE models with 4--8× larger base models than the current state-of-the-art. We also describe memory optimizations in the optimizer step, and communication optimizations that eliminate unnecessary data movement. We implement our approach in DeepSpeed and achieve speedups of 26% over a baseline (i.e. without our communication optimizations) when training a 40 billion parameter MoE model (6.7 billion base model with 16 experts) on 128 V100 GPUs.
TL;DRDeepSpeed-TED is presented, a novel, three-dimensional, hybrid parallel algorithm that combines data, tensor, and expert parallelism to enable the training of MoE models with 4--8× larger base models than the current state-of-the-art. (Semantic Scholar)