12 — “How to Scale Your Model” (Google DeepMind / JAX Scaling Book)¶

The book: https://jax-ml.github.io/scaling-book/ Repo: https://github.com/jax-ml/scaling-book Authors: Jacob Austin, Sholto Douglas, Roy Frostig, Anselm Levskaya, Charlie Chen, Sharad Vikram, Federico Lebron, Peter Choy, Vinay Ramasesh, Albert Webson, Reiner Pope (Google DeepMind, 2025) Reading time: 12–20 focused hours.

1. Why this book — even if you never touch a TPU¶

Two great free books define the modern distributed-training curriculum:

Book	Framing	Hardware	Best at
HF Ultra-Scale Playbook	PyTorch + GPU + NCCL	H100/A100 clusters	5D-parallelism recipes, real-run reality
DeepMind Scaling Book	JAX + TPU + XLA	v4/v5e/v5p pods	First-principles roofline thinking for whole systems

The Scaling Book’s superpower is that it teaches you to reason about accelerator + interconnect together as one machine. Every chapter starts with a rooflines-and-bandwidth derivation, then does the arithmetic to predict what a specific config will do, then compares to what actually happens. That predict-then-measure discipline — the same one you saw at the kernel scale in Phase 2 — is here elevated to the fleet scale.

Ignore the fact that the substrate is TPU. Every idea (all-reduce bandwidth, activation memory, MFU, sharding decisions) maps cleanly to H100/B200 + NCCL. The TPU framing is actually pedagogically easier because pods are more homogeneous than GPU clusters and the arithmetic is less noisy.

The one book that would most improve your napkin-math skills as a fleet operator? This one.

2. Structure of the book¶

The book is ~11 chapters, each self-contained but building on the previous. Recommended order:

Ch	Title	What you extract
1	How to Think About TPUs	Roofline for a full chip; MXU vs VMEM vs HBM; TPU v5p numbers (2.5e14 bf16 FLOP/s per core, ~1 TB/s HBM, ICI at ~90 GB/s per link)
2	How to Think About Sharding	Named-axis (`GSPMD`) sharding, `PartitionSpec`, when a matmul needs comms
3	Rooflines for Transformers	Per-layer arithmetic intensity; when a config is HBM-bound vs ICI-bound vs FLOP-bound
4	Sharding a Transformer for Training	Which axes go where; TP vs data vs FSDP-analog under names
5	How to Handle a Transformer’s Activations	Recomputation, activation checkpointing, sequence-length blow-up
6	How to Train LLaMA-3 on TPU	Concrete walkthrough with real numbers
7	All About Rooflines for Inference	Prefill vs decode roofline separation, batching
8	How to Serve a Transformer on TPU	Sharded inference, how disaggregation ideas apply to TPU
9	How to Profile TPU Code	XLA/HLO reading, TensorBoard profiler, comm/compute overlap
10	Programming TPUs in JAX	`jit`, `shard_map`, `pjit`, mesh construction
11	JAX + XLA under the hood	Compiler intermediates, why HLO IR matters

Chapters 1–3 and 7 are the highest-leverage reads for someone whose day job is GPU inference — they are hardware-independent enough to translate line-by-line into H100 arithmetic.

3. Chapter-by-chapter cheat sheet¶

Ch 1 — How to Think About TPUs¶

Draw the diagram from memory:

    +-----+  +-----+
    | MXU |  | MXU |    ← matrix multiply units (systolic array, bf16/int8)
    +-----+  +-----+
       |        |
    +----------------+
    |     VMEM       |  ← ~128 MiB on-chip programmer-managed scratchpad
    +----------------+
       |
    +----------------+
    |     HBM        |  ← ~1 TB/s
    +----------------+
       |
    +--- ICI ---+       ← inter-chip interconnect (3D torus, ~90 GB/s per link)

GPU translation:

TPU concept	GPU analogue
MXU	tensor core (`wmma`/`wgmma`)
VMEM	shared memory + L1
HBM	HBM (same)
ICI (torus)	NVLink domain + IB fabric
Pod	multi-node H100 cluster

The ridge point of a v5p chip is 5e14 FLOPs / 2 * (1e12 B/s) ≈ 250 FLOPs/byte — remarkably close to H100 (~295 FLOPs/byte bf16). This is not a coincidence; it’s what makes memory-bandwidth-bound decode a fundamental fact of both platforms.

Ch 2 — How to Think About Sharding¶

Introduces GSPMD’s named-axis mesh — the abstraction PyTorch’s DTensor and DeviceMesh copied.

mesh = jax.sharding.Mesh(devices, ('data', 'model'))
P = jax.sharding.PartitionSpec

# X is [batch, seq, hidden]. Batch sharded over 'data', hidden over 'model'.
x_sharding = jax.sharding.NamedSharding(mesh, P('data', None, 'model'))

Key insight: a matmul between two arrays is comms-free iff its contracting axis is un-sharded on both sides. Otherwise you get an all-reduce (contracting axis sharded on both), all-gather (one side needs to see the whole thing), or reduce-scatter (a form of Megatron TP).

GPU translation: this is torch.distributed._tensor.DTensor and torch.distributed.device_mesh.DeviceMesh. Same math. When you use FSDP2 + TP in PyTorch, you are building the same object PJIT builds behind the scenes.

Ch 3 — Rooflines for Transformers¶

The book derives, term by term, the FLOPs and byte counts for every op in a transformer layer. Result: a per-op table showing which ops are compute-bound and which are memory-bound at a given batch/seq.

Take this table to your GPU work. Re-derive it for Llama-3-70B on H100 and compare. You should find:

QKV/O projections: compute-bound at batch × seq > ~4096; memory-bound below.
Attention (with FlashAttention): compute-bound at long seq, memory-bound at short.
FFN/SwiGLU: compute-bound at nontrivial batch × seq.
Layernorm/residual: always memory-bound (this is why fusion exists).

This chapter is the arithmetic scaffolding for every “why is my kernel slow” conversation you’ll ever have.

Ch 4 — Sharding a Transformer for Training¶

Walks through why TP shards on hidden, DP shards on batch, and PP shards on layer. Derives the communication cost of each and shows when each pays. Introduces fully sharded activations (Megatron-3 sequence parallelism) as a named-axis construction.

The neat trick: everything is stated as “which mesh axis does each tensor axis map onto?” Once you internalize that, TP + FSDP + PP + CP is just picking a mapping. No more religious wars about which is “the right” parallelism — it’s a linear-algebra choice.

Ch 5 — Handling Activations¶

Activation memory is the sneaky killer. This chapter shows:

Activation memory = O(layers × batch × seq × hidden × precision).
For Llama-3-70B at 8k seq, batch 2, bf16: ~40 GB per replica without checkpointing.
Selective activation recomputation (Korthikanti et al. 2022, arxiv:<phone_number_or_numberic_id_or_random_id_150>) drops this by ~5x by re-computing only the cheap ops.
Full activation checkpointing: drops to O(sqrt(layers)) at ~30% throughput cost.

Rule to memorize: if your training run OOMs and you’ve already turned on FSDP2, the next lever is activation checkpointing. If it still OOMs, drop batch. Only then consider PP or CP.

Ch 6 — Training LLaMA-3 on TPU¶

The concrete walkthrough. Values to steal:

MFU target: 40–55% is realistic for well-tuned dense transformer training. Above 60% is exceptional.
ICI bandwidth budget = per-step comms / per-step compute; if > ~20%, you are comm-bound.
Micro-batch size chosen so that per-chip compute > all-gather/reduce-scatter latency.

Reproduce on GPU: pick a 7B, train for 100 steps on 8 H100s with FSDP2, log MFU. You should land in the 35–50% range depending on seq length.

Ch 7 — Rooflines for Inference¶

This chapter is closest to your day job.

Prefill roofline: FLOPs ≈ 2 · P · S (P params, S tokens), bytes ≈ 2 · P (weights streamed once). Arithmetic intensity = S. Above the ridge (~250) → compute-bound.
Decode roofline: FLOPs ≈ 2 · P (one token), bytes ≈ 2 · P + KV. Intensity ≈ 1. Always memory-bound at batch 1. Batching lifts it: intensity = B, so B = 32 puts you well above the ridge.
Throughput/latency trade curve derived from these two.

The whole DistServe/Mooncake argument (Phase 6.09) is a direct consequence of this chapter’s math: two workloads with opposite ridge positions should not share hardware.

Ch 8 — Serving a Transformer on TPU¶

TPU-specific serving patterns, but the ideas transfer:

Prefill sharded aggressively (TP=8 or higher within a pod slice) for latency.
Decode with modest TP, larger batches for throughput.
Continuous batching maps onto TPU’s scan primitive.

Not as detailed as vLLM internals but useful for the shape of a well-designed serving system.

Ch 9 — Profiling TPU Code¶

TensorBoard profiler traces are the TPU analogue of Nsight Systems. The chapter’s method — find the longest stall, ask what it’s waiting on, fix, repeat — is the same skill you already have from GPU profiling.

Transferable habit: the book insists on a “step-time budget” per training step: expected FLOPs / peak FLOPs + expected comms / peak comms + overhead. Anything above budget is a bug or a design flaw. Apply this to any GPU training run you profile.

Ch 10 — Programming TPUs in JAX¶

jit, pjit, shard_map. Skim if you don’t intend to write JAX. The shard_map idea (explicit per-shard code, like SPMD) is worth understanding because DTensor’s model is similar.

Ch 11 — JAX + XLA Under the Hood¶

HLO IR, compilation stages, common failure modes (“XLA sharded my tensor wrong and now there’s a hidden all-reduce”). Analogous to reading torch.compile’s fx graphs. Skim.

4. Reading strategy¶

If you have 4 hours: Ch 1, Ch 3, Ch 7. This is the arithmetic backbone.

If you have 12 hours: add Ch 2, Ch 4, Ch 5, Ch 8. You now understand the full training + serving picture.

If you have 20 hours: read all 11 chapters. Do every worked example on paper. Reproduce two of them on a rented H100.

Do NOT: try to read it linearly and take notes in JAX. You’ll get stuck on syntax and miss the arithmetic. Read for the derivations first.

5. Exercises¶

Ridge-point derivation both stacks. For H100 SXM (989 TFLOPs bf16 / 3.35 TB/s HBM) and TPU v5p (5e14 FLOPs / 1e12 B/s HBM), compute the ridge in FLOPs/byte. Explain in one paragraph why they’re similar and what would have to change hardware-wise to make one wildly different.
Reproduce Ch 3 on Llama-3-70B, H100. Fill in the per-op FLOPs/bytes table for one transformer block at prefill S=2048 and decode B=1. Circle every op that’s memory-bound. This is a job-study-quality artifact.
Ch 4 mapping exercise. For an 8×H100 node running a 70B, write out the mesh-axis mapping for (a) TP=8, no FSDP; (b) TP=4, FSDP=2; (c) TP=2, FSDP=4. Predict which is fastest for a 4M-token training batch.
Ch 7 decode roofline. For Llama-3-70B fp8 on 1×H100 (80GB, 3.35 TB/s), predict batch-1 and batch-32 decode tokens/sec. Then rent an H100 and measure. Report your gap.

7. What this book is NOT¶

Not a serving-engine book. vLLM/SGLang internals are not in scope.
Not a quantization book. FP8 is mentioned but not derived.
Not a MoE playbook. EP appears but sparingly.
Not a rollout-infra book. RLHF/GRPO systems are outside scope.

For those, see Phase 6.09, 6.15, 6.16, and Phase 5.

8. Big takeaways to internalize¶

A distributed system has a roofline just like a kernel does. Comms-per-step vs comms-per-second is the same shape of graph as FLOPs vs FLOPs/sec.
Named-axis sharding is the abstraction that unifies every parallelism dimension. Once you see it, TP/PP/EP/DP/CP become tuple assignments, not architectures.
Predict before you measure, at every scale. The book’s discipline is the same as this whole roadmap’s: hypothesis → arithmetic → measurement → gap explanation.
The ridge-point argument is universal. It’s why disaggregated PD works, why decode batches, why memory bandwidth is the protagonist of inference — and it’s derived cleanly in Ch 3 and Ch 7 in a way you can steal for studies and design docs.

Next: 13_pretraining_small_model.md — putting the theory into a 124M–1B run on FineWeb-Edu.