Efficient Training Techniques

Chapter 16's 6ND rule counts the FLOPs a training run requires; Chapter 18 spread those FLOPs across thousands of GPUs. This chapter asks a different question: of the compute a GPU is theoretically capable of, how much is actually used for useful work? The answer is often surprisingly low — and closing that gap is worth millions of dollars on a frontier run.

Model FLOPs Utilization (MFU) is the fraction of a GPU's peak theoretical throughput that a training run actually achieves on useful model computation. A well-optimized large-model run reaches perhaps 40–55% MFU; a naive one can be far lower. The gap between achieved and peak is the target of every technique in this chapter.

Why isn't MFU 100%? Because GPUs spend much of their time NOT doing matrix multiplications: waiting on memory, launching kernels, communicating between devices, and running non-matmul operations (softmax, normalization, activation) that the hardware's tensor cores cannot accelerate. The techniques in this chapter attack each of these inefficiencies.

To understand efficiency, you must understand where time actually goes on a GPU. The central fact: arithmetic is cheap, but moving data between levels of the memory hierarchy is expensive. Most deep-learning operations are bottlenecked not by compute but by memory bandwidth — they are 'memory-bound', spending their time waiting for data rather than calculating.

Each level down is far larger but far slower. SRAM is ~6× the bandwidth of HBM but thousands of times smaller. The art of efficient GPU code is keeping data in the fast levels as long as possible — doing as much work as you can on each byte before it has to travel back to slow HBM.

An operation's 'arithmetic intensity' — FLOPs performed per byte of memory moved — determines whether it is compute-bound or memory-bound. Large matrix multiplications have high intensity (lots of compute per byte) and are compute-bound: they actually use the tensor cores. Elementwise operations (add, activation, normalization) have low intensity and are memory-bound: they spend their time shuttling data, leaving the tensor cores idle.

Chapter 11 introduced gradient checkpointing; here we examine it as a core efficiency technique with quantitative trade-offs. The backward pass needs the forward activations to compute gradients, and storing them all costs memory proportional to depth × batch × sequence. Activation recomputation discards most of these and recomputes them on demand during the backward pass.

The standard choice — checkpointing at layer boundaries — stores only the input to each Transformer block and recomputes the block's internals during backprop. This typically costs about 33% extra compute for a large reduction in activation memory, enabling bigger models or longer sequences on the same GPU.

A refinement (Korthikanti et al., 2022) recognizes that not all activations are equally costly to recompute. The attention softmax is cheap to recompute but expensive to store; the big matmuls are the opposite. Selective recomputation stores the expensive-to-recompute activations and recomputes only the cheap-to-recompute ones, getting most of the memory savings for far less than 33% overhead.

Kernel fusion is the single most important idea in GPU efficiency. A 'kernel' is one GPU function launch. Running a chain of separate kernels — say, a matmul, then a bias add, then a GELU — forces each intermediate result to be written to slow HBM and read back by the next kernel. Fusion combines them into ONE kernel that keeps intermediates in fast registers/SRAM, eliminating the round trips.

Consider computing GELU(x @ W + b). Unfused, this is three kernels: the matmul writes its output to HBM, the bias-add reads and rewrites it, and GELU reads and rewrites it again — the data crosses the slow HBM boundary multiple times. A fused kernel computes all three in one pass, with the intermediate values living only in fast on-chip memory. For memory-bound elementwise chains, fusion can give large speedups.

Historically, writing fused GPU kernels meant CUDA C++ — difficult, low-level, and slow to iterate. Triton (Tillet et al., 2019), an open-source language from OpenAI, lets you write high-performance fused kernels in Python-like syntax. It handles the hardest parts — memory coalescing, SRAM management, and parallelization — automatically, while you express the computation at the level of blocks of data.

In Triton, you write a kernel that operates on BLOCKS of data rather than individual elements. The kernel is launched over a grid of program instances, each handling one block. You explicitly load blocks from HBM into SRAM, compute on them, and write results back — but Triton manages the low-level details. This block-level model maps naturally onto the GPU's architecture.

The real payoff is fusing a chain that would otherwise round-trip to HBM. A fused 'linear + bias + activation' kernel loads the inputs once, performs the matmul, adds the bias, applies the activation — all while the data stays in SRAM — and writes only the final result. The same pattern, scaled up with the online-softmax trick, is exactly how FlashAttention is implemented in Triton.

Chapter 19 introduced FlashAttention as an architecture-level efficiency win. Here we revisit it as the canonical example of the efficiency principles of this chapter: it is a fused kernel that exploits the memory hierarchy through tiling and an online softmax, avoiding the slow-memory round trips that make naive attention inefficient.

The obstacle to fusing attention is the softmax: it seems to need the entire row of scores at once (to compute the max and the sum for normalization). FlashAttention's key algorithmic insight is the ONLINE softmax — a way to compute the softmax incrementally, block by block, by maintaining a running maximum and running sum and rescaling as new blocks arrive. This lets attention be tiled and fused without ever materializing the full score row.

Each block of scores lives only in SRAM, is consumed immediately, and is discarded. The running statistics (m, l, O) are tiny. The result is bit-for-bit the same as standard attention, but the HBM traffic drops from O(T²) to O(T). This is the union of every idea in this chapter: fusion, tiling, the memory hierarchy, and a clever algorithm to make it all possible.

Chapter 15 covered bf16 mixed precision. The frontier of low-precision training is fp8 — eight-bit floating point — supported natively on Hopper (H100) and newer GPUs. fp8 roughly doubles throughput and halves memory versus bf16, but its tiny dynamic range demands careful scaling to avoid overflow and underflow.

fp8 has two variants: E4M3 (4 exponent bits, 3 mantissa) offers more precision for forward activations, while E5M2 (5 exponent, 2 mantissa) offers more range for gradients. Training in fp8 keeps a higher-precision master copy and uses per-tensor scaling factors to map values into fp8's narrow range — the same mixed-precision philosophy as bf16, pushed further.

Writing kernels by hand is powerful but laborious. Graph compilers automate much of it: they capture the model's computation as a graph, then apply optimizations — fusing operations, eliminating redundant work, generating efficient kernels — automatically. PyTorch's torch.compile is the most widely used; it can deliver substantial speedups with a single line of code.

The techniques in this chapter are complementary, and a well-tuned training run uses most of them together. Here is how they stack into a coherent, efficient configuration — the practical synthesis of the chapter.

This chapter has focused on efficient PRETRAINING. A related family of techniques makes FINE-TUNING efficient, and they share the same spirit — doing more with less compute and memory. They become central in Part V (alignment), but it is worth previewing them here as efficiency techniques.

Full fine-tuning updates all of a model's billions of parameters — expensive in compute and memory, and producing a full-size copy per task. LoRA (Hu et al., 2021) freezes the pretrained weights and learns only small low-rank update matrices added alongside them. A weight update ΔW is factored as the product of two thin matrices BA, so instead of training d×d parameters you train only 2×d×r for a small rank r.

Exercises 1–10 are pen-and-paper or derivations; 11–20 require code.

Further reading: “FlashAttention” and “FlashAttention-2/3” (Dao et al., 2022–2024). “Reducing Activation Recomputation in Large Transformer Models” (Korthikanti et al., 2022) for selective recomputation. “Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations” (Tillet et al., 2019) and the Triton tutorials. The torch.compile and TorchInductor documentation. “LoRA” (Hu et al., 2021) and “QLoRA” (Dettmers et al., 2023). NVIDIA's Transformer Engine documentation for fp8 training.

Next → Chapter 21: Evaluation During Pretraining

You now have an efficient, distributed training run consuming clean, curated data — but how do you know if it is working? Chapter 21, the close of Part IV, covers evaluation during pretraining: tracking loss and perplexity, the benchmark suites (MMLU, HellaSwag, and others) used to measure capability, the perils of benchmark contamination, and how to read the signals of a months-long run to decide whether to continue, adjust, or stop. We close Part IV by turning the training process from a black box into something you can measure and steer.