13  Computational Performance

The same model and optimization method can differ substantially in wall-clock time and memory use across implementations. The difference depends on how well the program uses arithmetic units, limits data movement, amortizes framework overhead, and allocates device memory. This chapter develops a quantitative method for analyzing these effects.

The roofline model compares arithmetic throughput with memory bandwidth. Arithmetic intensity, measured in FLOPs per byte moved, determines which resource bounds an operation. A third regime appears when dispatch overhead dominates both. The practical procedure is to measure a program, classify its limiting regime, apply a corresponding optimization, and measure again. Section 13.1 introduces this procedure and the remaining sections apply it.

The sections proceed from individual operations to full models on multiple GPUs. Section 13.1 establishes the roofline, the three regimes, and the measurement discipline required by asynchronous framework dispatch. Section 13.2 explains where the roofline’s two numbers come from — the memory hierarchy, the tensor cores and their numerical formats, interconnects, and energy costs 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 can reduce both, and the section contrasts torch.compile’s bytecode capture with jax.jit’s tracing. Section 13.4 turns to memory use: the allocation 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 communicated bytes and latency; 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 complete procedure on a GPT model. The case study uses each framework in its usual idiom — torch.compile, autocast, and DDP on the PyTorch side; jax.jit, explicit bf16 threading, and declarative sharding on the JAX side — and the cumulative results differ in how compilation changes execution and how data parallelism is launched.

The build machine determines the constants in the measurements below. It contains 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 configuration represents the interconnect available in many workstations. 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. The same accounting applies to machines with one or two GPUs.

What This Chapter Is Not

This chapter has a limited systems scope. Multi-node training is the province of the Language Models part, which has data large enough to warrant it: splitting a model across machines with tensor, pipeline, or expert parallelism, and the network fabrics that make it possible, all belong there. Section 13.6 introduces the required concepts but does not cover multi-node execution. 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 belong to the Language Models part as well: continuous batching, paged key–value caches, speculative decoding, all the systems that turn a trained model into a low-latency service. 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 and hardware-selection guidance both live in the Tools appendix (Section 30.6, Section 30.4): which distributed framework to reach for at which scale, how to checkpoint a long run, and how to launch across a cluster. This chapter develops the concepts at notebook scale, while the appendix identifies systems for datacenter-scale use.

Resources and Further Reading

The references below follow the chapter’s progression from the roofline and performance regimes through compilation, memory, collectives, and the case study. All are freely accessible online. Resources on kernel authoring, multi-node parallelism, and serving are identified explicitly because those topics lie outside this chapter’s scope.

Books and long-form

Courses and video lectures

Foundational and current papers

Tutorials, notes, and lore