[Paper Review] Mesh-TensorFlow: Deep Learning for Supercomputers

Noam Shazeer, Youlong Cheng+10 more|arXiv (Cornell University)|Nov 5, 2018

Advanced Neural Network Applications21 references52 citations

TL;DR

Mesh-TensorFlow introduces a language to specify distributed tensor computations across multi-dimensional processor meshes, enabling scalable model- and data-parallel training of large models such as Transformers on TPUs, achieving state-of-the-art results.

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 .

Motivation & Objective

• Motivate scalable training beyond pure data-parallelism to address memory bottlenecks and latency in large DNNs.
• Introduce Mesh-TensorFlow as a language to specify distributed tensor computations across a multi-dimensional processor mesh.
• Show how to compile Mesh-TensorFlow graphs into SPMD programs with collective communications.
• Demonstrate practical benefits by training Transformer models with billions of parameters on TPU clusters.

Proposed method

• Define named tensor dimensions and a multi-dimensional mesh of processors.
• Specify a global computation layout mapping tensor dimensions to mesh dimensions.
• Represent each tensor as a slice per processor and implement operations as local computations with possible collectives (Allreduce).
• Use einsum-style operations (Einsum) and reductions to express matrix multiplications and contractions across distributed shards.
• Provide layouts (data-parallel, model-parallel, and mixed) and analyze their performance trade-offs in terms of computation, communication, and memory.

Experimental results

Research questions

• RQ1 Can Mesh-TensorFlow express and efficiently execute a broad class of distributed tensor computations beyond data-parallelism?
• RQ2 How do different distribution layouts (data-parallel, model-parallel, and hybrid) affect communication, memory, and scalability on large TPU meshes?
• RQ3 What performance and model-size benefits can be achieved by applying Mesh-TensorFlow to Transformer-like architectures on large clusters?

Key findings

• A Mesh-TensorFlow graph compiles into an SPMD program with parallel operations and MPI-like collectives.
• Data-parallel, model-parallel, and hybrid layouts enable training of Transformer models up to several billions of parameters on TPU meshes.
• Training Transformer models with up to 5 billion parameters on up to 512 cores achieved state-of-the-art results on WMT’14 En–Fr translation and the One Billion Word language modeling benchmark.
• Using multi-dimensional meshes (e.g., 2D 512-core TPUs) maintained substantial compute efficiency (over 50% of peak) while scaling model size and attention heads.
• The approach allows combining data- and model-parallelism to scale both batch size and model dimensions proportionally with the number of processors.