Scalable and Efficient MoE Training for Multitask Multilingual Models

By Young Jin Kim, Ammar Ahmad Awan, Alexandre Muzio, Andres Felipe Cruz Salinas, Liyang Lu, Amr Hendy, Samyam Rajbhandari, Yuxiong He, Hany Hassan Awadalla. 2021

101 citations recorded by the source corpus.

The Mixture of Experts (MoE) models are an emerging class of sparsely activated deep learning models that have sublinear compute costs with respect to their parameters. In contrast with dense models, the sparse architecture of MoE offers opportunities for drastically growing model size with significant accuracy gain while consuming much lower compute budget. However, supporting large scale MoE training also has its own set of system and modeling challenges. To overcome the challenges and embrace the opportunities of MoE, we first develop a system capable of scaling MoE models efficiently to trillions of parameters. It combines multi-dimensional parallelism and heterogeneous memory technologies harmoniously with MoE to empower 8x larger models on the same hardware compared with existing work. Besides boosting system efficiency, we also present new training methods to improve MoE sample efficiency and leverage expert pruning strategy to improve inference time efficiency. By combining the efficient system and training methods, we are able to significantly scale up large multitask multilingual models for language generation which results in a great improvement in model accuracy. A model trained with 10 billion parameters on 50 languages can achieve state-of-the-art performance in Machine Translation (MT) and multilingual natural language generation tasks. The system support of efficient MoE training has been implemented and open-sourced with the DeepSpeed library.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Natural Language Processing

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Natural Language Processing

Citations in the corpus

Unified Scaling Laws for Routed Language Models (2022, methodology): “We note there are several libraries available for supporting large-scale Routing, including DeepSpeed [Kim et al., 2021, Rajbhandari et al., 2022] and GSPMD [Xu et al., 2021].”
A Review of Sparse Expert Models in Deep Learning (2022, background): “This system yielded 10 × inference improvements (Rajbhandari et al., 2022) and state-of-the-art translation (Kim et al., 2021) – increasing the practicality of these models for production services.”
The Efficiency Misnomer (2021, methodology): “FLOPs: A widely used metric as the proxy for the computational cost of a model is the number of ﬂoating-point multiplication-and-addition operations (Johnson, 2018; Kim et al., 2021; Arnab et al., 2021; Tay et al., 2021a;b; Narayanan et al., 2021; Liu et al., 2021).”
Mixture of Quantized Experts (MoQE): Complementary Effect of Low-bit Quantization and Robustness (2023, methodology): “The Mixture-of-Experts (MoE) architecture efficiently increase the number of model parameters, while maintaining a sub-linear increase in computational requirements by activating only a few small number of experts at a time (Lepikhin et al., 2020; Fedus et al., 2021; Kim et al., 2021; Artetxe et al.”
DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale (2022, methodology): “To achieve low latency and high throughput at an unprecedented scale for MoE, we design our inference system to minimize the critical data path per device, maximize the achievable aggregate memory bandwidth, and offer ample aggregate memory simultaneously to enable massive model sizes by using (1) e”
Tutel: Adaptive Mixture-of-Experts at Scale (2022, methodology): “While GShard [18] provides a computation logic that ensures algorithmic correctness of MoE, several popular MoE frameworks [17, 33] follow the same logic but perform poorly at a large scale.”
Designing Effective Sparse Expert Models (2022, methodology): “…that the recent signiﬁcant successes of sparse expert models have often been in settings with a lot of data and no distribution shift – two examples being language modeling/span corruption and machine translation (Shazeer et al., 2017; Lepikhin et al., 2020; Kim et al., 2021; Fedus et al., 2021).”
ST-MoE: Designing Stable and Transferable Sparse Expert Models (2022, methodology): “…that the recent signiﬁcant successes of sparse expert models have often been in settings with a lot of data and no distribution shift – two examples being language modeling/span corruption and machine translation (Shazeer et al., 2017; Lepikhin et al., 2020; Kim et al., 2021; Fedus et al., 2021).”
Taming Sparsely Activated Transformer with Stochastic Experts (2021, methodology): “Moreover, our model achieves state-of-the-art results on these tasks — the same BLEU score that is achieved by the Z-code MoE model (Kim et al., 2021) with 5.”
A Survey on Mixture of Experts in Large Language Models (2024, unknown): “Kim et al. [74] suggest randomizing token prioritization within sequences to mitigate routing bias towards early-positioned tokens.”