mlsanityData to 14 Oct 2026

Scalable Training of Language Models using JAX pjit and TPUv4

arXiv 2204.06514 ↗First posted 13 Apr 2022cs.LG cs.CLcited by 9

Not on the map. Why?

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.

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

8 of 15 references.

References · 8

  1. 2017Attention Is All You Needcited by 60,361
  2. 2014Adam: A Method for Stochastic Optimizationcited by 37,240
  3. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,599
  4. 2020Bootstrap your own latent: A new approach to self-supervised Learningcited by 4,750
  5. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,509
  6. 2017Ray: A Distributed Framework for Emerging AI Applicationscited by 527
  7. 2021GSPMD: General and Scalable Parallelization for ML Computation Graphscited by 111
  8. 2019Benchmarking TPU, GPU, and CPU Platforms for Deep Learningcited by 52

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2022: 1; 2023: 2; 2024: 3; 2025: 1; 2026: 2.

9 papers.

  1. 2024Aya 23: Open Weight Releases to Further Multilingual Progresscited by 123
  2. 2023The Data Provenance Initiative: A Large Scale Audit of Dataset Licensing & Attribution in AIcited by 73
  3. 2023Intriguing Properties of Quantization at Scalecited by 49
  4. 2024To Code, or Not To Code? Exploring Impact of Code in Pre-trainingcited by 47
  5. 2022In Defense of Cross-Encoders for Zero-Shot Retrievalcited by 21
  6. 2024On the Limitations of Compute Thresholds as a Governance Strategycited by 18
  7. 2026Tiny Aya: Bridging Scale and Multilingual Depthcited by 15
  8. 2025One Tokenizer To Rule Them All: Emergent Language Plasticity via Multilingual Tokenizerscited by 12
  9. 2026Soft-SVeRL: Self-Verified Reinforcement Learning with Soft Rewardscited by 1