Part I · Chapter 10

Training at Scale: Stability and Distributed Parallelism

AuthorChangkun Ou

Reading time~20 min

The architectures from the previous chapters only matter if the run can be executed. A frontier model does not fit on one accelerator, and the training job that produces it spans thousands of them for weeks. Once the work is split across that many devices, three resources start to fight: the memory each shard must hold, the communication that keeps the shards consistent, and the failures that a fleet of that size produces as its steady state. Training systems spend those three against each other so achieved throughput stays high and a dead GPU costs minutes, not days. The scoreboard for that work is model FLOPs utilization: the number that says how much of the machine is doing useful model math after communication, bubbles, recomputation, and failures take their cut.

Three forces that fight

Two facts force everything that follows. A frontier model's parameters, gradients, and optimizer state do not fit in one device's memory, and the arithmetic to train it does not finish on one device in any acceptable time. So the work is split across thousands of accelerators, and the moment it is split, three constraints start to fight.

Memory comes first: weights, the Adam moments (the optimizer's per-weight running averages of the gradient) whose state dominates (see Chapter 5), activations (intermediate values saved on the forward pass for use in the backward pass), and gradients must each find a home, and peak activation memory grows with sequence length. Communication is the next pressure, since every split introduces traffic to keep the shards consistent, and that traffic competes with compute for the same time budget. Then there is failure. At thousands of nodes for weeks, a dead GPU or NIC is the steady state, not the exception, and a single one kills the whole collective. The job of the training-systems layer is to spend memory, communication, and recovery against each other so that achieved throughput stays high and a failure costs minutes, not days. No single way of splitting the work resolves all three at once, which is why the design offers several orthogonal cuts and composes them.

Four cuts, and what each one costs

Each axis answers a different scaling problem, shards a different resource, and arrived in the literature when the previous generation hit a wall. Read in that order they are less a menu than a record of how the toolkit was assembled, one cut at a time.

Tensor parallelism (TP) came first. It shards individual matmuls inside a layer, splitting attention heads and the feed-forward columns and rows across devices. Megatron-LM showed how to do this with only two all-reduces per layer, which made a multi-billion-parameter model trainable once it no longer fit on one device (Shoeybi et al. 2019). The cost is that it is communication-heavy: each TP layer needs an all-reduce (or a reduce-scatter and all-gather pair) on the critical path of every forward and backward step. That cost is only affordable over the fastest interconnect, so TP stays inside a node.

Pipeline parallelism (PP) arrived in parallel and cuts a different resource. It splits the layer stack into stages placed on different nodes and streams micro-batches through them, so stage two works on micro-batch one while stage one works on micro-batch two. GPipe established the micro-batching that keeps the stages busy (Huang et al. 2019), and PipeDream generalized the schedule to cut the bubble (Narayanan et al. 2019). PP is bandwidth-cheap, moving only the activations at stage boundaries, but it pays a pipeline bubble: the fill and drain at the ends of each step where some stages sit idle. Figure 10.1 shows where that idle time lives.

Figure 10.1. Pipeline bubble over four stages and four micro-batches: the fill at the start and the drain at the end leave stages idle. After Huang et al., 2019.

Data parallelism (DP) replicates the model and shards the batch across replicas; each replica computes gradients on its slice and an all-reduce averages them. Plain DP does not save memory, since every replica holds a full copy. Zero Redundancy Optimizer (ZeRO) observed that the optimizer state, gradients, and parameters in standard DP are redundantly replicated, and sharded them across the DP group in three progressive stages (Rajbhandari et al. 2019); ZeRO-Infinity extended the idea to offload onto CPU and NVMe (Rajbhandari et al. 2021), and PyTorch fully sharded data parallel (FSDP) is the native version of the same full sharding idea, gathering each layer's weights only when that layer runs (Zhao et al. 2023). Sharding drops per-device memory without changing the math: a layer's parameters are all-gathered just before they are needed and freed right after. This is the sharding Chapter 5 points to when it notes that the optimizer state is the dominant memory term.

Sequence and context parallelism (SP/CP) shard along the sequence dimension, the axis that activation memory lives on. SP splits the parts of a layer that TP leaves replicated, the norms and dropouts, along the sequence to cut activation memory, introduced alongside selective recomputation that redoes only the cheap-to-redo, memory-heavy operations rather than the whole layer (Korthikanti et al. 2022). CP, of which ring attention is the canonical form, shards the attention computation itself so a context too long for one device can be processed in blocks that pass key and value tiles around a ring, letting the context grow nearly without bound (Liu et al. 2023).

For each axis the design question is the same: what does it shard, what does it cost in communication, and which tier of the interconnect can it afford to use. Figure 10.2 lays the answers side by side.

Figure 10.2. The four parallelism axes plus expert parallelism, each by what it shards, the collective it runs, and the interconnect tier it can afford.

Mixture-of-experts layers add a fifth cut. The experts are sharded across devices, and each token is dispatched to its chosen experts by an all-to-all collective, then the results are combined back by a second all-to-all, the dispatch pattern traced to GShard (Lepikhin et al. 2020). The systems layer owns only this execution: the two all-to-all exchanges per MoE layer, the load-imbalance stragglers when a hot expert draws more than its share of tokens, the capacity padding that bounds per-expert work, and the overlap of dispatch and combine with expert compute. The routing decision and the load-balancing loss that create this traffic pattern belong to Chapter 9.

Paying in bits: the precision axis

The arithmetic does not need to run in fp32. Mixed precision keeps the forward and backward matmuls and the activations in bf16, while holding an fp32 master copy of the weights and fp32 optimizer state so that small updates do not vanish when added to a large weight, the recipe established by Micikevicius et al. (Micikevicius et al. 2017). bf16's wide exponent is what makes this safe: it retired the loss scaling (multiplying gradients up so the small ones do not underflow) that fp16 needed to keep gradients off the floor of its narrow range.

fp8 (the E4M3 and E5M2 formats) pushes the matmuls further on Hopper-class and newer hardware, roughly doubling matmul throughput and halving the memory the operands occupy; the formats paper defined E4M3 and E5M2 for the generation that made fp8 matmuls native (Micikevicius et al. 2022). The price is range: fp8 has so few exponent bits that the operands must be rescaled per tensor, or more finely, to keep values inside the representable band. It is applied selectively, the matmuls in fp8 while the numerically sensitive operations, the normalizations, residual adds, and the attention softmax, stay in higher precision. Pushed too far, fp8 does not crash; it quietly costs quality, which is what makes it hard to use well.

That boundary has since moved. DeepSeek-V3 trained at frontier scale with fp8 matmuls in late 2024, which settled that the format can carry a full pre-training run (DeepSeek-AI 2024). The live edge is now four bits. fp4 formats such as NVFP4 and MXFP4 are microscaling formats: values are stored in four bits, and each small block of them, sixteen or thirty-two values, carries its own scale factor to recover locally the range four bits cannot hold, with Blackwell-class hardware running the four-bit matmul natively. NVIDIA has trained a 12B model over ten trillion tokens in NVFP4 to match an fp8 baseline, leaning on Hadamard transforms and stochastic rounding to keep the narrow format unbiased (NVIDIA 2025), and the design space of what makes four-bit training hold together is being mapped systematically (Hu et al. 2025). The hardware side of this bet, why the compute frontier wants fewer operand bytes at all, is Chapter 66's subject.

Figure 10.4 shows why the exponent bits matter. bf16 keeps fp32's eight exponent bits, so it spans the same dynamic range and never needed fp16's loss scaling; fp8 spends its few bits on a tradeoff between range (E5M2) and precision (E4M3). fp4's E2M1 layout is the floor: one mantissa bit, two exponent bits, and the per-block scale factor carrying the rest.

Figure 10.3. Where each format spends its bits. After one sign bit, the rest split between exponent (dynamic range) and mantissa (precision). bf16 keeps fp32's eight exponent bits, so it spans the same range and skips fp16's loss scaling; fp8 has so few exponent bits it must rescale per tensor, and E5M2 buys range while E4M3 buys precision. Pick a format to see the split. Illustrative.

Figure 10.4. Bit layouts of the training precisions. Sign, exponent, and mantissa fields drawn to width. bf16 inherits fp32's eight exponent bits, which is why it spans the same range without loss scaling. fp4's E2M1 keeps almost none, which is why the microscaling formats attach a scale factor to each small block of values.

Hiding the traffic

Every cut above generates traffic, and that traffic runs through a small set of collectives: all-reduce, reduce-scatter, all-gather, and all-to-all. NCCL (RCCL on AMD) implements them and maps each onto the topology, choosing a ring or a tree algorithm and routing over NVLink within a node or the network between nodes (NVIDIA 2024).

The collectives are pure overhead unless they run while the device is doing something else, so overlap is the central throughput technique. FSDP prefetches the next layer's parameters with an all-gather while the current layer computes (Zhao et al. 2023). The gradient reduce-scatter is folded into the backward pass, PP communication into stage compute, and EP all-to-all into the expert matmuls. Whatever cannot be hidden is exposed communication, and it shows up directly as lost throughput.

Model FLOPs utilization is the scoreboard that says whether the overlap worked. It is achieved FLOPs over peak hardware FLOPs, counting only the FLOPs the model mathematically requires. It is distinct from hardware FLOPs utilization (HFU), which also counts the redundant FLOPs of activation recomputation, so HFU exceeds MFU whenever recomputation is on. MFU is the single number that says the systems work is done, and its budget is a sum of losses: exposed communication, the pipeline bubble, recomputation, and unfused kernels each take a slice. A low MFU is diagnosed from a per-step timeline, not from the loss curve. It is the thread that ties the whole teardown together, since every cut and every precision choice ends up scored here.

What's contested

How low training precision can go is unsettled, but the question's frontier has moved. bf16 is the conservative default; fp8 became the workhorse once DeepSeek-V3 carried it through a frontier-scale run (Micikevicius et al. 2022; DeepSeek-AI 2024); and the live boundary is now fp4, where NVFP4-class microscaling and Blackwell's native four-bit matmuls make the throughput case but the recipes are young (NVIDIA 2025; Hu et al. 2025). The disagreement is the same at every rung: which operations can move down, how fine the scaling must be, and how many trillions of tokens a run can sustain before the narrowed range costs final quality. The failure mode is what makes this hard to settle empirically: precision pushed too far does not crash, it quietly degrades the model, so the cost is invisible until a late evaluation. Treat the boundary as a per-operation, per-recipe choice that must be validated, not a setting you can copy.

Constraint arrow

The interconnect dictates the parallelism layout. Tensor parallelism needs an all-reduce on the critical path of every layer, which is only affordable over the fastest links, so the size of the NVLink domain sets the maximum TP group; beyond it, TP traffic crosses the network and stalls (Shoeybi et al. 2019). Pipeline and data parallelism, which move less and can hide it, are the axes that span nodes, and expert parallelism gets its own all-to-all group. The bandwidth hierarchy of Chapter 62 is therefore upstream of every mesh decision in this chapter: a lower layer's wires decide which split an upper layer may use where.

Turning the dials

Each axis spends a different resource, and the layout that wins is the one that hides the most communication under compute for a specific model shape on a specific interconnect. By 2021 the axes were being combined deliberately: the PTD-P work mapped data, tensor, and pipeline parallelism onto a single GPU cluster as a composed 3D layout, establishing the ordering that frontier runs still follow (Narayanan et al. 2021). The dials below are what an engineer actually turns inside that layout.

Parallelism layout, and its non-portability. TP cuts per-device memory and latency but burns intra-node bandwidth, so it caps at the NVLink domain. PP is bandwidth-cheap but pays a pipeline bubble, mitigated by more micro-batches and interleaved schedules (Narayanan et al. 2019), with diminishing returns once the bubble is already small (Figure 10.5). DP and ZeRO scale the batch, but the optimizer-state shard and the all-gather traffic grow with the DP degree (Rajbhandari et al. 2019). The right layout is not portable: change the model size, the sequence length, or the cluster, and the optimum moves.

Figure 10.5. Schematic illustration of how the pipeline bubble fraction shrinks as the number of micro-batches per step grows, for several stage counts, after Huang et al. (2019) and Narayanan et al. (2021). Idealized fraction (p-1)/(m+p-1); not measured data. The gain is steep at first and flattens, so beyond a handful of micro-batches more buys little.

Run the exact formula and watch the diminishing returns: change the stage counts or the range of micro-batches and see where the curve flattens.

import numpy as np
import matplotlib.pyplot as plt

m = np.arange(1, 33)  # micro-batches per step
for p in [4, 8, 16, 32]:  # pipeline stages
    bubble = (p - 1) / (m + p - 1)
    plt.plot(m, bubble, marker="o", markersize=3, label=f"p={p} stages")

print("p=8 stages:", {int(mm): round((8 - 1) / (mm + 8 - 1), 3) for mm in [1, 4, 8, 16, 32]})
plt.xlabel("micro-batches per step (m)")
plt.ylabel("pipeline bubble fraction")
plt.title("Bubble fraction (p-1)/(m+p-1)")
plt.legend()
plt.grid(alpha=0.3)
plt.show()

Figure 10.6. The pipeline bubble. Stages process micro-batches in a wavefront, idle during the fill and drain at each end, so the bubble fraction is (p-1)/(m+p-1). Add stages and the idle wedge grows; add micro-batches per step and it shrinks, steeply at first then flattening, which is why a handful of micro-batches is usually enough. Illustrative.

Memory versus compute. Activation and gradient checkpointing trade extra forward FLOPs in the backward pass for lower peak activation memory. Selective recomputation redoes only the cheap, memory-heavy operations (Korthikanti et al. 2022). More recomputation lets you fit a bigger model or a longer sequence, at the cost of MFU, which is exactly why HFU exceeds MFU when recompute is on.
Precision versus stability. bf16 is the safe default, and fp32 master weights are non-negotiable for the optimizer (Micikevicius et al. 2017). fp8 is the proven workhorse: it buys throughput and memory but narrows range, so it is applied with per-tensor or finer scaling and kept away from the sensitive operations (Micikevicius et al. 2022; DeepSeek-AI 2024). fp4 is the live boundary, affordable only with per-block scaling and the newest hardware (NVIDIA 2025). Pushed too far, any of them produces silent quality loss, not a crash.
Checkpoint cadence. Frequent checkpoints shrink the work lost per failure but tax storage bandwidth and can stall the run. Asynchronous and sharded checkpointing is what makes a short cadence affordable, and the optimum depends on cluster mean time between failures and write cost (Mohan et al. 2021).

Laying out a real run, and how it breaks

Composing the axes

A frontier run maps the device mesh as a product of axes, conventionally (DP x TP x PP), adding EP for MoE and layering SP on top of TP. The ordering rule of thumb follows directly from the constraint arrow and from the PTD-P result: put TP inside a node under NVLink, run PP and DP across nodes where their cheaper communication can hide, and give EP its own all-to-all group (Narayanan et al. 2021). The layout is co-designed with the model shape from Chapter 8 and Chapter 9, and it shifts with model size, sequence length, and interconnect.

Figure 10.7 sketches the conventional placement.

Figure 10.7. The conventional device-mesh placement: communication-heavy TP and SP stay inside a node under NVLink, while the cheaper-to-hide DP, PP, and EP span nodes.

Frameworks

The frameworks differ mostly in which layout and which hardware each makes ergonomic. Megatron-LM and Megatron-Core carry TP, PP, SP, and EP on the NVIDIA stack (Shoeybi et al. 2019). DeepSpeed packages ZeRO and offload (Rajbhandari et al. 2019). PyTorch FSDP is native sharded DP (Zhao et al. 2023). JAX and XLA express sharding as GSPMD-style annotations on TPU, where the JAX pjit interface drives the partitioner (Xu et al. 2021). In practice a frontier stack combines pieces, for example Megatron-Core for model parallelism plus a custom data and checkpoint plane.

Fault tolerance

At thousands of nodes for weeks, recovery is a first-class subsystem. Checkpointing cadence is set by the write cost against the work lost per failure; asynchronous and sharded checkpointing, as in CheckFreq, keep the write off the critical path so a short cadence stays affordable (Mohan et al. 2021). Elastic restart detects a dead node, fences it, and resumes from the last checkpoint, ideally onto spare capacity without a full restart. Stragglers and silent data corruption are harder than clean crashes, because the run keeps going while it is slow or subtly wrong.

One fault-tolerance concern reaches into the data plane: a restart must resume the exact same data order, which means the sampler and shuffle state are part of the checkpoint. Without that, a resume re-feeds or skips data and breaks the mixture contract from Chapter 6, surfacing as an unexplained loss discontinuity at every restart. The storage fabric that streams the corpus and the cluster orchestration that schedules the run live in Chapter 65; the resumability contract is the piece that stays here.

Failure modes

The symptoms cluster. A low MFU is the catch-all, traced by a per-step profile to exposed communication, an oversized bubble, too much recomputation, or unfused kernels. When a communication bottleneck dominates, the network is the limiter: TP spilling past the NVLink domain, a DP all-gather that cannot hide under the backward pass, or MoE all-to-all stalling on a hot expert. An out-of-memory error points instead to the wrong layout for the budget: activation memory at long sequence length, optimizer state under-sharded, or a fragmentation cliff that appears only at a certain micro-batch count. And a node failure with no fast detection, fencing, and recent checkpoint rolls the run back to its last save, where the gap is pure waste.

The fastest attention kernel does not change this accounting: FlashAttention keeps attention exact and makes the score matrix affordable by never writing it to high-bandwidth memory, a kernel-level win that sits underneath these parallelism axes rather than replacing one of them (Dao et al. 2022). Figure 10.8 contrasts the two scaling regimes: a materialized score matrix whose memory grows quadratically with sequence length against the linear growth of a kernel that streams the scores.

Figure 10.8. Schematic illustration of peak attention memory versus sequence length, after Dao et al. (2022). A materialized score matrix grows as O(L²) while an IO-aware kernel that never writes the scores to HBM grows as O(L). Idealized, arbitrary units; not measured data.

Drag the exponent to feel the gap between the two regimes: at exponent 2 the memory is the materialized O(L²) score matrix, at exponent 1 it is the FlashAttention O(L) stream.

Figure 10.9. Peak attention memory versus sequence length: exponent 2 is the materialized O(L²) score matrix, exponent 1 is the FlashAttention O(L) stream.

Recipe-level stability, the z-loss and QK-norm and loss-spike recovery that keep the optimizer from diverging, belongs to Chapter 5. The stability this chapter owns is numerical, set by the precision choice, and run-level, set by fault tolerance.