Scalable and Efficient MoE Training for Multitask Multilingual Models

Abstract

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.

References in corpus (10)

• Scaling Laws for Neural Language Models
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Multilingual Denoising Pre-training for Neural Machine Translation
• Beyond English-Centric Multilingual Machine Translation
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• On Layer Normalization in the Transformer Architecture
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• Mesh-TensorFlow: Deep Learning for Supercomputers
• Deep Encoder, Shallow Decoder: Reevaluating Non-autoregressive Machine Translation

Cited by in corpus (5)

• Alexa Teacher Model: Pretraining and Distilling Multi-Billion-Parameter Encoders for Natural Language Understanding Systems
• A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training
• Taming Sparsely Activated Transformer with Stochastic Experts
• MPipeMoE: Memory Efficient MoE for Pre-trained Models with Adaptive Pipeline Parallelism
• The Efficiency Misnomer