Training infrastructure

Chapter 9 ended with the complete conceptual story of distributed training. Choose a compute budget. Allocate it per Chinchilla, about $20$ tokens per parameter. Combine data, tensor, pipeline, and FSDP parallelism into a 3D layout across thousands of GPUs. Train. That is the math.

This chapter is the reality. What GPU does the cluster actually use, H100, A100, MI300X, and what does each offer in memory, bandwidth, and FLOPs? How are those GPUs wired together inside a node (NVLink) and across nodes (InfiniBand)? What software turns the abstract parallelism strategies of Chapter 9 into running kernels, PyTorch FSDP, DeepSpeed, Megatron-LM? When does a team write its own GPU kernel in Triton, and what does FlashAttention’s success story actually demonstrate? And what does the whole thing cost?

A modern LLM training run is a systems-engineering problem as much as a machine-learning one. FlashAttention did not change the math of attention; it changed how attention is computed, and that alone made long-context LLMs practical. DeepSpeed did not invent new optimizers; it implemented existing math in a way that runs $70$B-parameter models on hardware that previously could not hold them. Modern training is unrecognizable from five years ago, not because the math changed but because the engineering did.

This chapter is denser on practical reality than earlier ones. Readers who have followed Chapters 1 through 9 already have the conceptual training story end to end. Chapter 10 is optional depth, written for engineers at AI labs, researchers building their own training stacks, and anyone whose job involves the question “why is our MFU only $30\%$?”. Read it if that is you. Skip to Chapter 11 if not, the rest of the tutorial does not assume Chapter 10’s content.

The setup: why infrastructure matters

The headline claim of this chapter is one sentence: at frontier scale, the bottleneck on LLM training is as much systems as algorithms. FlashAttention, DeepSpeed, and Triton each made possible model scales that were “obviously” impossible the year before, not by changing the underlying mathematics but by changing how that mathematics gets executed on real hardware.

The shift is visible in where the research field’s attention has moved. The years from 2017 through 2020 were dominated by architectural innovation: the transformer itself, multi-head attention, positional encodings, scaling-up of decoder-only stacks. The years from 2020 through 2024 have been increasingly dominated by systems innovation: ZeRO and FSDP for memory sharding, FlashAttention for I/O-aware attention, Triton as a higher-level kernel language, mixed precision moving from FP16 to BF16 to FP8, NCCL tuning for the cross-node collectives that make 3D parallelism actually run. Most of the practical performance gain from GPT-3 (2020) to Llama-3 (2024) came from infrastructure improvements, not architectural redesign. The decoder-only transformer of GPT-3 and the decoder-only transformer of Llama-3 are the same architecture in any meaningful sense; what changed was how they are trained.

What this chapter covers is the practical stack, the hardware, the interconnects, the software frameworks, the custom kernels, the optimization techniques like activation checkpointing and mixed precision, and the resulting cost economics. What it deliberately does not cover is GPU microarchitecture in depth, warp scheduling, SIMT lanes, register allocation, the kind of detail that fills the CUDA programming guide. Those topics belong in a graduate course on GPU computing. The intent here is to give a reader enough working knowledge to reason about the trade-offs between configurations without turning them into a GPU kernel engineer.

The hardware: GPUs and what they offer

GPU spectrum (as of 2026-07-13)

Frontier LLM training runs on one of a small handful of accelerators. The list is short because the engineering investment per platform is enormous and few vendors can support it, and the roster turns over roughly annually, so treat the specific models and numbers below as a snapshot rather than a fixed inventory.

The NVIDIA H100 (Hopper architecture, $2022$–present) was the workhorse of the 2023–2024 training generation and remains one of the most widely deployed accelerators in production, even though it is no longer NVIDIA’s current-generation part. $80$ GB of HBM3 memory at $\sim 3$ TB/s bandwidth, around $989$ TFLOPS of BF16 matmul throughput on the tensor cores, $\sim 1979$ TFLOPS at FP8. NVLink 4.0 provides $900$ GB/s of GPU-to-GPU bandwidth within a node. Llama-3, GPT-4-class models, Mistral’s larger runs, all trained on H100 clusters as their default substrate, and a large fraction of the world’s installed GPU base is still H100.

The NVIDIA H200 ($2024$–present) is the same Hopper compute die with a memory upgrade: $141$ GB of HBM3e at $\sim 4.8$ TB/s, same FLOPs as the H100. The extra memory matters most for the workloads where $80$ GB is the binding constraint, long-context training and large-batch or long-context inference, and clusters that already have H100 nodes can often drop in H200s without touching the rest of the stack.

The NVIDIA B200 (Blackwell architecture, $2024$–present) began NVIDIA’s current Blackwell generation of training and inference GPUs. A dual-die package with $180$ GB of HBM3e at $\sim 8$ TB/s of bandwidth, native FP4 tensor-core support delivering on the order of $9$ PFLOPS of dense FP4 per GPU ($\sim 18$ PFLOPS with structured sparsity), and NVLink 5.0 at $1.8$ TB/s of per-GPU bidirectional bandwidth, double the H100’s NVLink 4.0. The GB200 superchip pairs two B200 GPUs with a Grace CPU, the same coupling pattern as GH200 one generation back; the GB200 NVL72 rack links $72$ B200 GPUs and $36$ Grace CPUs into a single NVLink domain. NVIDIA followed in 2025 with Blackwell Ultra (B300/GB300), a $288$ GB HBM3e refresh with higher FP4 throughput and a $1400$ W TDP that requires liquid cooling. Given the pace of releases, this is a moving target: NVIDIA first unveiled its next platform, Vera Rubin (Rubin GPU, $288$ GB of HBM4), at GTC 2025, with partner availability slated for the second half of $2026$, so B200/GB300 may already be one generation behind the frontier by the time you read this.

The NVIDIA A100 (Ampere, $2020$) is now two generations behind the frontier but still runs a meaningful share of production training. $40$ or $80$ GB of HBM2e, $\sim 2$ TB/s of memory bandwidth, $\sim 312$ TFLOPS BF16, NVLink 3.0 at $600$ GB/s. GPT-3, the original Chinchilla, and LLaMA 1 all trained on A100s, and plenty of production training still runs on A100 clusters that were stood up before H100 became available.

The NVIDIA GH200 (Grace Hopper Superchip, $2023$–present) tightly couples an H100-class GPU with a Grace CPU and gives each GPU $96$ or $144$ GB of HBM3 plus a unified memory model. It was the bridge design between plain H100 nodes and the Grace-coupled GB200 generation described above.

The AMD MI300X (CDNA 3, $2023$) was AMD’s first credible frontier-scale alternative to NVIDIA. $192$ GB of HBM3, about $2.4\times$ what an H100 offers, at $5.3$ TB/s, and roughly $1.3$ PFLOPS of BF16. AMD followed with the MI325X (CDNA 3, late $2024$), a memory-upgraded variant at $256$ GB of HBM3e and $6$ TB/s, and then the MI355X (CDNA 4, general availability October $2025$), which reached $288$ GB of HBM3e at $8$ TB/s and added native FP4/FP6 support, putting it roughly in the same generation as Blackwell. AMD’s software stack (ROCm) has historically lagged CUDA, but adoption is growing, partly because the larger memory per GPU makes it easier to fit big models with less aggressive sharding.

Memory hierarchy

A GPU is not a flat memory device. Every kernel author has to think about a hierarchy that spans roughly five orders of magnitude in bandwidth.

• Registers, thousands of bytes per thread, single-cycle access, the fastest tier.
• Shared memory / L1 cache, $\sim 228$ KB of configurable shared memory per streaming multiprocessor (SM) on an H100 (up from $\sim 128$ KB in older architectures), addressable from any thread in a block.
• L2 cache, $\sim 50$ MB total on an H100, shared across all SMs.
• HBM (high-bandwidth memory), $80$ GB on an H100, $\sim 3$ TB/s aggregate, the GPU’s main RAM.
• Host memory (CPU RAM), hundreds of GB but cross-PCIe, an order of magnitude slower than HBM for GPU access.

The performance lesson is that keeping working data in the fastest tier that fits dominates everything else. FlashAttention’s central trick (Section 6) is to keep attention’s intermediates in shared memory rather than streaming the $N \times N$ attention matrix through HBM. The asymptotic memory-bandwidth savings are what made it practical to push context lengths from $2$K to $128$K without melting the GPU.

Cluster scale

The unit of a modern training cluster is the node, typically $8$ GPUs connected by NVLink through an NVSwitch, plus high-bandwidth network NICs (one or two InfiniBand HCAs per GPU). $128$ nodes is a $1024$-GPU cluster, sufficient for most $7$B–$30$B training runs. $1024$ nodes is an $8192$-GPU cluster; that was the working scale of major lab training circa $2024$. GPT-4-class frontier models reportedly trained on $10{,}000$ to $25{,}000$ GPUs. The frontier keeps moving: as of 2026-07-13, individual clusters have scaled well past that, xAI’s Colossus site has been reported at roughly $555{,}000$ GPUs across its expansions, and Microsoft’s first Blackwell campus for OpenAI’s Stargate effort is reported at $450{,}000$ GPUs, so treat any specific cluster-size figure as a snapshot of a number that is still climbing.

What turns those individual GPUs into a working cluster is the network.

The interconnects: NVLink and InfiniBand

NVLink within a node

NVLink is NVIDIA’s GPU-to-GPU interconnect for within-node communication. The generations matter: NVLink 3.0 on A100 delivers $600$ GB/s of bidirectional bandwidth per GPU; NVLink 4.0 on H100 delivers $900$ GB/s; NVLink 5.0 on B200/GB200 delivers $1.8$ TB/s. These are per-GPU numbers, each GPU can simultaneously send and receive at near-NVLink-peak.

The piece that makes a multi-GPU node work end-to-end is NVSwitch, a chip that provides full-mesh NVLink connectivity. With NVSwitch, every GPU in an 8-GPU node can talk to every other GPU at full NVLink bandwidth, no contention, no hop penalties. Without NVSwitch, GPUs are wired in a fixed topology (typically a hybrid mesh-and-cube) and some pairs are farther apart than others. Modern training nodes all use NVSwitch; the topology question disappears within a node.

InfiniBand across nodes

Once a workload crosses the node boundary, GPU-to-GPU communication leaves NVLink and goes over the network. The dominant fabric is InfiniBand: a low-latency, high-bandwidth switched network. Older HDR/NDR generations run at $25$ to $50$ GB/s per port; the current XDR generation (NVIDIA Quantum-X800, paired with ConnectX-8 NICs, rolling out through 2025–2026 alongside Blackwell clusters) reaches $800$ Gb/s, i.e. $\sim 100$ GB/s, per port. A less common alternative is RoCE (RDMA over Converged Ethernet), which gives similar bandwidths over Ethernet hardware.

The critical feature for training is GPU-direct RDMA: data moves directly from one GPU’s HBM, into the NIC, across the fiber, into the destination NIC, and into the destination GPU’s HBM, without ever staging through CPU memory or hitting a system-call boundary. This is the path NCCL uses when running collectives across nodes; everything else would be much too slow.

Why the gap dictates parallelism

The bandwidth ratio between NVLink and InfiniBand is the single most consequential number in a cluster’s design. NVLink $\sim 900$ GB/s, InfiniBand $\sim 25$–$50$ GB/s, a $\sim 20$ to $40\times$ gap. Communication patterns that are fine within a node become bottlenecks across nodes.

This is exactly why Chapter 9 noted that tensor-parallel rank is typically capped at $8$: TP requires an all-reduce per layer per forward and per backward pass, on activation-sized tensors. That volume is only tolerable at NVLink bandwidth. Pipeline-parallel rank, which communicates only at stage boundaries, can comfortably span many nodes over InfiniBand. Data-parallel rank is even more forgiving, DP communicates once per step, and the all-reduce can be tuned to fit available cross-node bandwidth.

The practical rule that falls out: TP within a node, PP across nodes, DP wherever capacity exists. The Megatron-Turing $530$B configuration of TP$=8$, PP$=35$, DP$=8$ is not arbitrary, each factor is placed on the network layer that can actually sustain its communication.

NCCL: collectives on GPUs

NCCL (the NVIDIA Collective Communications Library, pronounced “nickel”) implements collective operations, all-reduce, all-gather, broadcast, reduce-scatter, all-to-all, on GPU networks. Every modern PyTorch distributed primitive calls NCCL under the hood when the backend is nccl; the same is true for the distributed paths in DeepSpeed and Megatron-LM. On GPU clusters, NCCL is the only collective implementation that matters in practice.

The reason for a dedicated library, rather than MPI or some generic message-passing layer, is that collectives must be implemented carefully against the actual hardware topology. NCCL is topology-aware: it detects the NVLink-and-InfiniBand layout of the cluster, chooses the right collective algorithm for the data size and rank count, and uses GPU-direct paths for cross-node communication. A generic MPI all-reduce on a GPU cluster is typically several times slower than the equivalent NCCL call, because the MPI implementation does not have the GPU memory-layout and NIC-offload knowledge that NCCL has.

The default all-reduce algorithm, the one that dominates DP gradient reduction, is the ring all-reduce that Chapter 9 introduced. Each GPU sends a chunk of its tensor to its right neighbor and receives a chunk from its left, accumulating partial sums; after $n - 1$ rounds, every GPU has the sum of every chunk. The total per-GPU communication is $2 \cdot (n-1)/n \cdot \text{tensor\_size}$, which approaches $2 \cdot \text{tensor\_size}$ in the large-$n$ limit, and crucially, that cost is independent of $n$. The ring is nearly optimal for moderately-sized clusters. For very large rank counts (above $\sim 32$ nodes), NCCL switches to a tree all-reduce that has lower latency at the cost of slightly higher bandwidth utilization.

Practical NCCL tuning is a real engineering activity at scale. Environment variables like NCCL_ALGO (choosing ring vs tree), NCCL_PROTO (choosing the wire protocol, Simple, LL, LL128), NCCL_IB_HCA (binding to specific InfiniBand NICs), and NCCL_NSOCKS_PERTHREAD interact with cluster topology in non-obvious ways. A well-tuned NCCL configuration is commonly $20$–$30\%$ faster than the defaults on the same hardware, and finding the right configuration is part of the production-readiness work for any new cluster.

The runnable code below simulates the result of a ring all-reduce, not the round-by-round communication, but the final accumulation that every rank ends up with, and computes the communication-cost saving relative to the naive all-to-all approach.

NCCL is the substrate. The frameworks built on top of it orchestrate everything else.

Training frameworks: which one to pick

PyTorch FSDP (the default)

PyTorch FSDP (torch.distributed.fsdp.FullyShardedDataParallel) is PyTorch’s native implementation of ZeRO-3, exposed through a clean wrapper API that requires almost no changes to a single-GPU training loop. It has become the modern default for most teams below frontier scale.

FSDP’s strengths are integration and ergonomics. It is documented as a first-class PyTorch feature, plays well with the rest of the ecosystem (checkpointing utilities, profilers, the rest of torch.distributed), and gets bug-fixes on the PyTorch release cadence. For any model in the $1$B–$30$B range that fits on cluster sizes of $8$–$256$ GPUs, FSDP is almost always the right starting choice.

Its limitation is that it implements only the data-parallel-with-sharding axis. Tensor parallelism and pipeline parallelism are not native to FSDP, combining FSDP with TP or PP requires either composing in another library (Megatron’s TP into an FSDP-wrapped model) or stepping out of FSDP entirely.

DeepSpeed

DeepSpeed (Microsoft) is the original ZeRO implementation and remains the framework with the most aggressive set of memory-savings features. ZeRO-Infinity offloads parameters and optimizer state to CPU and even NVMe, making it possible to train models far larger than the cluster’s GPU memory at the cost of higher communication. 1-bit Adam compresses the all-reduce of optimizer-state updates by a factor of $\sim 32$, which matters on bandwidth-constrained clusters.

DeepSpeed is the right pick when memory pressure is severe, when training a model that simply will not fit under FSDP at any rank count, or when running on an interconnect whose all-reduce bandwidth is the dominant bottleneck.

Megatron-LM

Megatron-LM (NVIDIA) is the reference implementation of tensor parallelism and pipeline parallelism. The TP recipe Chapter 9 described, the column-then-row FFN split, the head-dimension attention split, is Megatron’s recipe; almost every other TP implementation follows the same pattern.

Megatron is the right pick for any model in the $30$B-and-up range that needs 3D parallelism. It is usually combined with another library (Megatron-DeepSpeed, Megatron-Core integrated with PyTorch FSDP) to add a ZeRO-style sharding axis on top of TP and PP. The Megatron-Turing $530$B model, the BLOOM 176B run, and most frontier internal stacks use Megatron-derived 3D parallelism.

NeMo and Lightning

NVIDIA NeMo is not a fifth parallelism strategy competing with the three above; it is an opinionated wrapper around Megatron-Core, the same TP/PP engine underlying Megatron-LM, packaged with production-tested model recipes, data loaders, and launch configs for training and fine-tuning LLMs on NVIDIA hardware. Where raw Megatron-LM asks a team to hand-write its own 3D-parallelism configuration, NeMo trades some of that flexibility for a curated set of known-good configs, at the cost of being more tightly coupled to NVIDIA’s release cadence than Megatron-LM or FSDP alone.

PyTorch Lightning solves a different problem, it does not implement a new sharding strategy at all. Lightning is a training-loop framework that wraps whatever backend a team chooses, DDP, FSDP, DeepSpeed, behind a single Trainer API, so switching strategies is a config change rather than a training-loop rewrite. Lightning’s FSDP integration is a thinner, more ergonomic layer over the same PyTorch FSDP described above, popular with teams that want FSDP’s memory sharding without hand-rolling the wrapping and checkpointing boilerplate that raw FullyShardedDataParallel requires.

When to use what

The decision flow is roughly:

• Under $1$B parameters, single node: vanilla DP via PyTorch’s DistributedDataParallel. There is no memory pressure that justifies the complexity of anything else.
• $1$B to $30$B parameters: PyTorch FSDP. Simple to set up, well-tested, and the memory sharding makes the model fit on commodity hardware.
• $30$B to $100$B parameters: Megatron-LM for tensor and pipeline parallelism, with an FSDP-style data axis layered on top. The model is too big for FSDP alone.
• $100$B+ parameters: Megatron-DeepSpeed or an equivalent combination. Full 3D parallelism with all three of FSDP, TP, and PP, plus the more aggressive memory tricks DeepSpeed offers.

These frameworks orchestrate the training loop, but for the most performance-critical operations they often invoke custom GPU kernels rather than relying on PyTorch’s defaults, hand-written or compiler-generated kernels like the ones Triton and FlashAttention produce.

Custom kernels: Triton and FlashAttention

Why custom kernels

PyTorch’s default kernels, the ones generated automatically when an nn.Module is called, are good. They use cuBLAS and cuDNN for the big matmuls, and they are correct in the sense that they produce the right answer. They are not always optimal. There are three reasons a team writes a custom kernel.

Operation fusion. A typical PyTorch sequence like linear → bias → activation runs as three separate kernel launches with intermediate tensors written to HBM between them. A fused kernel does all three in a single launch and never materializes the intermediates. The win is exactly the HBM bandwidth that was saved.

Memory hierarchy exploitation. PyTorch’s default kernels generally treat the GPU as a flat memory device. A hand-written kernel can deliberately keep working data in shared memory across a sequence of operations that would otherwise spill to HBM on every step. This is the lever FlashAttention pulls.

Eliminating dispatch overhead. Each PyTorch operation has Python-level dispatch cost, small per op, large in aggregate over a training step that issues thousands of kernel launches. A custom kernel collapses many small ops into one launch.

Triton

The traditional path to a custom GPU kernel is CUDA C++, verbose, low-level, and requiring significant GPU-architecture expertise to get right. Triton, originally developed as an academic project (Tillet et al. 2019) and later adopted and popularized by OpenAI, is the practical alternative: a Python-embedded language for writing GPU kernels at a higher level than CUDA, with a JIT compiler that produces competitive PTX.

The Triton model is straightforward. A function decorated with @triton.jit is compiled to a GPU kernel. The kernel processes work in tiles, indexed by program IDs (tl.program_id). Inside the kernel, the programmer issues explicit memory loads (tl.load) and stores (tl.store), with masks for boundary handling. Most of CUDA’s complexity, warp scheduling, shared-memory layout, register pressure, is abstracted away or handled by the compiler.

The result is that a Triton implementation of a non-trivial kernel is roughly an order of magnitude shorter than the equivalent CUDA, with performance that is competitive for most patterns and sometimes superior for patterns the CUDA compiler does not optimize well. For many internal teams, Triton has effectively replaced direct CUDA for new kernel work.

FlashAttention

The canonical success story of custom-kernel work in deep learning is FlashAttention (Dao et al. 2022, arxiv.org/abs/2205.14135).

Standard attention computes $S = QK^\top \in \mathbb{R}^{N \times N}$ in HBM, runs a softmax on it ($S \to P$, also in HBM), and then computes $O = PV$. The total HBM I/O is $O(N^2 + Nd)$, dominated by the $N \times N$ attention matrix, which has to be written, read for the softmax, and read again for the matmul against $V$. The memory footprint is $O(N^2)$. For $N = 32{,}000$, the attention matrix alone is over $2$ GB in FP16.

FlashAttention processes attention in tiles. Each query block is loaded into SRAM. The corresponding key and value tiles are streamed through SRAM; the partial $S$, the online-softmax statistics, and the partial output $O$ are all kept on-chip. The full $N \times N$ attention matrix is never written to HBM. The total HBM I/O drops to $O(N \cdot d \cdot N / M)$ where $M$ is SRAM size, much smaller for large $N$. The memory footprint drops from $O(N^2)$ to $O(N d)$.

The result is $2$ to $4\times$ speedup at standard context lengths ($2$K to $8$K) and $10\times$ or more at long contexts ($32$K to $128$K). Crucially, the output is identical to standard attention, no approximation, same math, just a different implementation. FlashAttention-2 (Dao 2023, arxiv.org/abs/2307.08691) refined the work-partitioning across GPU threads and is the current default in PyTorch’s scaled_dot_product_attention when running on supported hardware. This section covers the kernel-engineering side, tiling, SRAM, HBM I/O; Chapter 17 revisits FlashAttention from the inference-serving angle, how it interacts with the KV cache and continuous batching once the model is actually being served.

Hand-writing a kernel is one route to fused, hardware-aware execution, the one a human takes when they have noticed a specific opportunity like FlashAttention’s. The other route is to let a compiler find fusion opportunities automatically, across a much wider range of code than any one hand-written kernel covers.

Compiler stacks: torch.compile and JAX/XLA

torch.compile

torch.compile (PyTorch 2.0, 2023) takes a different path from Triton and FlashAttention. Instead of a human noticing a fusion opportunity and hand-writing a kernel for it, torch.compile traces a model’s ordinary Python and generates fused kernels automatically. Two subsystems do the work. TorchDynamo hooks into the CPython frame evaluation API and traces the operations a model actually executes into an FX graph, falling back to normal eager execution wherever it hits code it cannot safely trace, a graph break. TorchInductor takes that graph and lowers it to fused, hardware-specific kernels, generating Triton kernels for GPU targets and C++ for CPU targets. A meaningful share of the “someone should write a custom kernel for this” work described above is now something torch.compile does automatically for common patterns, no @triton.jit required.

JAX, XLA, and Google’s stack

Everything in this chapter so far is the PyTorch ecosystem. Google’s training stack is built on a different foundation. JAX is a NumPy-like array library with automatic differentiation, composed with XLA (Accelerated Linear Algebra), a compiler that lowers JAX’s traced computation to fused kernels for TPUs and GPUs alike. Where PyTorch executes eagerly by default and compiles opt-in via torch.compile, JAX is compile-first: a function wrapped in jax.jit is traced once, and the entire computation, not individual ops, is handed to XLA, which fuses and schedules it as a single optimized program.

JAX’s parallelism story runs through GSPMD (General and Scalable Parallelism for ML Dataflow Graphs), XLA’s automatic partitioner. A user annotates how tensors should be sharded across a device mesh, and GSPMD’s compiler pass derives the collective operations, all-gathers, reduce-scatters, all-to-alls, needed to execute the sharded computation, a role analogous to what Megatron-LM’s hand-written TP/PP code does explicitly on the PyTorch side. T5X and its successor MaxText are Google’s reference training codebases built on JAX, XLA, and GSPMD, the rough equivalent of Megatron-LM or NeMo for teams working in the JAX ecosystem; Gemini and PaLM are both reported to train on TPU pods through this stack.

For a reader coming from the PyTorch side, the practical takeaway is not “learn JAX.” It is recognizing that a technical report mentioning TPU pods, pjit, or GSPMD sharding annotations is describing the same 3D-parallelism ideas Chapter 9 introduced, expressed through a compiler-first framework rather than PyTorch’s eager-plus-library approach.

Compilers and hand-written kernels both exist to make a training step run efficiently on the hardware it is given. The remaining practical techniques work at a different level: trading compute for memory, and choosing a numeric precision.

Practical concerns: checkpointing, mixed precision, activations

Activation checkpointing

The forward pass of a transformer produces intermediate activations at every layer, the attention output, the post-LayerNorm vectors, the FFN intermediate, and so on. Every one of those activations has to be available during the backward pass to compute gradients. For a $70$B-class model with $80$ layers, $8192$ hidden dimension, sequence length $2048$, and a micro-batch of $4$, the per-micro-batch activations sum to over $100$ GB.

Activation checkpointing breaks this memory cost by not storing all activations. The training pipeline checkpoints only the activations at certain layer boundaries; during the backward pass, the intermediate activations between checkpoints are recomputed from the nearest checkpoint by re-running a slice of the forward pass.

The classic trade is $\sim 30\%$ extra forward-pass compute for $\sim 10\times$ memory reduction. For any model large enough that activation memory is the binding constraint, which is essentially every model above $\sim 10$B, that is a trade you take every time. Selective activation checkpointing (Korthikanti et al. 2022, arxiv.org/abs/2205.05198) refines this further: rather than checkpoint everything, recompute only the most memory-intensive activations (the QKV projections and the FFN intermediate), which delivers roughly half the memory reduction of full checkpointing ($\sim 5\times$) at a much smaller $\sim 10\%$ extra compute cost.

Mixed precision

Three numeric formats matter for current training.

FP32, $32$-bit IEEE float, the long-standing default for both compute and storage. The current role of FP32 is as the master precision: master weights and optimizer state ($m$, $v$ in AdamW) are kept in FP32 for numerical stability.

BF16 (bfloat16), $16$-bit with $8$ exponent bits and $7$ mantissa bits. Same dynamic range as FP32, half the storage, twice the matmul throughput on tensor cores. BF16 is the current default for forward and backward pass. The migration from FP16 to BF16 around $2020$–$2022$ was driven by FP16’s underflow issues at large scale, gradient magnitudes can fall below FP16’s representable range, and loss scaling is required to keep training stable. BF16’s wider exponent range eliminates the underflow problem entirely; the standard recipe is now FP32 master + BF16 compute.

FP8, $8$-bit floats, in two formats (E4M3 for forward, E5M2 for backward). Hopper-generation GPUs (H100 and later) support FP8 in hardware. The throughput gain is roughly $2\times$ over BF16, at the cost of more careful loss-scaling and per-tensor scale management. FP8 is in production at the frontier, DeepSeek-V3 ($2024$) trained the majority of its matmuls in FP8, while Llama-3 used BF16 throughout pretraining and reserved FP8 for post-training inference quantization of the $405$B model, but the bookkeeping required to make FP8 numerically stable during training is non-trivial.

Step timing and overlap

A single training step has phases: forward pass, backward pass, optimizer step. Between and within those phases, the cluster issues collective communications, DP gradient all-reduce, FSDP all-gather and reduce-scatter, TP per-layer all-reduce, PP activation send/receive. A naive implementation executes compute and communication sequentially; a well-tuned implementation overlaps them, issuing the next communication while the current compute is still running, prefetching the next FSDP-sharded parameters one layer ahead of when the kernel needs them.

The metric that captures all of this is MFU (Model FLOPs Utilization): actual FLOPs delivered to the model, divided by the cluster’s theoretical peak FLOPs. MFU is the single most useful number for grading a training run’s engineering quality.

• Un-tuned PyTorch training: $20$–$30\%$ MFU.
• Well-tuned production training: $40$–$50\%$ MFU.
• Industry-leading transformer training: $50$–$55\%$ MFU.

The difference between $25\%$ MFU and $50\%$ MFU is mostly the difference between sequential execution and overlapped execution. The widget below makes the two regimes side-by-side visible.

The runnable below makes the activation-checkpointing trade concrete, three regimes (none, selective, full) for a $70$B-class layer stack, and prints the per-strategy memory footprint.

All of that engineering optimizes a training run that is actually running. At frontier scale, keeping it running at all is its own discipline.

Observability and fault tolerance at scale

Watching the run: W&B and TensorBoard

A frontier training run can last weeks to months; nobody watches it by staring at a terminal. Weights & Biases (W&B) and TensorBoard are the two dominant tools for training observability. Both log scalar time series, loss, gradient norm, learning rate, MFU, tokens/sec, and render them as a live dashboard an engineer checks periodically rather than continuously. W&B additionally tracks run configuration, hyperparameters, and per-GPU hardware utilization, and supports comparing many runs side by side, which matters when a lab is running several ablations against the same frontier candidate at once.

The signal that catches problems early is rarely the training loss curve alone. A loss spike, a sudden jump that does not recover on its own, is the most common early warning that something has gone wrong, a corrupted data shard, a numerical instability in mixed-precision training, or a hardware fault silently corrupting a gradient. Gradient norm tracking often catches instability before it reaches the loss: a norm that grows unboundedly over several steps usually precedes a visible divergence. Per-GPU MFU dashboards catch a different failure mode entirely, a single straggler node running at half speed quietly drags down the whole cluster’s aggregate throughput well before the effect is visible in the loss curve.

Fault recovery: distributed checkpoints and elastic restart

At the scale of a $1024$-plus GPU cluster running for weeks, hardware failure is not an edge case, it is a certainty. A bad GPU, a flaky NIC, or a power event will eventually take down some fraction of the fleet, and the run has to survive it without losing weeks of progress.

The standard defense is periodic distributed checkpointing: at a fixed interval, often every $30$ minutes to a few hours depending on cluster size and observed failure rate, every rank writes its shard of the model weights, optimizer state, and data-loader position to durable storage, coordinated so all ranks checkpoint a consistent view of the same training step. When a node fails, an orchestrator (Slurm or Kubernetes, typically) detects the failure and relaunches the job, and every rank reloads its shard from the most recent checkpoint, resuming from that step rather than from scratch.

The checkpoint interval is itself a trade-off. Checkpointing too rarely means a failure loses more compute, hours of progress vanish if the last checkpoint is stale. Checkpointing too often wastes compute on I/O: writing a $70$B-parameter model’s full state, weights, gradients, optimizer moments, is itself a write of several hundred gigabytes repeated across every rank, competing with training for compute and network bandwidth while it happens. Production runs tune the interval against the cluster’s observed failure rate; at $8192$-plus GPUs, a hardware fault every few hours is typical, which pushes checkpoint intervals toward the more frequent end.

Elastic training goes further: rather than stopping the entire job when a GPU fails, the framework detects the failure, drops the affected rank, and continues on the reduced GPU count, or waits for a replacement node to rejoin, without a full restart. PyTorch’s torchelastic and Megatron’s built-in fault-tolerance support are examples. This matters more as cluster size grows, since the expected time between failures shrinks roughly in proportion to GPU count.

None of this is optional at frontier scale. A training run with no fault tolerance is a run that, sooner or later, loses everything to a single failed GPU. The cost of the checkpointing infrastructure is trivial next to the cost of restarting a \50$M run from step zero, which is exactly the arithmetic Section 10 makes concrete.

Cost economics: and what’s next

The cost of frontier training

GPU-hour pricing moves fast and varies enormously by provider; as of 2026-07-13, on-demand rates across cloud providers run roughly \1$to$$3$/hr for A100s,$$2$to$$7$/hr for H100s, and$$5$to$$7$/hr for current-generation B200s (neocloud spot pricing is often$30$–$60%$ below these on-demand figures, and hyperscalers like AWS and Azure typically price several times above specialized GPU-cloud providers for the same hardware). These are directional numbers from a market that resets every few months, not a citable price list, treat any single figure quoted here as stale by the time you read it and check current listings before budgeting a real run. Multiply by the GPU-hours that a frontier training run consumes and the totals get large fast.

Public and rumored estimates: LLaMA-2 70B consumed roughly $1.7$M GPU-hours, around \8$M of compute at$$5$/hr. **GPT-4** is rumored to have used$\sim 30$M GPU-hours and over$$100$M in compute. **Llama-3 405B** used about$31$M H100 GPU-hours for pretraining, around$$100$M–$$150$M. As models scale toward the next compute level ($10^25$to$10^26$FLOPs of training), individual runs approach$$1$B in compute alone, a cost regime where infrastructure efficiency dominates everything else in the training budget.

The pragmatic implication is that **MFU matters more than \/$GPU-hour**. A cluster running at$30%$MFU on$$5$/hr GPUs delivers fewer useful FLOPs per dollar than the same cluster running at$50%$MFU on$$8$/hr GPUs. The arithmetic of a$$50$M training run says that improving MFU from$25%$to$50%$is worth$$25$M, far more than the salaries of the engineers who would do that work. This is why every major lab now has a dedicated training-infrastructure team whose explicit job is MFU optimization.

Exercises

Four exercises that build on the chapter’s machinery. Each is a self-contained problem with a starting template; hints are collapsed by default, try the problem first.

Exercise 1 (easy): NVLink vs InfiniBand bandwidth math

Compute the time for a single gradient all-reduce on a $7$B-parameter model across two scenarios: $8$ GPUs on one node (NVLink) vs $8$ GPUs across two nodes (InfiniBand). Verify the $\sim 36\times$ speed difference.

def all_reduce_time(data_size_bytes, num_gpus, bandwidth_gbps):
    per_gpu_bytes = 2 * (num_gpus - 1) / num_gpus * data_size_bytes
    bandwidth_bytes_s = bandwidth_gbps * 1e9
    return per_gpu_bytes / bandwidth_bytes_s

GRAD_SIZE = 2 * 7e9   # BF16 = 2 bytes/param
NUM_GPUS = 8

nvlink_time = all_reduce_time(GRAD_SIZE, NUM_GPUS, 900)
print(f"NVLink (~900 GB/s):     {nvlink_time*1000:>7.2f} ms")

ib_time = all_reduce_time(GRAD_SIZE, NUM_GPUS, 25)
print(f"InfiniBand (~25 GB/s):  {ib_time*1000:>7.2f} ms")

print(f"\nSpeed gap: {ib_time / nvlink_time:.0f}x slower across nodes")

Exercise 2 (medium): FlashAttention I/O complexity

Compute the HBM I/O for standard attention vs FlashAttention at increasing sequence lengths. Verify that the gap grows asymptotically.

def flashattention_hbm_io(N, d, sram_size=100_000, dtype_bytes=2):
    tiles = max(1, N / sram_size)
    return 4 * N * d * dtype_bytes * tiles

for N in [512, 2048, 8192, 32768, 131072]:
    std = standard_attention_hbm_io(N, 64)
    flash = flashattention_hbm_io(N, 64)
    speedup = std / flash
    print(f"{N:>8} {std/1e6:>12.1f} MB {flash/1e6:>12.1f} MB {speedup:>10.1f}x")

     512          2.4 MB          0.3 MB        9.0x
    2048         34.6 MB          1.0 MB       33.0x
    8192        541.1 MB          4.2 MB      129.0x
   32768       8606.7 MB         16.8 MB      513.0x
  131072     137506.1 MB         88.0 MB     1563.3x

Exercise 3 (medium): Activation memory with checkpointing

Compute activation memory for a $70$B-parameter transformer ($80$ layers, $d_\text{model}=8192$) at various sequence lengths, with and without activation checkpointing. Identify the sequence length where checkpointing becomes necessary on an $80$ GB H100.

def activation_memory(num_layers, batch_size, seq_len, d_model, with_checkpointing, dtype_bytes=2):
    if with_checkpointing:
        return num_layers * batch_size * seq_len * d_model * dtype_bytes
    per_layer = 10 * batch_size * seq_len * d_model * dtype_bytes
    return num_layers * per_layer

NUM_LAYERS, D_MODEL, BATCH_SIZE, H100_MEM_GB, NUM_GPUS = 80, 8192, 4, 80, 32
STATE_GB = (70e9 * 18) / NUM_GPUS / 1e9  # 39.4 GB/GPU

for seq_len in [512, 1024, 2048, 4096, 8192, 16384]:
    mem_no = activation_memory(NUM_LAYERS, BATCH_SIZE, seq_len, D_MODEL, False) / 1e9
    mem_yes = activation_memory(NUM_LAYERS, BATCH_SIZE, seq_len, D_MODEL, True) / 1e9
    available_gb = H100_MEM_GB - STATE_GB
    print(seq_len, mem_no, mem_yes, mem_no < available_gb, mem_yes < available_gb)

 seq_len   no checkpt   with checkpt    fits no/yes
     512      26.8 GB        2.7 GB    yes /    yes
    1024      53.7 GB        5.4 GB     no /    yes
    2048     107.4 GB       10.7 GB     no /    yes
    4096     214.7 GB       21.5 GB     no /    yes
    8192     429.5 GB       42.9 GB     no /     no
   16384     859.0 GB       85.9 GB     no /     no

Exercise 4 (hard): MFU calculation + cost projection

Compute MFU and dollar cost for a realistic training run: Llama-3 $70$B trained on $1.4$T tokens across $1024$ H100s. Compare against industry-typical $40$–$50\%$ MFU.

def estimate_training_run(model_params, tokens, num_gpus, peak_flops_per_gpu, mfu, hourly_cost_per_gpu):
    total_flops = 6 * model_params * tokens
    achievable_flops_per_sec = num_gpus * peak_flops_per_gpu * mfu
    time_hours = total_flops / achievable_flops_per_sec / 3600
    cost = time_hours * num_gpus * hourly_cost_per_gpu
    return {"hours": time_hours, "cost": cost}

def mfu_from_throughput(tokens_per_sec, model_params, num_gpus, peak_flops_per_gpu):
    achieved_flops_per_sec = tokens_per_sec * 6 * model_params
    theoretical_flops_per_sec = num_gpus * peak_flops_per_gpu
    return achieved_flops_per_sec / theoretical_flops_per_sec

MODEL_PARAMS, TOKENS, NUM_GPUS, PEAK_FLOPS_H100 = 70e9, 1.4e12, 1024, 989e12

for mfu in [0.30, 0.40, 0.50]:
    result = estimate_training_run(MODEL_PARAMS, TOKENS, NUM_GPUS, PEAK_FLOPS_H100, mfu, hourly_cost_per_gpu=5.0)
    print(f"MFU={mfu*100:.0f}%: {result['hours']/24:.1f} days, cost ${result['cost']/1e6:.1f}M")

observed_mfu = mfu_from_throughput(1.0e6, MODEL_PARAMS, NUM_GPUS, PEAK_FLOPS_H100)
print(f"Observed MFU from 1M tokens/sec: {observed_mfu*100:.1f}%")

MFU=30%: 22.4 days, cost $2.8M
MFU=40%: 16.8 days, cost $2.1M
MFU=50%: 13.4 days, cost $1.7M
Observed MFU from 1M tokens/sec: 41.5%

Ten chapters in, the entire training-side story is on the page. Architecture (Chapters 1 through 6), data pipeline (Chapter 7), training loop (Chapter 8), scaling math and parallelism strategies (Chapter 9), and the practical infrastructure that runs those strategies on real clusters (Chapter 10). With these, the system diagrams of any frontier training run, Llama-3, GPT-4-class, Gemini, are readable at every level: the architecture, the data mix, the loss and the optimizer, the parallelism layout, the GPU hardware, the interconnect, the framework, the kernels, the cost.

Chapter 11 begins a new arc. The chapters that follow leave the standard dense transformer behind and explore alternative architectures, Mixture of Experts (MoE), state-space models like Mamba, and other designs that trade off the assumptions of the architecture from Chapters 1 through 6 in interesting ways. After that the tutorial moves to post-training (supervised fine-tuning, alignment, parameter-efficient methods), then to inference, capabilities (reasoning, tool use, retrieval, multimodal), safety, evaluation, and agents. The training half of the book is complete; the capabilities half begins.