13  Computational Performance

The book so far has made two kinds of choice. We chose models — which architecture, how deep, what kind of attention — and we chose data — how much, how batched, how tokenized. This chapter is about the third choice, the one every training run makes whether or not you attend to it: how to use the machine. Two models with identical loss curves can differ by an order of magnitude in wall-clock time and in whether they fit in memory at all — and holding the model and its loss curve fixed, that difference is systems, not mathematics. It is whether the code kept the arithmetic units fed, moved as few bytes as it could, stayed out of the interpreter’s way, and spent its memory where it counted.

There is one map to start from, and one method. The map is the roofline: a computation is limited either by how fast the machine can do arithmetic or by how fast it can move bytes, and a single ratio — arithmetic intensity, the FLOPs performed per byte moved — tells you which, before you run anything. The method is a loop: measure what the program actually does, classify which of three regimes it is in (compute-bound, bandwidth-bound, or overhead-bound), fix the binding constraint with the technique that targets it, and re-measure, because every fix moves the bottleneck. Section 13.1 builds both the map and the method on the book’s own build machine; every section after it is the loop applied to one regime.

The seven sections form a ladder from a single operation to a full model on many GPUs. Section 13.1 establishes the roofline, the three regimes, and the measurement discipline the frameworks’ asynchronous dispatch makes non-optional. Section 13.2 explains where the roofline’s two numbers come from — the memory hierarchy, the tensor cores and their format ladder, the interconnects, and the energy budget underneath them all — using our own four-GPU box as the worked example. Section 13.3 targets the bandwidth and overhead regimes: capturing the compute graph and letting a compiler fuse it can collapse both — though not every technique pays on every model, and Section 13.7 measures one that costs time when its constraint does not bind, which is why diagnosis comes first — contrasting torch.compile’s bytecode capture with jax.jit’s tracing. Section 13.4 turns to space: the memory anatomy of a training step, mixed precision, activation checkpointing, and gradient accumulation — the techniques that decide whether a model fits. Section 13.5 builds data parallelism from scratch, derives the ring allreduce, and measures the communication bill; Section 13.6 replaces the hand-rolled version with production data parallelism and contrasts PyTorch’s explicit collectives with JAX’s declarative sharding. Finally Section 13.7 runs the whole method on a real Transformer, taking one of the book’s own GPT models down a measured waterfall of every technique the chapter taught. That capstone runs in both frameworks, each in its own idiom — torch.compile, autocast, and DDP on the PyTorch side; jax.jit, explicit bf16 threading, and declarative sharding on the JAX side — and the two waterfalls diverge exactly where the frameworks do: what a “compile” rung even means, and how data parallelism is launched.

A word on the machine this chapter is built on, because it shapes what you will see. The book’s build box is four consumer RTX 4090 GPUs with no NVLink and — a deliberate market segmentation — no peer-to-peer transfer: every byte between two GPUs is staged through host memory over PCIe — tens of gigabytes per second at best, roughly two orders of magnitude below a datacenter NVLink fabric. This is not a handicap to apologize for; it is a teaching instrument. Most readers’ multi-GPU machines look like ours, not like a datacenter rack with a terabyte-per-second fabric, and a slow interconnect makes the accounting of parallel training impossible to ignore. The constants in this chapter are ours; the reasoning transfers to any machine, and the sections are written so that a reader on two GPUs, or one, sees the same story.

What This Chapter Is Not

The performance story is large, and this chapter draws sharp borders. Multi-node training — splitting a model across machines with tensor, pipeline, or expert parallelism, and the network fabrics that make it possible — is the province of the Language Models part, which has data large enough to warrant it; Section 13.6 builds the bridge and stops at the water’s edge. Kernel authoring in CUDA, Triton, or Pallas is fenced off book-wide (Section 5.4); we teach how to get performance from the operations you already have, and point to the resources below for those who want to write their own. Serving engines — continuous batching, paged key–value caches, speculative decoding, the systems that turn a trained model into a low-latency service — belong to the Language Models part as well; this chapter teaches the inference economics (the prefill-versus-decode roofline reading of Section 13.2) but not the engines that exploit them. Quantization as compression for inference is likewise deferred; Section 13.2 teaches formats as training precisions only. The production library map — which distributed framework to reach for at which scale, how to checkpoint a long run, how to launch across a cluster — and buying advice both live in the Tools appendix (Section 29.6, Section 29.4): this chapter earns the concepts at notebook scale, and the appendix names the products at datacenter scale.

Resources and Further Reading

The references below follow the chapter’s arc — the roofline and the regimes, the compiler and the memory, the collectives and the case study. All are freely accessible online. Where this chapter fences off a topic — kernel authoring, multi-node parallelism, serving — the resource that owns it is flagged as such.

Books and long-form

Courses and video lectures

Foundational and current papers

Tutorials, notes, and lore