[Paper Review] Scalable and Efficient MoE Training for Multitask Multilingual Models

Young Jin Kim, Ammar Ahmad Awan+7 more|arXiv (Cornell University)|Sep 22, 2021

Topic ModelingComputer Science26 references33 citations

TL;DR

The paper introduces DeepSpeed MoE, a scalable system for training large Mixture-of-Experts models across multitask multilingual settings, along with training techniques (RTS, AoE, pruning) and shows strong MT and multilingual generation results with Z-code M3 up to 10B parameters.

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.

Motivation & Objective

• Enable scalable training of multitask multilingual Mixture-of-Experts models reaching billions/trillions of parameters.
• Develop training and system techniques to improve sample efficiency and runtime efficiency.
• Showcase practical models (Z-code M3) achieving strong MT and multilingual generation across 50 languages.

Proposed method

• Develop DeepSpeed MoE system enabling five parallelism forms (data, expert, model, tensor-slicing, ZeRO) and ZeRO-Offload to exceed GPU memory limits.
• Introduce Random Token Selection to mitigate token-bias in MoE routing.
• Propose Aggregation of Experts (AoE) combining checkpoints to create larger expert pools for initialization and training.
• Explore experts pruning strategies (random and utilization-based) for faster inference.
• Train multitask multilingual models (MT, DAE, ELECTRA, MLM) within a single objective by summing task losses.
• Use transformer encoder-decoder architecture with MoE layers placed every other layer and a 250k SentencePiece vocabulary.

Experimental results

Research questions

• RQ1 How can MoE architectures be scaled to trillions of parameters for multitask multilingual training?
• RQ2 Can DeepSpeed MoE overcome GPU memory limits to enable larger base models and more experts?
• RQ3 Do MoE-based multitask multilingual training regimes improve downstream MT and NLG tasks compared to dense baselines?
• RQ4 What training techniques maximize sample efficiency and inference efficiency in large MoE models?
• RQ5 What is the impact of multitask objectives on multilingual translation and generation quality?

Key findings

• DeepSpeed MoE enables near-linear throughput scaling across GPUs and supports model sizes beyond GPU memory via ZeRO-Offload.
• RTS reduces token-position bias and improves convergence speed and regularization in MoE training.
• AoE allows creating larger effective expert pools by aggregating parameters from checkpoints to initialize larger models.
• Experts pruning yields smaller, efficient inference without sacrificing much performance, via random or utilization-based selection.
• Z-code M3 (10B parameters, 64 experts) outperforms dense baselines and smaller MoE configurations on MT and multilingual tasks, demonstrating strong multitask multilingual capabilities.
• Fine-tuned Z-code M3 models achieve notable improvements on downstream tasks such as Wikilingua and cross-lingual generation.