Mesh-TensorFlow: Deep Learning for Supercomputers

By Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter G. Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young and 2 more. 2018

439 citations recorded by the source corpus.

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 this https URL .

Subfields: Computational Mathematics, Computer Vision and Pattern Recognition, Hardware and Architecture, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computational Mathematics

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Efficiency & Compression

Citations in the corpus

Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (2019, methodology): “…research argues that general methods that can leverage additional computation ultimately win out against methods that rely on human expertise [Sutton, 2019; Hestness et al., 2017; Shazeer et al., 2017; Jozefowicz et al., 2016; Mahajan et al., 2018; Shazeer et al., 2018, 2017; Huang et al., 2018b].”
PaLM: Scaling Language Modeling with Pathways (2022, methodology): “Therefore, techniques have arisen for splitting model tensors across accelerators (Shazeer et al., 2018) or alternatively separating layers of the models across accelerators and then pipe-lining activations between the stages (Huang et al.”
Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers (2020, methodology): “We chose these tasks because accuracy continues to improve as models are made larger (Shazeer et al., 2018), trained for more steps (Liu et al.”
ALBERT: A Lite BERT for Self-supervised Learning of Language Representations (2019, background): “Existing solutions to the aforementioned problems include model parallelization (Shazeer et al., 2018; Shoeybi et al., 2019) and clever memory management (Chen et al.”
Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (2019, methodology): “For example, GPipe (Huang et al., 2018) and Mesh-Tensorﬂow (Shazeer et al., 2018) provide frameworks for model parallelism of different kinds.”
Simple, Scalable Adaptation for Neural Machine Translation (2019, background): “This is further exacerbated by the increasing model capacity of state-of-the-art NMT models (Shazeer et al., 2018; Bapna et al., 2018; Huang et al., 2018); full ﬁne-tuning is just too parameter inefﬁcient.”
Massively Multilingual Neural Machine Translation in the Wild: Findings and Challenges (2019, methodology): “Other than the systems and engineering challenges (Jouppi et al., 2017; Huang et al., 2018; Shazeer et al., 2018), training deep and high capacity neural networks poses signiﬁcant trainabil-ity challenges (Montavon et al., 2018).”
One Epoch Is All You Need (2019, background): “GPT (Radford et al., 2018) 100 SPN (Menick & Kalchbrenner, 2018) Not reported BERT (Devlin et al., 2018) 40 Mesh Transformer (Shazeer et al., 2018) 10 Transformer-XL (Dai et al., 2019) Not reported GPT-2 (Radford et al., 2019) Not reported (20 or 100) Sparse Transformer (Child et al., 2019) 70 - 120”
GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (2021, methodology): “We use RMSNorm as in (Zhang & Sennrich, 2019; Shazeer et al., 2018b) instead of standard LayerNorm (Ba et al., 2016).”
Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers (2021, methodology): “Our implementation and experiments are performed in Mesh Tensorflow1 (Shazeer et al., 2018) using the default T5 library2.”
ZeRO: Memory optimizations Toward Training Trillion Parameter Models (2019, methodology): “The largest models in the current literature, the 11B T5 model [5], and MegatronLM 8.3B [4], were both powered by model parallelism, implemented in Mesh-Tensorflow [7] and Megatron-LM[4], respectively.”
Scaling Up Models and Data with t5x and seqio (2022, methodology): “We validated these models by reproducing the T5 models from Raﬀel et al. (2020), originally implemented in Mesh TensorFlow Transformer.”
Pathways: Asynchronous Distributed Dataflow for ML (2022, positive): “P ATHWAYS builds extensively on prior systems, including XLA (TensorFlow, 2019) to represent and execute TPU computations, TensorFlow graphs and executors (Abadi et al., 2016) to represent and execute distributed CPU computations, and Python programming frameworks including JAX (Bradbury et al., 201”
GSPMD: General and Scalable Parallelization for ML Computation Graphs (2021, methodology): “but partitions the graph into pipeline stages . vmap() and TensorFlow’s vectorized_map() .”
PanGu-α: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation (2021, methodology): “One regime is op-level model parallelism [17, 18, 19, 20, 21, 22, 23], which partitions its involved tensors of each operator (layer), as shown in Figure 4(b).”
ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep learning (2021, background): “When models do not fit in device memory, model parallelism3 (MP) [7, 17, 18] and”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, methodology): “Mesh-TensorFlow [30] proposes a language for easily specifying parallelization strategies that combine data and model parallelism.”