# Exploring the limits of Concurrency in ML Training on Google TPUs

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 7 Nov 2020.

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

- [Understanding Data Storage and Ingestion for Large-Scale Deep Recommendation Model Training](https://mlsanity.com/paper/2108.09373/) (2021)
- [RecShard: Statistical Feature-Based Memory Optimization for Industry-Scale Neural Recommendation](https://mlsanity.com/paper/2201.10095/) (2022)
- [OneFlow: Redesign the Distributed Deep Learning Framework from Scratch](https://mlsanity.com/paper/2110.15032/) (2021)
- [Concurrent Adversarial Learning for Large-Batch Training](https://mlsanity.com/paper/2106.00221/) (2021)
- [Demystifying BERT: Implications for Accelerator Design](https://mlsanity.com/paper/2104.08335/) (2021)
- [Efficient Strong Scaling Through Burst Parallel Training](https://mlsanity.com/paper/2112.10065/) (2021)
- [Accelerating Large Batch Training via Gradient Signal to Noise Ratio (GSNR)](https://mlsanity.com/paper/2309.13681/) (2023)
- [Learning What to Predict: Downstream-Guided Task Design for Continued Pretraining](https://mlsanity.com/paper/2601.22108/) (2026)
- [An Efficient 2D Method for Training Super-Large Deep Learning Models](https://mlsanity.com/paper/2104.05343/) (2021)

## Abstract (arXiv)

Recent results in language understanding using neural networks have required training hardware of unprecedentedscale, with thousands of chips cooperating on a single training run. This paper presents techniques to scaleML models on the Google TPU Multipod, a mesh with 4096 TPU-v3 chips. We discuss model parallelism toovercome scaling limitations from the fixed batch size in data parallelism, communication/collective optimizations,distributed evaluation of training metrics, and host input processing scaling optimizations. These techniques aredemonstrated in both the TensorFlow and JAX programming frameworks. We also present performance resultsfrom the recent Google submission to the MLPerf-v0.7 benchmark contest, achieving record training times from16 to 28 seconds in four MLPerf models on the Google TPU-v3 Multipod machine.