Source-linked AI summary
Mesh-TensorFlow: Deep Learning for Supercomputers
Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, Blake Hechtman
TL;DR
Large-scale DNN training needs distribution strategies beyond batch-splitting, whose memory and synchronization demands can limit very large models. Mesh-TensorFlow provides a language for splitting arbitrary tensor dimensions across multidimensional processor meshes and compiling the result to SPMD programs. Using it, the authors trained Transformer models with up to 5 billion parameters on 512-core TPU meshes and surpassed state-of-the-art results on two benchmarks.
Problem
Batch-splitting can be impossible or inefficient for very large models because of parameter or activation memory and parameter synchronization costs, while model-parallel strategies are complicated to specify and implement.
Method
Mesh-TensorFlow specifies distributed tensor computations by mapping arbitrary tensor dimensions onto multidimensional processor meshes and compiling graphs into SPMD programs with collective communication.
Results
The approach trained Transformer models with up to 5 billion parameters on TPU meshes of up to 512 cores, surpassing state-of-the-art results on WMT’14 English-to-French translation and one-billion-word language modeling.
Takeaways & Limitations
Mesh-TensorFlow supports large-scale Transformer training with combined data-parallel and model-parallel computation.
Takeaways & Limitations
The paper focuses on clusters of identical, reliable processors with local memory, and physical network topology affects communication performance.
Abstract
from arXiv · showhide
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 .
1 Introduction
Batch-splitting is widely used because it applies broadly and supports SPMD programming, but it can be impractical for very large models. Mesh-TensorFlow addresses these distribution challenges with a general tensor-computation language and demonstrates large-scale Transformer training.
- Motivation: Batch-splitting can become impossible or inefficient when parameter and activation memory or parameter synchronization time grows too large.Model-parallelism can address these issues, but existing approaches may be complicated to specify, implement, compile, and optimize.
- Approach: Mesh-TensorFlow lets users split arbitrary tensor dimensions across dimensions of a multidimensional processor mesh.Its graphs compile into SPMD programs combining parallel operations with collective communication primitives such as Allreduce.
- Approach: Mesh-TensorFlow supports an efficient data-parallel and model-parallel implementation of the Transformer sequence-to-sequence model.The implementation uses collective communication primitives including Allreduce.
- Results: 5 billion parameters were trained on TPU meshes of up to 512 cores, surpassing state-of-the-art results on WMT’14 English-to-French translation and one-billion-word language modeling.The paper focuses on clusters of identical, reliable processors with local memory; physical network topology still affects communication performance.
4 Mesh-TensorFlow: Beyond Batch Splitting
Mesh-TensorFlow extends batch-splitting by assigning named tensor dimensions to named mesh dimensions. This representation preserves SPMD execution while allowing tensors and operations to be distributed across broader processor layouts.
- Generalization: Mesh-TensorFlow generalizes batch-splitting by allowing computation to be split across different tensor dimensions.Each tensor is represented by a slice on every processor, and each operation runs on every processor with occasional collective communication.
- Named dimensions: Named tensor dimensions let the same logical dimension, such as batch, be split consistently across tensors and operations.A tensor cannot contain two identically named dimensions.
- Processor mesh: Mesh-TensorFlow organizes processors as an n-dimensional mesh whose dimensions are also named.This replaces an unstructured processor set with a multidimensional naming abstraction.
- Computation layout: A computation layout maps tensor dimensions to mesh dimensions to specify which dimensions are split across processors.For example, mapping batch to an all_processors mesh dimension expresses batch-splitting while leaving other tensors replicated.
- Tensor representation: Each tensor’s layout restricts the global layout to its dimensions, producing processor-specific stripes or full replication.The current implementation requires mapped tensor dimensions to divide evenly by their corresponding mesh dimensions.
6 Operation Implementation
Mesh-TensorFlow implements distributed tensor operations as parallel per-processor computations with collective communication when layouts split reduced dimensions. Named tensor dimensions and computation layouts let users combine data- and model-parallelism across multidimensional processor meshes, while layout legality and cost determine efficiency.
- Operation implementations: Mesh-TensorFlow operations run in parallel on every processor and use collective communication when reduced tensor dimensions are split.Component-wise operations are local; reductions and einsums may require MPI-allreduce across corresponding mesh dimensions.
- Operation implementations: Einsum expresses matrix multiplication, reductions, and broadcasts through named dimensions, with allreduce across mesh dimensions corresponding to reduced-out tensor dimensions.The operation broadcasts inputs to their union of dimensions, multiplies component-wise, and reduces dimensions absent from the output.
- Operation implementations: Reshape can require network communication when input and output layouts differ, including allgather when a split input dimension becomes unsplit.Changing dimension names can alter layouts even when dimension sizes remain unchanged.
- Language and compilation: The language extends TensorFlow with named tensor dimensions and statically known shapes, while its Python implementation lowers graphs to TPU SPMD or multi-CPU/GPU code.Users define mathematical computations separately from layouts, which can preserve identical results while changing performance characteristics.
- Example layouts: Data-parallel layouts replicate parameters and split activations across batch, whereas model-parallel layouts split hidden activations and parameters across the hidden dimension.Data-parallel training communicates parameter gradients; model-parallel computation communicates reductions over split hidden dimensions.
- Example layouts: A two-dimensional mesh combines batch and hidden-unit splitting, using processor rows for batch fractions and columns for hidden-unit fractions.These layouts use partitioned allreduces over the mesh dimension associated with the reduced tensor dimension.
- Layout efficiency and legality: Efficient layouts split expensive operations across mesh dimensions; replicating all computation saves neither time nor memory, and some layouts are illegal when two tensor dimensions share one mesh dimension.A stated rule is that each expensive einsum should have an input dimension split across each batch dimension.
- Layout efficiency and legality: 28: A three-dimensional mesh can increase processor count cubically while requiring only linear increases in batch and layer sizes for efficiency.The two-dimensional layout permits quadratic processor scaling with linear increases in batch size and hidden-layer size.
9 Model-Parallel "Transformer"
Mesh-TensorFlow enables model-parallel Transformer layouts by splitting vocabulary, feed-forward, and attention-head dimensions across TPU meshes. Scaling these dimensions supported efficient training of billion-parameter models and improved benchmark results.
- Model-parallel layout: Mesh-TensorFlow splits vocabulary, feed-forward hidden-layer, and attention-head dimensions across all processors.The layout works because each expensive operation contains exactly one split dimension, while no tensor has more than one.
- Model scaling: Scaling hidden layers and attention heads with larger TPU clusters produced similar performance characteristics and benefited model quality.The vocabulary size was not increased.
- Data and model parallelism: The combined layout splits the batch across one TPU-mesh dimension and model dimensions across the other.Performance remains constant when batch size scales with r and model dimensions scale with c.
- Model scaling: 512 cores supported feed-forward hidden dimensions up to 262144 and 256 attention heads while maintaining over 50% computational efficiency.The largest models achieved 6 PFLOP/s out of a maximum 11.5 PFLOP/s on 16x32 TPUv2 meshes.
- Experiments and results: The largest billion-word language model had 4.9B parameters, achieved dev-perplexity 24.0, and reached 23.5 after multiplying logits by 0.9.It trained for 13 hours on a 512-core TPUv2 cluster and was reported as the best published result on the dataset.
- Experiments and results: The largest WMT14 English-to-French model had 2.9B parameters and achieved BLEU score 43.9, reported as the best published result to date.The model trained for 22 hours on a 128-core TPUv2 cluster.
10 Related Work
Related work applies iteration-space partitioning and cost modeling to distributed deep learning, but Mesh-TensorFlow emphasizes named tensor dimensions and broader mapping exploration.
- Iteration-space partitioning: Distributed matrix-multiplication methods partition iteration space to minimize communication, and Mesh-TensorFlow can express mappings such as 3D and 2.5D algorithms.It also targets rectangular, sparse, and direct-convolution mappings, sometimes with higher memory requirements.
- Iteration-space partitioning: Mesh-TensorFlow lets users name dimensions to simplify layout specification and mapping exploration across composed multiplications.Existing approaches often require specifying each matrix’s data layout separately.
- Deep-learning systems: Prior deep-learning work combined data and model parallelism analytically or used cost modeling, but explored narrower models or potentially suboptimal output-based mappings.One approach supported arbitrary processor grids without implementation; another was implemented but could be communication-suboptimal.
11 Future Work
Future work targets automated layout optimization, broader model and operation coverage, and SPMD execution on CPU/GPU clusters.
- Future work: Potential future work includes automated search for optimal computation layouts.
- Future work: The library needs implementations of additional models and operations, including convolutions on spatially partitioned tensors with halo-region communication.
- Future work: SPMD programming on CPU/GPU clusters is another identified development area.
12 Conclusion
The paper introduces Mesh-TensorFlow for broad SPMD distributed tensor computations and demonstrates it by training very large Transformers on up to 512-core clusters.
- Conclusion: Mesh-TensorFlow facilitates a broad class of SPMD distributed tensor computations.
- Conclusion: Transformer models with up to 5 billion parameters were trained on clusters of up to 512 cores, establishing state-of-the-art WMT14 En-Fr and One Billion Word results.
A Illustrations for the Two Fully-Connected Layers Example
The example progresses through two fully connected layers, then illustrates data-, model-, and mixed-parallel layouts and iteration-space partitionings. Mesh-TensorFlow supports fine-grained partitioning that can improve communication efficiency beyond owner-compute strategies.
- Overall computation: The computation multiplies x by w, adds bias, applies component-wise Relu to form h, and multiplies h by v to produce y.The intermediate matrix xw and component-wise operations are omitted from the figure.
- Data-parallel layout: Data parallelism splits the batch dimension across two processors while fully replicating w and v.The processor ranks label the matrix slices stored on each processor.
- Model-parallel layout: Model parallelism splits the hidden dimension across two processors while fully replicating x and y.
- Mixed layouts: Mixed parallelism arranges processors into 2-by-2 or 2-by-2-by-2 meshes and labels owned matrix slices with serialized processor ranks.These layouts use 4 and 8 processors, respectively.
- Iteration-space partitioning: Matrix multiplication partitions its iteration space into 1D, 2D, or 3D schemes according to the number of split axes.The 1D example splits j, the 2D example splits i and j, and the 3D example splits all three dimensions.
- Communication-aware parallelism: Mesh-TensorFlow expresses fine-grained parallelism by combining replicated tensors with multiple tensor layouts to split iteration spaces across multiple dimensions.This can extend beyond owner-compute strategies, which assign processors responsibility for computations associated with owned matrix chunks.
C Random Samples from Transformer Language Models Trained on Billion-Word Language Modeling Benchmark
Random samples from a billion-word language-model benchmark are seeded with “According to Ray Kurzweil” and continued by the model. The samples are mostly grammatical, while larger models exhibit more world knowledge.
- The samples are randomly generated from the Transformer language models described in the paper.All sentences use the initial words “According to Ray Kurzweil.”
- Larger models produce mostly grammatical sentences with more world knowledge.
- The displayed continuations range across military, economic, technology, scientific, and historical topics.
C.2 Model with 0.37B Parameters, PPL=28.9
The 0.37B-parameter model produces random continuations from the seed “According to Ray Kurzweil.” The examples are generally grammatical but contain factual and semantic inconsistencies.
- The continuations include claims about internet use, computing, employment, entertainment, and archaeology.
- The examples combine fluent phrasing with implausible or internally inconsistent associations.
C.3 Model with 1.28B Parameters, PPL=25.1
The 1.28B-parameter model is sampled from the same “According to Ray Kurzweil” prompt. Its examples remain largely grammatical and cover varied factual topics.
- The examples are presented as random continuations of the fixed seed “According to Ray Kurzweil.”
- The continuations address emissions, Google, computing history, biology, and video games.