  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 .
