Mesh-TensorFlow: Deep Learning for Supercomputers

Abstract

Batch-splitting (data-parallelism) is the dominant distributed Deep Neural Network (DNN) training strategy, due to its universal applicability and its amenability to Single-Program-Multiple-Data (SPMD) programming. However, batch-splitting suffers from problems including the inability to train very large models (due to memory constraints), high latency, and inefficiency at small batch sizes. All of these can be solved by more general distribution strategies (model-parallelism). Unfortunately, efficient model-parallel algorithms tend to be complicated to discover, describe, and to implement, particularly on large clusters. We introduce Mesh-TensorFlow, a language for specifying a general class of distributed tensor computations. Where data-parallelism can be viewed as splitting tensors and operations along the "batch" dimension, in Mesh-TensorFlow, the user can specify any tensor-dimensions to be split across any dimensions of a multi-dimensional mesh of processors. A Mesh-TensorFlow graph compiles into a SPMD program consisting of parallel operations coupled with collective communication primitives such as Allreduce. We use Mesh-TensorFlow to implement an efficient data-parallel, model-parallel version of the Transformer sequence-to-sequence model. Using TPU meshes of up to 512 cores, we train Transformer models with up to 5 billion parameters, surpassing state of the art results on WMT'14 English-to-French translation task and the one-billion-word language modeling benchmark. Mesh-Tensorflow is available at https://github.com/tensorflow/mesh .

References in corpus (6)

• Exploring the Limits of Language Modeling
• Minimizing Communication in Linear Algebra
• Beyond Data and Model Parallelism for Deep Neural Networks
• Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks
• Communication-Optimal Convolutional Neural Nets
• Communication-Avoiding Optimization Methods for Distributed Massive-Scale Sparse Inverse Covariance Estimation

Cited by in corpus (17)

• Scaling Laws for Neural Language Models
• BigSSL: Exploring the Frontier of Large-Scale Semi-Supervised Learning for Automatic Speech Recognition
• Charformer: Fast Character Transformers via Gradient-based Subword Tokenization
• Bayesian Layers: A Module for Neural Network Uncertainty
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
• FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver
• FastMoE: A Fast Mixture-of-Expert Training System
• PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
• Are Pre-trained Convolutions Better than Pre-trained Transformers?
• Scalable and Efficient MoE Training for Multitask Multilingual Models
• Accelerating Sparse Deep Neural Networks
• Bridging the Gap for Tokenizer-Free Language Models
• Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training
• Making EfficientNet More Efficient: Exploring Batch-Independent Normalization, Group Convolutions and Reduced Resolution Training
• A Linear Algebraic Approach to Model Parallelism in Deep Learning
• Benchmarking down-scaled (not so large) pre-trained language models