2020

GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

Lepikhin, Dmitry, Lee, HyoukJoong, Xu, Yuanzhong et al.

Understand

Neural network scaling has been critical for improving the model quality in many real-world machine learning applications with vast amounts of training data and compute.

Although this trend of scaling is affirmed to be a sure-fire approach for better model quality, there are challenges on the path such as the computation cost, ease of programming, and efficient implementation on parallel devices.
GShard is a module composed of a set of lightweight annotation APIs and an extension to the XLA compiler.
It provides an elegant way to express a wide range of parallel computation patterns with minimal changes to the existing model code.

Built on

Die grundlage der allgemeinen relativitätstheorie
Albert Einstein · 1923
Earlier work this paper cites.
A Cellular Computer to Implement the Kalman Filter Algorithm
Lynn Elliot Cannon · 1969
Earlier work this paper cites.
Transfer of training: A review and directions for future research
Timothy T Baldwin and J Kevin Ford · 1988
Earlier work this paper cites.
Special-purpose digital hardware for neural networks: An architectural survey
Paolo Ienne, Thierry Cornu, and Gary Kuhn · 1996
Earlier work this paper cites.
Bleu: a method for automatic evaluation of machine translation
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu · 2002
Earlier work this paper cites.
Scalable parallel programming with cuda
John Nickolls, Ian Buck, Michael Garland, and Kevin Skadron · 2008
Earlier work this paper cites.
MPI: A Message-Passing Interface Standard. Version 2.2, September 4th 2009
MPI Forum · 2009
Earlier work this paper cites.
Large-scale deep unsupervised learning using graphics processors
Rajat Raina, Anand Madhavan, and Andrew Y Ng · 2009
Earlier work this paper cites.
Rectified linear units improve restricted boltzmann machines
Vinod Nair and Geoffrey E. Hinton · 2010
Earlier work this paper cites.
Large scale parallel document mining for machine translation
Jakob Uszkoreit, Jay M. Ponte, Ashok C. Popat, and Moshe Dubiner · 2010
Earlier work this paper cites.
Deep, big, simple neural nets for handwritten digit recognition
Dan Claudiu Cireşan, Ueli Meier, Luca Maria Gambardella, and Jürgen Schmidhuber · 2010
Earlier work this paper cites.
ImageNet classification with deep convolutional neural networks
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton · 2012
Earlier work this paper cites.
Deep neural networks for acoustic modeling in speech recognition: The shared views of four research groups
Geoffrey Hinton, Li Deng, Dong Yu, George E Dahl, Abdel-rahman Mohamed, Navdeep Jaitly, Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Tara N Sainath, et al · 2012
Earlier work this paper cites.
Large scale distributed deep networks
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc’aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al · 2012
Earlier work this paper cites.
Theano: new features and speed improvements
Original
Frédéric Bastien, Pascal Lamblin, Razvan Pascanu, James Bergstra, Ian Goodfellow, Arnaud Bergeron, Nicolas Bouchard, David Warde-Farley, and Yoshua Bengio · 2012
Earlier work this paper cites.
Estimating or propagating gradients through stochastic neurons for conditional computation, 2013
Yoshua Bengio, Nicholas Léonard, and Aaron Courville · 2013
Earlier work this paper cites.
Low-rank approximations for conditional feedforward computation in deep neural networks, 2013
Andrew Davis and Itamar Arel · 2013
Earlier work this paper cites.
Sequence to sequence learning with neural networks
Ilya Sutskever, Oriol Vinyals, and Quoc V Le · 2014
Earlier work this paper cites.
Neural machine translation by jointly learning to align and translate
Original
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio · 2014
Earlier work this paper cites.
Conditional computation in neural networks for faster models, 2015
Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup · 2015
Earlier work this paper cites.
Going deeper with convolutions
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich · 2015
Earlier work this paper cites.
Deep residual learning for image recognition
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun · 2016
Earlier work this paper cites.
Identity mappings in deep residual networks
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun · 2016
Earlier work this paper cites.
Tensorflow: a system for large-scale machine learning
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al · 2016
Earlier work this paper cites.
Multi-way, multilingual neural machine translation with a shared attention mechanism
Orhan Firat, Kyunghyun Cho, and Yoshua Bengio · 2016
Earlier work this paper cites.
Wide and deep learning for recommender systems
Heng-Tze Cheng, Mustafa Ispir, Rohan Anil, Zakaria Haque, Lichan Hong, Vihan Jain, Xiaobing Liu, Hemal Shah, Levent Koc, Jeremiah Harmsen, and et al · 2016
Earlier work this paper cites.
Google’s neural machine translation system: Bridging the gap between human and machine translation
Original
Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, et al · 2016
Earlier work this paper cites.
Listen, attend and spell: A neural network for large vocabulary conversational speech recognition
William Chan, Navdeep Jaitly, Quoc Le, and Oriol Vinyals · 2016
Earlier work this paper cites.
Wavenet: A generative model for raw audio
Original
Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu · 2016
Earlier work this paper cites.
Attention is all you need, 2017
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin · 2017
Earlier work this paper cites.

Similar

Outrageously large neural networks: The sparsely-gated mixture-of-experts layer
Original
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean · 2017
Cited alongside, same era.
High-dimensional dynamics of generalization error in neural networks, 2017
Madhu S. Advani and Andrew M. Saxe · 2017
Cited alongside, same era.
Deep learning scaling is predictable, empirically, 2017
Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md. Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou · 2017
Cited alongside, same era.
Automatic differentiation in pytorch
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer · 2017
Cited alongside, same era.
Google’s multilingual neural machine translation system: Enabling zero-shot translation
Melvin Johnson, Mike Schuster, Quoc V. Le, Maxim Krikun, Yonghui Wu, Zhifeng Chen, Nikhil Thorat, Fernanda Viégas, Martin Wattenberg, Greg Corrado, and et al · 2017
Cited alongside, same era.
Understanding deep learning requires rethinking generalization
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals · 2017
Cited alongside, same era.
Exploring generalization in deep learning, 2017
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nathan Srebro · 2017
Cited alongside, same era.
In-datacenter performance analysis of a tensor processing unit
Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al · 2017
Cited alongside, same era.
Automatic differentiation in pytorch
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer · 2017
Cited alongside, same era.
On the optimization of deep networks: Implicit acceleration by overparameterization
Original
Sanjeev Arora, Nadav Cohen, and Elad Hazan · 2018
Cited alongside, same era.
The lottery ticket hypothesis: Finding sparse, trainable neural networks
Original
Jonathan Frankle and Michael Carbin · 2018
Cited alongside, same era.
Bert: Pre-training of deep bidirectional transformers for language understanding
Original
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova · 2018
Cited alongside, same era.

Then

Gpipe: Efficient training of giant neural networks using pipeline parallelism
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, and Zhifeng Chen · 2019
Later among the works it cites.
Beyond human-level accuracy
Joel Hestness, Newsha Ardalani, and Gregory Diamos · 2019
Later among the works it cites.
Lingvo: a modular and scalable framework for sequence-to-sequence modeling
Original
Jonathan Shen, Patrick Nguyen, Yonghui Wu, Zhifeng Chen, Mia Xu Chen, Ye Jia, Anjuli Kannan, Tara Sainath, Yuan Cao, Chung-Cheng Chiu, et al · 2019
Later among the works it cites.
BlueConnect: Decomposing All-Reduce for Deep Learning on Heterogeneous Network Hierarchy
Minsik Cho, Ulrich Finkler, and David Kung · 2019
Later among the works it cites.
Massively multilingual neural machine translation
Original
Roee Aharoni, Melvin Johnson, and Orhan Firat · 2019
Later among the works it cites.
Language modeling with deep transformers
Kazuki Irie, Albert Zeyer, Ralf Schlüter, and Hermann Ney · 2019
Later among the works it cites.
Learning deep transformer models for machine translation
Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F. Wong, and Lidia S. Chao · 2019
Later among the works it cites.
The evolved transformer, 2019
David R. So, Chen Liang, and Quoc V. Le · 2019
Later among the works it cites.
Megatron-lm: Training multi-billion parameter language models using gpu model parallelism
Original
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro · 2019
Later among the works it cites.
Summarizing cpu and gpu design trends with product data
Original
Yifan Sun, Nicolas Bohm Agostini, Shi Dong, and David Kaeli · 2019
Later among the works it cites.
Beyond Data and Model Parallelism for Deep Neural Networks
Zhihao Jia, Matei Zaharia, and Alex Aiken · 2019
Later among the works it cites.
Zero: Memory optimization towards training a trillion parameter models
Original
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He · 2019
Later among the works it cites.
Fast transformer decoding: One write-head is all you need
Original
Noam Shazeer · 2019
Later among the works it cites.
Scaling laws for neural language models
Original
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei · 2020
Closest in time.
Language models are few-shot learners
Original
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al · 2020
Closest in time.
Language models are few-shot learners
Original
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al · 2020
Closest in time.
Scaling description of generalization with number of parameters in deep learning
Mario Geiger, Arthur Jacot, Stefano Spigler, Franck Gabriel, Levent Sagun, Stéphane d’ Ascoli, Giulio Biroli, Clément Hongler, and Matthieu Wyart · 2020
Closest in time.
Depth-adaptive transformer
Original
Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli · 2020
Closest in time.
Controlling computation versus quality for neural sequence models, 2020
Ankur Bapna, Naveen Arivazhagan, and Orhan Firat · 2020
Closest in time.
https://www.tensorflow.org/xla , 2019
XLA: Optimizing Compiler for TensorFlow · 2020
Closest in time.
Train ML models on large images and 3D volumes with spatial partitioning on Cloud TPUs
Youlong Cheng, HyoukJoong Lee, and Tamas Berghammer · 2020
Closest in time.
https://github.com/onnx/onnx , 2019
ONNX: Open Neural Network Exchange · 2020
Closest in time.
https://ai.googleblog.com/2019/10/exploring-massively-multilingual.html
Exploring massively multilingual, massive neural machine translation · 2020
Closest in time.
https://ai.googleblog.com/2020/06/recent-advances-in-google-translate.html
Recent advances in google translate · 2020
Closest in time.
https://cloud.google.com/tpu/docs/bfloat16 , 2020
Using bfloat16 with TensorFlow models · 2020
Closest in time.
https://aiimpacts.org/2019-recent-trends-in-gpu-price-per-flops/
2019 recent trends in GPU price per FLOPS · 2020
Closest in time.
Automatic cross-replica sharding of weight update in data-parallel training, 2020
Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Hongjun Choi, Blake Hechtman, and Shibo Wang · 2020
Closest in time.
Surprisal-triggered conditional computation with neural networks, 2020
Loren Lugosch, Derek Nowrouzezahrai, and Brett H. Meyer · 2020
Closest in time.
Rethinking parameter counting in deep models: Effective dimensionality revisited, 2020
Wesley J. Maddox, Gregory Benton, and Andrew Gordon Wilson · 2020
Closest in time.

Beyond the bibliography

alphaXiv searches the wider corpus for related work and actual follow-ups.

Open on alphaXiv

alphaXiv is searching for related work…

Reading the bibliography…