Mesh-TensorFlow: Deep Learning for Supercomputers
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 [16] 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 Englishto-French translation task and the one-billion-word language modeling benchmark. Mesh-Tensorflow is available at https://github.com/tensorflow/mesh .
collective communication primitives such as Allreduce. We use Mesh-TensorFlow to implement an efficient data-parallel, model-parallel version of the Transformer [16] 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.
Given the author's additional information about their plans to open-source the language described in this paper, I've upgrade my review. I'm still concerned that the original paper didn't mention anything about this intent to open-source; but I'm willing to give them the benefit of the doubt because if this is open-sourced, it could be very helpful to the community.  Here are some minor typos I found: Capitalize “equation”, “table”, “section” as needed throughout the paper. Line 112: “pleasant” to “pleasantly” (and throughout the paper where needed) 219: “tpu” to “TPU” 222: “(3)” to “Table 3”  ORIGINAL REVIEW This paper introduces a convenient notation for compactly describing mappings from a tensor representation of a learning problem’s data and model parameters to an n-dimensional mesh of compute nodes.  The notation is flexible enough to capture both data-parallel and model-parallel mappings, as well as mixtures of the two. The paper uses these mappings to derive first order approximations to the computation and communication times required to perform a single backpropagation update (assuming an optimized communication library like MPI).  The paper then uses these estimates to develop a strategy for simultaneously growing the model size and the number of compute nodes used such that communication does not become a major performance bottleneck.  The paper then uses this strategy to explore the parallel performance of the Transformer sequence transduction model.  The strategy allowed the authors to train a 5B parameter system on 256 nodes and get state-of-the-art performance on the WMT14 En-Fr data set.  When I first read this paper, I thought that the authors meant that they had created a new programming language for easily and efficiently mapping machine learning problems to parallel computer nodes.  I was hoping that this “language” would be made available to the reader and that it would include some automatic optimization of the mappings; but as I read more of the paper, I realized that the paper only presented a notation for representing the mapping.  If a new programming  language (or tool) has been created, the authors should make that more clear and maybe give pointers to it.  If it is only a notational convenience, the authors should make that more clear too.  If the contribution of Section 4 is only to introduce a notational convenience, the authors should explain the significance more clearly and/or concretely.  For example, the notation in Algorithm 1 is simply Einstein Summation Notation in another form.  It doesn’t add significantly to the reader’s understanding of the problem.  Furthermore, the observations made in Table 1 are fairly well-known from a parallel computing point of view - most practitioners realize this computation/communication trade off.  And for those NIPS readers who don’t, the paper has no explanation for where these numbers come from.  The paper could be improved if Section 4 was much shorter and the authors spent time explaining how they calculated the terms in Table 1. It seems that the major result of this paper is a better BLEU score from the ability to run a bigger model.  I’m not sure that is significant enough.  The paper would benefit from clarification in several places.  For example, “b”, “d_x” and “d_y” are never explicitly defined.  Equation (3a) is described and then never used again in the paper.  The paper’s “{b-->0}” notation is never explicitly defined.  Presumable “0” means dimension zero of the compute mesh; but it would be better it the authors specifically said it so that the reader doesn’t have to guess.  Likewise, “F” and “D(X,M)” are never defined.  (I know “F” is an arbitrary function, but the paper should say that instead of making the reader guess.) Also, for example, when two model dimensions map to a single mesh dimension, does that mean the mesh dimension is split 50/50?  It could be split based on some communication minimization, or other, algorithm.  The reader has to guess.  The text mentions four layouts but only shows three in Table 1.  In Table 3, why are two of the BLEU scores missing?  Was it that they couldn’t run?  Ran too slow? Wasn’t enough time? That should be explain.  In Section 5, the reader would benefit from some pictures showing, for example, the matrix multiplication, etc., to help them “see” the mapping.  “SOTA” should be replaced with “state-of-the-art” throughout the paper.