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

arXiv: https://arxiv.org/abs/2004.13336 · HTML: https://mlsanity.com/paper/2004.13336/

## Key facts

- Cited by 30 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 28 Apr 2020.

## Built on by (most cited first, 30 of 30)

- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [CoCa: Contrastive Captioners are Image-Text Foundation Models](https://mlsanity.com/paper/2205.01917/) (2022)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Efficient Large Scale Language Modeling with Mixtures of Experts](https://mlsanity.com/paper/2112.10684/) (2021)
- [Mobile Edge Intelligence for Large Language Models: A Contemporary Survey](https://mlsanity.com/paper/2407.18921/) (2024)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [BigSSL: Exploring the Frontier of Large-Scale Semi-Supervised Learning for Automatic Speech Recognition](https://mlsanity.com/paper/2109.13226/) (2021)
- [Vision Models Are More Robust And Fair When Pretrained On Uncurated Images Without Supervision](https://mlsanity.com/paper/2202.08360/) (2022)
- [FourCastNet 3: A geometric approach to probabilistic machine-learning weather forecasting at scale](https://mlsanity.com/paper/2507.12144/) (2025)
- [Galvatron: Efficient Transformer Training over Multiple GPUs Using Automatic Parallelism](https://mlsanity.com/paper/2211.13878/) (2022)
- [On the Factory Floor: ML Engineering for Industrial-Scale Ads Recommendation Models](https://mlsanity.com/paper/2209.05310/) (2022)
- [On the Generalization Mystery in Deep Learning](https://mlsanity.com/paper/2203.10036/) (2022)
- [Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning](https://mlsanity.com/paper/2201.12023/) (2022)
- [CommonCanvas: An Open Diffusion Model Trained with Creative-Commons Images](https://mlsanity.com/paper/2310.16825/) (2023)
- [Improving Automatic Parallel Training via Balanced Memory Workload Optimization](https://mlsanity.com/paper/2307.02031/) (2023)
- [Outlier-Safe Pre-Training for Robust 4-Bit Quantization of Large Language Models](https://mlsanity.com/paper/2506.19697/) (2025)
- [Training and Serving System of Foundation Models: A Comprehensive Survey](https://mlsanity.com/paper/2401.02643/) (2024)
- [PartIR: Composing SPMD Partitioning Strategies for Machine Learning](https://mlsanity.com/paper/2401.11202/) (2024)
- [Hardware Phi-1.5B: A Large Language Model Encodes Hardware Domain Specific Knowledge](https://mlsanity.com/paper/2402.01728/) (2024)
- [Hardware Scaling Trends and Diminishing Returns in Large-Scale Distributed Training](https://mlsanity.com/paper/2411.13055/) (2024)
- [End-to-end Adaptive Distributed Training on PaddlePaddle](https://mlsanity.com/paper/2112.02752/) (2021)
- [CowClip: Reducing CTR Prediction Model Training Time from 12 hours to 10 minutes on 1 GPU](https://mlsanity.com/paper/2204.06240/) (2022)
- [An Oracle for Guiding Large-Scale Model/Hybrid Parallel Training of Convolutional Neural Networks](https://mlsanity.com/paper/2104.09075/) (2021)
- [Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR](https://mlsanity.com/paper/2210.06352/) (2022)
- [Two-dimensional Sparse Parallelism for Large Scale Deep Learning Recommendation Model Training](https://mlsanity.com/paper/2508.03854/) (2025)
- [Partial Parameter Updates for Efficient Distributed Training](https://mlsanity.com/paper/2509.22418/) (2025)
- [PaSE: Parallelization Strategies for Efficient DNN Training](https://mlsanity.com/paper/2407.04001/) (2024)
- [Secure Multiparty Generative AI](https://mlsanity.com/paper/2409.19120/) (2024)
- [Raising Bars, Not Parameters: LilMoo Compact Language Model for Hindi](https://mlsanity.com/paper/2603.03508/) (2026)

## Abstract (arXiv)

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.