How To Scale Your Model
by Google DeepMind
The arithmetic behind why your training run is slow — rooflines, sharding and parallelism, from the people who scale Gemini.
Overview
Twelve chapters across four parts. Part 1, Preliminaries, covers roofline analysis, how to think about TPUs, and sharded matrices and how to multiply them. Part 2, Transformers, works through all the transformer math you need, how to parallelize a transformer for training, training LLaMA 3 on TPUs, all about transformer inference, and serving LLaMA 3 on TPUs. Part 3, Practical Tutorials, covers profiling TPU code with the JAX/TensorBoard profiler and reading what the XLA compiler emitted, plus programming TPUs in JAX. Part 4 closes with conclusions and further reading, and a bonus twelfth chapter, How to Think About GPUs, published 18 August 2025: H100, B200 and GB200 streaming-multiprocessor and memory hierarchies, NVLink node and cluster topologies, AllGather/AllReduce/AllToAll measured against their advertised numbers, roofline comparisons for data, tensor, expert and pipeline parallelism, and real sharding configurations from DeepSeek-V3 and LLaMA-3. Written by ten Google DeepMind researchers including Jacob Austin, Sholto Douglas, Roy Frostig, Anselm Levskaya and Reiner Pope, published 4 February 2025 under an MIT license with the source on GitHub. Chapters end with worked problems, and the JAX code runs on free Colab TPUs.
At a Glance
- Topic
- ML
- Level
- Advanced
- Format
- Book
- Cost
- Free
- Duration
- 12 chapters, ~12-18 hours to read carefully and attempt the worked problems
- Provider
- Google DeepMind
- Hands-on
- Yes — code/exercises
- Certificate
- None
What You’ll Learn
- ✓Do roofline analysis to predict whether a workload is compute or bandwidth bound
- ✓Read TPU and GPU hardware topology and reason about interconnect bandwidth limits
- ✓Express sharded matrix multiplication and predict the collectives each sharding implies
- ✓Derive FLOPs, memory and communication cost for a transformer forward and backward pass
- ✓Choose data, tensor, expert and pipeline parallelism from first principles rather than folklore
- ✓Size KV cache and batch to hit a latency or throughput target when serving
- ✓Profile JAX programs with the TensorBoard profiler and read what XLA actually emitted
- ✓Compare H100 and B200 NVLink collectives against their advertised numbers before trusting them
Highlights
- •Written by ten Google DeepMind researchers who scale production models, with worked LLaMA-3 training and serving configurations
- •Chapter 12, How to Think About GPUs (added 18 August 2025), makes the book useful even if you never touch a TPU
- •Every chapter ends with worked problems, and the Part 3 JAX code runs on free Colab TPUs
- •Honest about vendor numbers — it measures AllReduce at roughly 370GB/s against a claimed 450GB/s
- •Free, MIT-licensed, and the source lives on GitHub so corrections and new chapters are public
Who It’s For
Best For
- ✓Engineers whose training or serving runs are slower than the FLOPs say they should be
- ✓ML infrastructure and performance engineers sizing clusters and parallelism strategies
- ✓Researchers moving from single-GPU experiments to multi-node training
- ✓Anyone who wants to read a model card's hardware section and judge whether the numbers are plausible
Prerequisites
- •Solid understanding of the transformer architecture and attention
- •Comfort with linear algebra and back-of-the-envelope arithmetic — this book is mostly math
- •Some JAX or NumPy exposure for Part 3; the GPU chapter is an appendix, not a starting point
FAQ
What is How To Scale Your Model?
A free, blog-style textbook from Google DeepMind on how large language models actually run on TPU and GPU hardware, and how to parallelize them during training and inference so they stay efficient at scale. Written for engineers who can already build a transformer but cannot yet predict how fast it will run. After reading it you can size parallelism strategies, KV caches and clusters from first principles.
Is How To Scale Your Model free?
How To Scale Your Model is free to access.
What level is How To Scale Your Model for?
How To Scale Your Model is aimed at a advanced audience. Recommended background: Solid understanding of the transformer architecture and attention, Comfort with linear algebra and back-of-the-envelope arithmetic — this book is mostly math, Some JAX or NumPy exposure for Part 3; the GPU chapter is an appendix, not a starting point.
How long does How To Scale Your Model take?
Expect roughly 12 chapters, ~12-18 hours to read carefully and attempt the worked problems. Most learners work through it at their own pace.
What will I learn from How To Scale Your Model?
You'll learn: Do roofline analysis to predict whether a workload is compute or bandwidth bound; Read TPU and GPU hardware topology and reason about interconnect bandwidth limits; Express sharded matrix multiplication and predict the collectives each sharding implies; Derive FLOPs, memory and communication cost for a transformer forward and backward pass; Choose data, tensor, expert and pipeline parallelism from first principles rather than folklore; Size KV cache and batch to hit a latency or throughput target when serving; Profile JAX programs with the TensorBoard profiler and read what XLA actually emitted; Compare H100 and B200 NVLink collectives against their advertised numbers before trusting them.
Topics
Sources
This page was written from 4 sources, 2 on domains other than jax-ml.github.io.