GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
arXiv:2006.16668
Abstract
Neural network scaling has been critical for improving the model quality in many real-world machine learning applications with vast amounts of training data and compute. Although this trend of scaling is affirmed to be a sure-fire approach for better model quality, there are challenges on the path such as the computation cost, ease of programming, and efficient implementation on parallel devices. GShard is a module composed of a set of lightweight annotation APIs and an extension to the XLA compiler. It provides an elegant way to express a wide range of parallel computation patterns with minimal changes to the existing model code. GShard enabled us to scale up multilingual neural machine translation Transformer model with Sparsely-Gated Mixture-of-Experts beyond 600 billion parameters using automatic sharding. We demonstrate that such a giant model can efficiently be trained on 2048 TPU v3 accelerators in 4 days to achieve far superior quality for translation from 100 languages to English compared to the prior art.
References in corpus (10)
- Sequence to Sequence Learning with Neural Networks
- Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation
- Language Models are Few-Shot Learners
- Scaling Laws for Neural Language Models
- Theano: new features and speed improvements
- Massively Multilingual Neural Machine Translation in the Wild: Findings and Challenges
- Lingvo: a Modular and Scalable Framework for Sequence-to-Sequence Modeling
- Fast Transformer Decoding: One Write-Head is All You Need
- Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training
- Surprisal-Triggered Conditional Computation with Neural Networks