GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (2020-06-30T00:00:00.000000Z)

TL;DR

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 and it is demonstrated 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.

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.

Authors

M. Krikun

4 papers

Orhan Firat

8 papers

Dmitry Lepikhin

2 papers

TL;DR

Abstract

Authors

References99 items

Surprisal-Triggered Conditional Computation with Neural Networks

Language Models are Few-Shot Learners

Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training

Rethinking Parameter Counting in Deep Models: Effective Dimensionality Revisited

Controlling Computation versus Quality for Neural Sequence Models

5分で分かる!? 有名論文ナナメ読み：Jacob Devlin et al. : BERT : Pre-training of Deep Bidirectional Transformers for Language Understanding

Scaling Laws for Neural Language Models

Summarizing CPU and GPU Design Trends with Product Data

Fast Transformer Decoding: One Write-Head is All You Need

Unsupervised Cross-lingual Representation Learning at Scale

BlueConnect: Decomposing all-reduce for deep learning on heterogeneous network hierarchy

Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer

Depth-Adaptive Transformer

ZeRO: Memory Optimization Towards Training A Trillion Parameter Models

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Massively Multilingual Neural Machine Translation in the Wild: Findings and Challenges

Learning Deep Transformer Models for Machine Translation

Language Modeling with Deep Transformers

NAS-FPN: Learning Scalable Feature Pyramid Architecture for Object Detection

CondConv: Conditionally Parameterized Convolutions for Efficient Inference

Massively Multilingual Neural Machine Translation

Lingvo: a Modular and Scalable Framework for Sequence-to-Sequence Modeling

Mixture Models for Diverse Machine Translation: Tricks of the Trade

The Evolved Transformer

Scaling description of generalization with number of parameters in deep learning

GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

Mesh-TensorFlow: Deep Learning for Supercomputers

Optimizing data-intensive computations in existing libraries with split annotations

An analytic theory of generalization dynamics and transfer learning in deep linear networks

SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing

Training Deeper Neural Machine Translation Models with Transparent Attention

Beyond Data and Model Parallelism for Deep Neural Networks

Relay: a new IR for machine learning frameworks

PipeDream: Fast and Efficient Pipeline Parallel DNN Training

Glow: Graph Lowering Compiler Techniques for Neural Networks

Exploring the Limits of Weakly Supervised Pretraining

Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks

The Lottery Ticket Hypothesis: Training Pruned Neural Networks

On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization

Measuring the Intrinsic Dimension of Objective Landscapes

Natural TTS Synthesis by Conditioning Wavenet on MEL Spectrogram Predictions

State-of-the-Art Speech Recognition with Sequence-to-Sequence Models

Automatic differentiation in PyTorch

High-dimensional dynamics of generalization error in neural networks

Exploring Generalization in Deep Learning

Attention is All you Need

In-datacenter performance analysis of a tensor processing unit

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Google’s Multilingual Neural Machine Translation System: Enabling Zero-Shot Translation

Understanding deep learning requires rethinking generalization

Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation

WaveNet: A Generative Model for Raw Audio

Identity Mappings in Deep Residual Networks

Multi-Way, Multilingual Neural Machine Translation with a Shared Attention Mechanism

Deep Residual Learning for Image Recognition

Conditional Computation in Neural Networks for faster models

Listen, attend and spell: A neural network for large vocabulary conversational speech recognition

Going deeper with convolutions

Sequence to Sequence Learning with Neural Networks

Neural Machine Translation by Jointly Learning to Align and Translate

Low-Rank Approximations for Conditional Feedforward Computation in Deep Neural Networks

Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation

ImageNet classification with deep convolutional neural networks

Large Scale Distributed Deep Networks

Theano: new features and speed improvements

Large Scale Parallel Document Mining for Machine Translation

Rectified Linear Units Improve Restricted Boltzmann Machines

Deep, Big, Simple Neural Nets for Handwritten Digit Recognition

Large-scale deep unsupervised learning using graphics processors

Bleu: a Method for Automatic Evaluation of Machine Translation

Special-purpose digital hardware for neural networks: An architectural survey

MPI: A Message-Passing Interface Standard

TRANSFER OF TRAINING: A REVIEW AND DIRECTIONS FOR FUTURE RESEARCH

Forum

XLA: Optimizing Compiler for TensorFlow