Elamur Research
☰
home / papers / 2018 / Mesh-TensorFlow: Deep Learning for Supercomputers

Mesh-TensorFlow: Deep Learning for Supercomputers

Written for the Research library (claude-opus-5-5) from the paper text, in our own words: the license "arXiv — non-exclusive license to distribute" does not allow redistributing the text itself. Original — arxiv.org/pdf/1811.02084

Summary

Data-parallel training, which splits only the batch across devices, cannot handle models whose parameters exceed a single device's memory. It also scales poorly when batches are small. Model-parallel schemes avoid these limits, but they are usually hard to design and code by hand. The authors present Mesh-TensorFlow, a TensorFlow-like Python library in which every tensor dimension carries a name. Hardware is treated as a logical n-dimensional grid of processors. A single "computation layout" assigns chosen tensor dimensions to grid dimensions, so data parallelism, model parallelism, or a mix of both can be selected without changing model code. The library lowers the graph to an SPMD program in which each processor runs local operations and communicates through collectives such as Allreduce, allgather and alltoall. Using this approach, the team trained Transformer models with up to roughly 5B parameters on TPUv2 clusters of up to 512 cores. These models set new records on WMT14 En-Fr and the One Billion Word benchmark. The work shows that scaling model size, not only data, can be made practical.

Key points

  • The main abstraction is a mapping from named tensor dimensions to named mesh dimensions. Einsum and reductions are implemented as local operations followed by an Allreduce, and reshapes may require allgather or alltoall.
  • On the One Billion Word benchmark, a 4.9B-parameter decoder trained for 13 hours on a 512-core TPUv2 cluster. It reached a dev perplexity of 24.0, which fell to 23.5 when the logits were scaled by 0.9. The previous best single DNN scored 28.0.
  • On WMT14 En-Fr, a 2.9B-parameter model trained for 3 epochs, taking 22 hours on 128 cores. It scored 43.9 BLEU, measured with sacrebleu. Gains on WMT14 En-De were smaller, likely because its training data is limited.
  • On the Tensor2Tensor dataset languagemodel_wiki_noref_v128k_l1k, which contains about 5B tokens of Wikipedia text, subword perplexity kept improving with model size, reaching 6.01 at 4.9B parameters.
  • The code is released at github.com/tensorflow/mesh. Planned future work includes automatic layout search, spatially partitioned convolutions, and SPMD support for CPU/GPU clusters.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Computing & GPUs

Cited by

132 papers of our library cite it (of the 140673 whose bibliographies we read); 10 of them have a page here, 122 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.1811.02084 · arXiv 1811.02084

@misc{shazeer2018mesh,
  title = {{Mesh-TensorFlow: Deep Learning for Supercomputers}},
  author = {Shazeer, Noam and Cheng, Youlong and Parmar, Niki and Tran, Dustin and Vaswani, Ashish and Koanantakool, Penporn and Hawkins, Peter and Lee, HyoukJoong and Hong, Mingsheng and Young, Cliff and Sepassi, Ryan and Hechtman, Blake},
  year = {2018},
  doi = {10.48550/arxiv.1811.02084},
  eprint = {1811.02084},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.1811.02084}
}

References

26 references, 0 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Jarle Berntsen. “Communication efficient matrix multiplication on hypercubes”. In: Parallel computing 12.3 (1989), pp. 335–342. find
  2. M. Wolfe. “More Iteration Space Tiling”. In: Proceedings of the 1989 ACM/IEEE Conference on Supercomputing. Supercomputing ’89. Reno, Nevada, USA: ACM, 1989, pp. 655–664. ISBN : 0-89791-341-8. DOI : 10.1145/76263.76337. URL : http://doi.acm.org/10.1145/76263.76337. — [doi:10.1145/76263.76337](https://doi.org/10.1145/76263.76337) ↗
  3. Alok Aggarwal, Ashok K. Chandra, and Marc Snir. “Communication Complexity of PRAMs”. In: Theor. Comput. Sci. 71.1 (Mar. 1990), pp. 3–28. ISSN : 0304-3975. DOI : 10.1016/0304-3975(90)90188-N. URL : http://dx.doi.org/10.1016/0304-3975(90)90188-N. — [doi:10.1016/0304-3975(90)90188-n](https://doi.org/10.1016/0304-3975%2890%2990188-n) ↗
  4. Dror Irony, Sivan Toledo, and Alexander Tiskin. “Communication lower bounds for distributed-memory matrix multiplication”. In: Journal of Parallel and Distributed Computing 64.9 (2004), pp. 1017–1026. find
  5. Pitch Patarasuk and Xin Yuan. “Bandwidth optimal all-reduce algorithms for clusters of workstations”. In: 69 (Feb. 2009), pp. 117–124. find
  6. Nikhil Jain and Yogish Sabharwal. “Y.: Optimal bucket algorithms for large mpi collectives on torus interconnects”. In: In: Proceedings of the 24th ACM International Conference on Supercomputing. 2010, pp. 27–36. find
  7. Grey Ballard et al. “Minimizing communication in numerical linear algebra”. In: SIAM Journal on Matrix Analysis and Applications 32.3 (2011), pp. 866–901. find
  8. Edgar Solomonik and James Demmel. “Communication-Optimal Parallel 2.5D Matrix Multiplication and LU Factorization Algorithms”. In: Euro-Par. Vol. 6853. 2011, pp. 90–109. ISBN : 978-3-642-23396-8. find
  9. Jeffrey Dean et al. “Large Scale Distributed Deep Networks”. In: Proceedings of the 25th International Conference on Neural Information Processing Systems - Volume 1. NIPS’12. Lake Tahoe, Nevada: Curran Associates Inc., 2012, pp. 1223–1231. URL : http://dl.acm. org/citation.cfm?id=2999134.2999271. ↗
  10. Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. “Imagenet classification with deep convolutional neural networks”. In: Advances in neural information processing systems. 2012, pp. 1097–1105. find
  11. Grey Ballard et al. “Communication optimal parallel multiplication of sparse random matrices”. In: Proceedings of the twenty-fifth annual ACM symposium on Parallelism in algorithms and architectures. ACM. 2013, pp. 222–231. find
  12. Michael Christ et al. Communication Lower Bounds and Optimal Algorithms for Programs That Reference Arrays - Part 1. Tech. rep. UCB/EECS-2013-61. EECS Department, University of California, Berkeley, May 2013. URL : http : / / www . eecs . berkeley . edu / Pubs / TechRpts/2013/EECS-2013-61.html. find
  13. James Demmel et al. “Communication-optimal parallel recursive rectangular matrix multiplication”. In: Parallel & Distributed Processing (IPDPS), 2013 IEEE 27th International Symposium on. IEEE. 2013, pp. 261–272. find
  14. Edgar Solomonik et al. “A massively parallel tensor contraction framework for coupled-cluster computations”. In: Journal of Parallel and Distributed Computing 74.12 (2014), pp. 3176– 3190. find
  15. Justus A. Calvin, Cannada A. Lewis, and Edward F. Valeev. “Scalable Task-based Algorithm for Multiplication of Block-rank-sparse Matrices”. In: Proceedings of the 5th Workshop on Irregular Applications: Architectures and Algorithms. IA3 ’15. Austin, Texas: ACM, 2015, 4:1–4:8. ISBN : 978-1-4503-4001-… — [doi:10.1145/2833179.2833186](https://doi.org/10.1145/2833179.2833186) find
  16. Martin Abadi et al. TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems. Software available from tensorflow.org. 2015. URL : https://www.tensorflow.org/. ↗
  17. Rafal Jozefowicz et al. “Exploring the Limits of Language Modeling”. In: CoRR abs/1602.02410 (2016). arXiv: 1602.02410. URL : http://arxiv.org/abs/1602.02410. ↗
  18. Penporn Koanantakool et al. “Communication-Avoiding Parallel Sparse-Dense Matrix-Matrix Multiplication”. In: 2016 IEEE International Parallel and Distributed Processing Symposium (IPDPS). May 2016, pp. 842–853. DOI : 10.1109/IPDPS.2016.117. — [doi:10.1109/ipdps.2016.117](https://doi.org/10.1109/ipdps.2016.117) find
  19. Amir Gholami et al. “Integrated Model and Data Parallelism in Training Neural Networks”. In: arXiv preprint arXiv:1712.04432 (2017). — [arXiv:1712.04432](https://arxiv.org/abs/1712.04432) find
  20. Noam Shazeer et al. “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”. In: CoRR abs/1701.06538 (2017). arXiv: 1701.06538. URL : http://arxiv. org/abs/1701.06538. find
  21. Ashish Vaswani et al. “Attention Is All You Need”. In: CoRR abs/1706.03762 (2017). arXiv: 1706.03762. URL : http://arxiv.org/abs/1706.03762. ↗
  22. James Demmel and Grace Dinh. “Communication-Optimal Convolutional Neural Nets”. In: arXiv preprint arXiv:1802.06905 (2018). — [arXiv:1802.06905](https://arxiv.org/abs/1802.06905) find
  23. Zhihao Jia, Matei Zaharia, and Alex Aiken. “Beyond Data and Model Parallelism for Deep Neural Networks”. In: arXiv preprint arXiv:1807.05358 (2018). find
  24. Zhihao Jia et al. “Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks”. In: arXiv preprint arXiv:1802.04924 (2018). find
  25. Peter Jin, Boris Ginsburg, and Kurt Keutzer. Spatially Parallel Convolutions. 2018. URL : https://openreview.net/forum?id=S1Yt0d1vG. ↗
  26. Penporn Koanantakool et al. “Communication-Avoiding Optimization Methods for Distributed Massive-Scale Sparse Inverse Covariance Estimation”. In: Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics. Ed. by Amos Storkey and Fernando Perez-Cruz. Vol. 84.… find