Ch 11: Conclusions — The Scaling Book

Chapter 0: Key Takeaways

Let us crystallize the core lessons from the entire book. If you remember nothing else, remember these:

The Three Challenges

Challenge	Nature	Primary Tools
Memory	Hard limit — if it does not fit, training cannot proceed	Sharding (ZeRO, TP, PP), activation recomputation, mixed precision
Compute efficiency	Optimization target — minimize idle hardware	Large batch sizes, fusions, collective matmuls, pipeline schedules
Communication	Scaling bottleneck — GPUs/TPUs waiting on each other	Overlapping comms with compute, topology-aware sharding, quantized gradients

The Roofline Mindset

The single most important skill from this book is roofline reasoning: for any operation, calculate:

T = max(T_compute, T_{communication}) = max(FLOPs / peak_FLOPS, bytes / bandwidth)

If your measured time matches this, there is nothing left to optimize. If it is far off, something is wrong and the profiler will tell you what.

The 5D Parallelism Framework

Dimension	Splits	Best For	Communication Cost
Data (DP)	Batch	Scaling throughput	AllReduce gradients
Tensor (TP)	Hidden dim	Fitting large layers (intra-node)	AllReduce / AllGather per layer
Context (CP)	Sequence	Long sequences	Ring Attention for K/V
Pipeline (PP)	Layers	Multi-node models	Activations between stages
Expert (EP)	MoE experts	Scaling capacity cheaply	AllToAll token routing

The fundamental trade-off: You can often trade one of {computation, communication, memory} for another. Activation recomputation trades extra compute for less memory. Tensor parallelism trades more communication for distributed memory. Finding the right balance is the art of distributed training.

Check: Which of these is a hard limit (not just an optimization target)?

Memory — if it does not fit, training cannot proceed at all Compute efficiency — low utilization means training cannot converge Communication overhead — any latency above zero blocks training

Chapter 1: Acknowledgments

This book represents a significant collective investment from many people at Google DeepMind. The authors wish to acknowledge key contributors:

Contributor	Role
James Bradbury, Reiner Pope, Blake Hechtman	Originally derived many core ideas; early pioneers of the systems view of the Transformer
Sholto Douglas	Wrote the first version; responsible for the overall narrative and kicking off the project
Jacob Austin	Led the transformation from rough notes to polished artifact; editing, formatting, release coordination
Anselm Levskaya, Charlie Chen	Created most figures and animations
Charlie Chen	Wrote the inference section and drew inference figures
Roy Frostig	Publication, editing, and many other steps

Additional reviewers who provided critical feedback include Zak Stone, Nikhil Sethi, Caitlin Stanton, Alek Dimitriev, Sridhar Lakshmanamurthy, Albert Magyar, Diwakar Gupta, Jeff Dean, Corry Wang, Matt Johnson, Peter Hawkins, and many others. Ruiqi Gao helped with HTML formatting.

Open-source roots: This book grew out of internal Google DeepMind documentation about TPU performance engineering. The authors made a deliberate choice to open-source it, believing that the knowledge of how to train models at scale should not remain behind closed doors.

Check: Who originally derived many of the core ideas about the systems view of the Transformer that this book is built on?

James Bradbury, Reiner Pope, and Blake Hechtman The authors of the original Transformer paper (Vaswani et al.) The JAX development team at Google Brain

Chapter 2: Further Reading

The authors curate a reading list organized by area. Each resource offers a different perspective or deeper dive on topics covered in this book.

Hardware Deep Dives

Resource	Focus
TPU Deep Dive	In-depth TPU architecture in the spirit of this book
Domain Specific Architectures for AI Inference	Hardware and model deep dive for inference workloads
A Domain-Specific Supercomputer for Training DNNs	One of the original TPU papers with details about the Google TPU program

Performance Engineering

Resource	Focus
Making Deep Learning Go Brrrr From First Principles	GPU/PyTorch-focused rooflines and performance engineering
How to Optimize a CUDA Matmul Kernel	Step-by-step CUDA kernel optimization worklog; excellent for GPU vs TPU contrast
Rafi Witten's High Performance LLMs 2024	Stanford course on TPU performance engineering with slides on GitHub

JAX Programming

Resource	Focus
Writing TPU Kernels with Pallas	Custom TPU kernels, lower-level details not covered here
Distributed Arrays and Automatic Parallelism	Guide to parallelism APIs in JAX; good for implementing ideas from this book

Broader Resources

Resource	Focus
Efficiently Scaling Transformer Inference (arXiv 2211.05102)	Detailed mathematics of Transformer inference; major inspiration for this book
HuggingFace Ultra-Scale Playbook	GPU analog covering PyTorch parallelism and memory-saving techniques
Transformer Inference Arithmetic	Blog with many of the same ideas and excellent illustrations
Stanford CS336	Fantastic Stanford course on LLM training/serving with exercises (Assignments 1 & 2 especially relevant)
Stas Bekman's ML Engineering Handbook	Practical guide: negotiating with cloud providers, cluster management, empirical throughput measurements

The authors note: "There remains a lot of room for comprehensive writing in this area, so we hope this manuscript encourages more of it! We also believe that this is a fruitful area to study and research. In many cases, it can be done even without having many hardware accelerators on hand."

Check: Which resource is described as the "GPU analog" to this (TPU-focused) scaling book?

Stanford CS336 The HuggingFace Ultra-Scale Playbook Making Deep Learning Go Brrrr From First Principles

Chapter 3: Research Frontiers

The field of distributed training is rapidly evolving. Here are some areas where active research is pushing the boundaries:

Hardware Evolution

Blackwell GPUs (B200, GB200): 2x FLOPs over Hopper, NVL72 domains with 72 GPUs in a single NVLink domain, new Tensor Memory (TMEM) for feeding larger Tensor Cores
TPU v6e and beyond: Continuing to increase MXU size and interconnect bandwidth
Custom silicon: Companies like MatX, Groq, Cerebras designing domain-specific accelerators that may change the roofline calculus entirely

Algorithmic Frontiers

Zero-bubble pipeline parallelism: Schedules that eliminate pipeline bubbles entirely through careful interleaving of forward/backward micro-batches
Sequence parallelism at scale: Ring Attention and variants enabling million-token context windows
Mixed-precision training: FP4 training on Blackwell, with open questions about convergence and accuracy
Communication compression: Quantized gradients, error feedback, and 1-bit Adam for reducing bandwidth requirements

Systems Frontiers

Automatic parallelism discovery: Can compilers find optimal 5D configurations automatically?
Heterogeneous training: Mixing GPU and TPU types, or using CPUs for offloading
Fault tolerance at scale: When training on 10,000+ accelerators, hardware failures are routine. Checkpointing and recovery strategies become critical.
Energy-efficient training: Understanding and reducing the carbon footprint of large-scale training runs

Inference and Serving Frontiers

Speculative decoding: Use a small draft model to propose tokens, verified in batch by the large model. Can 2-3x decode throughput with careful draft model selection.
KV cache compression: Quantizing or pruning the KV cache to serve longer contexts without proportionally more memory. Techniques like GQA (Grouped Query Attention) and MQA (Multi-Query Attention) are already standard.
Disaggregated serving: Separating prefill (compute-bound) from decode (memory-bound) onto different hardware, optimizing each independently.
Continuous batching: Dynamically adding and removing requests from a running batch, maximizing GPU utilization during decode.

Emerging Hardware Paradigms

Approach	Key Idea	Status
Wafer-scale (Cerebras)	Entire wafer as one chip; massive on-chip SRAM eliminates HBM bottleneck	Production, specialized
LPU (Groq)	Deterministic, SRAM-only architecture; eliminates memory hierarchy complexity	Production inference
MatX	Custom silicon optimized for dense linear algebra with focus on simplicity	In development
Photonic computing	Use light for matrix multiplication at near-zero energy	Research stage

The key open question: As models get larger and clusters get bigger, will the current 5D parallelism framework suffice, or will we need fundamentally new approaches? The interplay between hardware design, compiler technology, and training algorithms continues to be one of the most exciting areas in ML systems research.

Check: What new memory type does the Blackwell GPU architecture (B200) introduce to feed its larger Tensor Cores?

Extended SMEM (eSMEM) Tensor Memory (TMEM) Vector Memory (VMEM)

Chapter 4: Final Thoughts

Thank you for reading this book. Let us close with a decision tree that captures the practical wisdom of distributed training:

Your Situation	Primary Bottleneck	Start With
1 GPU, model fits	Throughput	Gradient accumulation, mixed precision, activation recomputation
1 node (8 GPUs/TPUs), model fits	Throughput	Data parallelism (DP=8)
1 node, model too big	Memory	Tensor parallelism or ZeRO-3
Multiple nodes	Communication	PP across nodes + TP within node
Very long sequences	Activation memory	Context parallelism (Ring Attention)
MoE architecture	Expert capacity	Expert parallelism (AllToAll routing)

The authors ran over 4,100 experiments with up to 512 GPUs. The benchmarks consistently show that no single technique dominates — the optimal configuration depends on model size, batch size, sequence length, and hardware topology. There is no substitute for understanding the fundamentals and reasoning from first principles.

The book's code implementations are available in two repositories:

Picotron: Self-contained, educational implementations of each technique in single files
Nanotron: Production-ready distributed training code

For questions or feedback, reach the corresponding author Jacob Austin at jacobaustin123 [at] gmail [dot] com, or contribute via GitHub issues and pull requests.

Citation:
Austin et al., "How to Scale Your Model", Google DeepMind, online, 2025.

Check: If your model does not fit on a single node (e.g., 8 GPUs), what is the recommended parallelism strategy?

Use data parallelism with a very large batch size Pipeline parallelism across nodes combined with tensor parallelism within each node Context parallelism to split the model across sequence chunks

Conclusions & Further Reading