Scaling LLMs, Distributed Systems, and Inference

Study notes and sources that build a first principles understanding of how large language models are trained and served across many accelerators.

The spine is How to Scale Your Model, the JAX and Google DeepMind "scaling book," worked through end to end. Each of its 13 parts has one set of notes here, written as a standalone explainer and cross checked against supplementary material (papers, blog posts, lectures). All 13 are complete.

---

What you can do after reading this

• Estimate the cost of training a model (FLOPs, wall clock, dollars, memory) from its shape, before touching a cluster.
• Choose a parallelism strategy (data, tensor, pipeline, expert, sequence, or a combination) for a given model, batch size, and pod, and say why.
• Explain why inference decode is slow and what actually fixes it (batching, KV cache reduction, quantization, speculative decoding).
• Estimate the latency, throughput, and cost per token of serving a model, and where it sits on the latency versus throughput curve.
• Read a profiler trace, compute the roofline for each op, and find where time is lost.
• Write sharded JAX: meshes, NamedSharding, shard_map, explicit collectives.
• Reason about the GPU version of all of the above and where it diverges from TPUs.

---

The core idea

An accelerator is three numbers: how fast it does arithmetic (FLOPs per second), how fast it moves data (bytes per second, for on chip memory and for the network between chips), and how much it holds (bytes).

Any operation takes at least

max( FLOPs / peak_FLOPs ,  bytes_moved / bandwidth )

The two terms can overlap, so the runtime is between that maximum and their sum. An operation is compute bound when its arithmetic intensity (FLOPs per byte moved) exceeds the hardware's ridge point (peak_FLOPs / bandwidth), which is roughly 200 to 300 on current chips. Otherwise it is memory bound or communication bound and the expensive compute units sit idle.

Almost every FLOP in a Transformer lives in a matmul, and the arithmetic intensity of a matmul is roughly its batch size in tokens. So the entire craft of scaling is: keep the token batch per chip large enough, and the communication small enough, that every operation stays above the ridge point as you spread the model over more and more chips. Each part of the book is one application of that single rule, to a different operation and a different link in the memory hierarchy.

---

The map

Foundations

Part	Notes	Source	What it gives you
0. Introduction	notes	index	the Transformer forward pass as a fixed graph whose cost we are going to predict
1. Rooflines	notes	roofline	the compute, memory, communication model and the ridge point
2. TPUs	notes	tpus	the hardware as a hierarchy: MXU, VMEM at 22x HBM, the ICI torus, DCN
3. Sharded Matmuls	notes	sharding	when a sharded matmul needs a collective, which one, and its cost

The Transformer as a cost model

Part	Notes	Source	What it gives you
4. Transformer Math	notes	transformers	parameters, training FLOPs (the 6ND rule), attention cost, KV cache size

Training at scale

Part	Notes	Source	What it gives you
5. Training	notes	training	data parallel, FSDP, tensor, pipeline, and how to combine them
6. Training LLaMA 3	notes	applied-training	the full arithmetic for a real 70B pretraining run

Inference and serving

Part	Notes	Source	What it gives you
7. Inference	notes	inference	prefill versus decode, why decode is memory bound, the KV cache as the constraint
8. Serving LLaMA 3 70B	notes	applied-inference	the full arithmetic for serving: slice size, latency, throughput, cost

Practice

Part	Notes	Source	What it gives you
9. Profiling	notes	profiling	roofline every op, compare to the trace, a gap is a bug
10. Programming in JAX	notes	jax-stuff	meshes, the three sharding paradigms, shard_map and collectives

Going further

Part	Notes	Source	What it gives you
11. Conclusions	notes	conclusion	the whole book on one page, plus an annotated reading list
12. GPUs	notes	gpus	the same rooflines on a fat tree of all to all islands, and how the recipe inverts

---

How to read this

• Fast path (the load bearing four): 1 Rooflines, 4 Transformer Math, 5 Training, 7 Inference. With these you can do most of the useful estimation.
• Full path: 0 through 12 in order. Each part assumes the ones before it.
• Inference only: 4, 7, 8.
• Just the JAX: 3 Sharding, 10 Programming in JAX.

---

Results worth memorizing

Result	Statement
Roofline	runtime is at least max(FLOPs / peak, bytes / bandwidth)
Ridge point	peak / bandwidth, about 240 FLOPs per byte on TPU v5e, about 295 on H100
Matmul intensity	roughly the token batch size, so a bf16 matmul needs batch above about 240 to be compute bound
Parameter count	about 12 * L * D², three quarters of it in the MLP blocks
Training FLOPs	about 6 * parameters * tokens, while sequence length T < 8D
Attention core FLOPs	about 12 * B * T² * N * H * L, dominant once T > 8D (about 64k tokens for a large model)
KV cache	about 2 * L * K * H * S bytes per sequence, divided by the grouped query attention group size
Training memory	about 16 bytes per parameter with Adam in mixed precision
Collective cost	AllGather and ReduceScatter cost bytes / bandwidth, AllReduce twice that, independent of shard count
Data parallel and FSDP	stay compute bound while tokens per chip exceed about 2550 (TPU v5p)
Tensor parallelism	saturates at 8 to 16 way
FSDP plus tensor parallelism	reaches a minimum of about 100 tokens per chip
Decode critical batch	about 240 concurrent tokens (bf16 v5e), about 120 with int8 weights
Decode latency	about (parameters + batch * KV_cache) / HBM_bandwidth, memory bandwidth bound
LLaMA 3 70B pretraining	70.4B parameters, 6.3e24 FLOPs over 15T tokens, about 44 days on one TPU v5p pod at 40% MFU

---

Core notation

Symbol	Meaning
D	model width (d_model), the residual stream dimension
F	MLP inner width (d_ff), usually about 4D
L	number of layers
N, K	query heads, key and value heads (K < N for grouped query attention)
H	per head width (d_head), usually N * H = D
V	vocabulary size
B, T, S	batch (sequences), query sequence length, key sequence length
HBM	high bandwidth memory, where the weights and KV cache live
ICI	inter chip interconnect, the TPU torus links
MFU	model FLOPs utilization, achieved useful FLOPs over peak FLOPs
prefill	the parallel forward pass over the prompt
decode	one token per step generation, the memory bound phase
ridge point	the hardware's FLOPs per byte, the compute versus memory boundary

---

Supplementary sources

Read alongside the parts they support. The Part 11 notes carry a fuller annotated list.

Rooflines and hardware

• Making Deep Learning Go Brrrr From First Principles, Horace He: compute, memory, and overhead bound, operator fusion.
• TPU Deep Dive, Henry Ko: the clearest companion to Part 2.

Transformer and inference math

• Transformer Inference Arithmetic, kipply: KV cache sizing and the memory bound of decode.
• Efficiently Scaling Transformer Inference, Pope et al.: the paper behind Parts 7 and 8.
• Decoder Only LLM Inference, Ted Kyi: the forward pass shape by shape.

Parallelism and training

• The Ultra Scale Playbook, Hugging Face: every parallelism form on GPU clusters, with measured numbers.
• Megatron LM: tensor parallelism for Transformers.
• ZeRO: sharded optimizer, gradient, and parameter states.
• GPipe and PipeDream: pipeline parallelism.

Attention and serving systems

• FlashAttention: memory aware attention.
• Orca and vLLM / PagedAttention: continuous batching and paged KV cache.

JAX

• Distributed arrays and automatic parallelization.
• shard_map.

Foundations

• Attention Is All You Need: the architecture.
• The Illustrated Transformer: the data flow.

---

Notes format

Every part N/notes_N.md is a standalone explainer, structured roughly as:

1. Purpose and how the part connects to the ones before it.
2. The mechanism, with formulas and at least one worked numeric example.
3. Extensions: what each supplementary source adds beyond the chapter.
4. Results carried forward: the formulas and numbers later parts depend on.
5. Open questions: further readings I need to do.