Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training

Abstract

In data-parallel synchronous training of deep neural networks, different devices (replicas) run the same program with different partitions of the training batch, but weight update computation is repeated on all replicas, because the weights do not have a batch dimension to partition. This can be a bottleneck for performance and scalability in typical language models with large weights, and models with small per-replica batch size which is typical in large-scale training. This paper presents an approach to automatically shard the weight update computation across replicas with efficient communication primitives and data formatting, using static analysis and transformations on the training computation graph. We show this technique achieves substantial speedups on typical image and language models on Cloud TPUs, requiring no change to model code. This technique helps close the gap between traditionally expensive (ADAM) and cheap (SGD) optimizers, as they will only take a small part of training step time and have similar peak memory usage. It helped us to achieve state-of-the-art training performance in Google's MLPerf 0.6 submission.

12 pages, 23 figures, 1 table

References in corpus (5)

• Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation
• One weird trick for parallelizing convolutional neural networks
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
• Adafactor: Adaptive Learning Rates with Sublinear Memory Cost
• Scale MLPerf-0.6 models on Google TPU-v3 Pods

Cited by in corpus (9)

• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• BigSSL: Exploring the Frontier of Large-Scale Semi-Supervised Learning for Automatic Speech Recognition
• Scalable Second Order Optimization for Deep Learning
• Improving Automatic Parallel Training via Balanced Memory Workload Optimization
• PaSE: Parallelization Strategies for Efficient DNN Training
• An Oracle for Guiding Large-Scale Model/Hybrid Parallel Training of Convolutional Neural Networks
• Exploring the limits of Concurrency in ML Training on Google TPUs
• PartIR: Composing SPMD Partitioning Strategies for Machine Learning
• Highly Available Data Parallel ML training on Mesh Networks