# Scalable Training of Language Models using JAX pjit and TPUv4

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 8 of its 15 references that are in the map.
- First posted on arXiv 13 Apr 2022.

## Builds on (most cited first, 8 of 8)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Ray: A Distributed Framework for Emerging AI Applications](https://mlsanity.com/paper/1712.05889/) (2017)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [Benchmarking TPU, GPU, and CPU Platforms for Deep Learning](https://mlsanity.com/paper/1907.10701/) (2019)

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

- [Aya 23: Open Weight Releases to Further Multilingual Progress](https://mlsanity.com/paper/2405.15032/) (2024)
- [The Data Provenance Initiative: A Large Scale Audit of Dataset Licensing & Attribution in AI](https://mlsanity.com/paper/2310.16787/) (2023)
- [Intriguing Properties of Quantization at Scale](https://mlsanity.com/paper/2305.19268/) (2023)
- [To Code, or Not To Code? Exploring Impact of Code in Pre-training](https://mlsanity.com/paper/2408.10914/) (2024)
- [In Defense of Cross-Encoders for Zero-Shot Retrieval](https://mlsanity.com/paper/2212.06121/) (2022)
- [On the Limitations of Compute Thresholds as a Governance Strategy](https://mlsanity.com/paper/2407.05694/) (2024)
- [Tiny Aya: Bridging Scale and Multilingual Depth](https://mlsanity.com/paper/2603.11510/) (2026)
- [One Tokenizer To Rule Them All: Emergent Language Plasticity via Multilingual Tokenizers](https://mlsanity.com/paper/2506.10766/) (2025)
- [Soft-SVeRL: Self-Verified Reinforcement Learning with Soft Rewards](https://mlsanity.com/paper/2605.28561/) (2026)

## Abstract (arXiv)

Modern large language models require distributed training strategies due to their size. The challenges of efficiently and robustly training them are met with rapid developments on both software and hardware frontiers. In this technical report, we explore challenges and design decisions associated with developing a scalable training framework, and present a quantitative analysis of efficiency improvements coming from adopting new software and hardware solutions.