GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

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 (27)

• 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
• Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
• Scaling Laws for Neural Language Models
• The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks
• Deep Big Simple Neural Nets Excel on Handwritten Digit Recognition
• Theano: new features and speed improvements
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Massively Multilingual Neural Machine Translation in the Wild: Findings and Challenges
• Lingvo: a Modular and Scalable Framework for Sequence-to-Sequence Modeling
• Glow: Graph Lowering Compiler Techniques for Neural Networks
• Language Modeling with Deep Transformers
• Scaling description of generalization with number of parameters in deep learning
• Adafactor: Adaptive Learning Rates with Sublinear Memory Cost
• On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization
• Relay: A New IR for Machine Learning Frameworks
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• Summarizing CPU and GPU Design Trends with Product Data
• Measuring the Intrinsic Dimension of Objective Landscapes
• An analytic theory of generalization dynamics and transfer learning in deep linear networks
• Rethinking Parameter Counting in Deep Models: Effective Dimensionality Revisited
• Fast Transformer Decoding: One Write-Head is All You Need
• Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training
• Controlling Computation versus Quality for Neural Sequence Models
• Surprisal-Triggered Conditional Computation with Neural Networks

Cited by in corpus (64)

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Transformers in Vision: A Survey
• LoRA: Low-Rank Adaptation of Large Language Models
• On the Opportunities and Risks of Foundation Models
• The Pile: An 800GB Dataset of Diverse Text for Language Modeling
• Beyond English-Centric Multilingual Machine Translation
• Deep Learning Based Text Classification: A Comprehensive Review
• Compute Trends Across Three Eras of Machine Learning
• Intelligent Computing: The Latest Advances, Challenges and Future
• The NLP Cookbook: Modern Recipes for Transformer based Deep Learning Architectures
• BigSSL: Exploring the Frontier of Large-Scale Semi-Supervised Learning for Automatic Speech Recognition
• MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer
• Carbon Emissions and Large Neural Network Training
• Spectral Bias and Task-Model Alignment Explain Generalization in Kernel Regression and Infinitely Wide Neural Networks
• ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale
• Gradient Vaccine: Investigating and Improving Multi-task Optimization in Massively Multilingual Models
• A Full-Stack Search Technique for Domain Optimized Deep Learning Accelerators
• M6: A Chinese Multimodal Pretrainer
• TensorOpt: Exploring the Tradeoffs in Distributed DNN Training with Auto-Parallelism
• A Survey on Symbolic Knowledge Distillation of Large Language Models
• Mask and Reason: Pre-Training Knowledge Graph Transformers for Complex Logical Queries
• Enabling Compute-Communication Overlap in Distributed Deep Learning Training Platforms
• FastMoE: A Fast Mixture-of-Expert Training System
• AMMUS : A Survey of Transformer-based Pretrained Models in Natural Language Processing
• A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training
• FlexMoE: Scaling Large-scale Sparse Pre-trained Model Training via Dynamic Device Placement
• Scalable and Efficient MoE Training for Multitask Multilingual Models
• Transforming Visual Scene Graphs to Image Captions
• Taming Sparsely Activated Transformer with Stochastic Experts
• Yuan 1.0: Large-Scale Pre-trained Language Model in Zero-Shot and Few-Shot Learning
• Safety Analysis in the Era of Large Language Models: A Case Study of STPA using ChatGPT
• XLM-T: Scaling up Multilingual Machine Translation with Pretrained Cross-lingual Transformer Encoders
• Reasoning Beyond Limits: Advances and Open Problems for LLMs
• Scaling Laws for Neural Machine Translation
• ReCycle: Resilient Training of Large DNNs using Pipeline Adaptation
• PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers
• M6-T: Exploring Sparse Expert Models and Beyond
• Scaling End-to-End Models for Large-Scale Multilingual ASR
• FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models
• Adaptive Sparse Transformer for Multilingual Translation
• Rethinking embedding coupling in pre-trained language models
• PaSE: Parallelization Strategies for Efficient DNN Training
• MPipeMoE: Memory Efficient MoE for Pre-trained Models with Adaptive Pipeline Parallelism
• Optimizing DNN Compilation for Distributed Training with Joint OP and Tensor Fusion
• SpeechMoE: Scaling to Large Acoustic Models with Dynamic Routing Mixture of Experts
• Auto-MAP: A DQN Framework for Exploring Distributed Execution Plans for DNN Workloads
• 8-bit Optimizers via Block-wise Quantization
• A Closer Look at Codistillation for Distributed Training
• Lite Training Strategies for Portuguese-English and English-Portuguese Translation
• Layered gradient accumulation and modular pipeline parallelism: fast and efficient training of large language models
• Automatic Graph Partitioning for Very Large-scale Deep Learning
• Go Wider Instead of Deeper
• Current Limitations of Language Models: What You Need is Retrieval
• Towards Structured Dynamic Sparse Pre-Training of BERT
• Towards More Effective and Economic Sparsely-Activated Model
• Meta Back-translation
• Mixture-of-Experts Graph Transformers for Interpretable Particle Collision Detection
• Memory-efficient array redistribution through portable collective communication
• Towards a Universal NLG for Dialogue Systems and Simulators with Future Bridging
• Low-Precision Hardware Architectures Meet Recommendation Model Inference at Scale
• AutoFL: Enabling Heterogeneity-Aware Energy Efficient Federated Learning
• A cost-benefit analysis of cross-lingual transfer methods
• CARLS: Cross-platform Asynchronous Representation Learning System
• Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference