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

By Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Hongjun Choi, Blake A. Hechtman, Shibo Wang. 2020

52 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture, Efficiency & Compression, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Efficiency & Compression

In Optimization & Training

Citations in the corpus

Efficient Large Scale Language Modeling with Mixtures of Experts (2021, background): “Recent work has explored sharding model parameters, optimizer state and gradients across workers (Xu et al., 2020; Rajbhandari et al., 2020), enabling training of models with more than one trillion parameters using only data parallelism without the added complexity introduced by model parallel…”
GSPMD: General and Scalable Parallelization for ML Computation Graphs (2021, methodology): “This sharding configuration combines the benefit of data parallelism, weight-update sharding [28, 38] and inlayer model parallelism [35], and only requires 7 annotations in the model.”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, background): “6 [28], sharded data parallelism [24, 44], where optimizer state is sharded over data-parallel workers, was introduced.”
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (2020, background): “Weight-update sharding (Xu et al., 2020) is another automatic parallelization transformation based on XLA, which mostly focuses on performance optimizations for TPU clusters, and conceptually can be viewed as a special case for GShard.”
PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, methodology): “Zero-Redundancy parallelism [26, 27, 29] shards parameters as well but communicates parameters on-demand to recover the unsharded form and executes the model as if it were replicated on every device.”
This paper is included in the Proceedings of the 16th USENIX Symposium on Operating Systems Design and Implementation. (unknown): “Besides operators in the forward or backward passes, one can also partition the operators from the weight update phase, yielding the weight update shard-ing or equivalently the ZeRO [44,56] technique, commonly comprehended as an optimization of data parallelism.”