Training at Scale: Stability and Distributed Parallelism

A training run becomes distributed for one of two reasons: the model state no longer fits on one accelerator, or one accelerator would take too long to finish the work. Distribution solves those problems by creating another one. Every shard must exchange data with its peers, and every participating worker becomes another place where the job can stop.

The design task is therefore an accounting problem. For each proposed layout, ask four questions:

1. What is replicated, and what is sharded?
2. What is the peak memory on each device, including temporary buffers?
3. Which messages lie on the critical path of a step?
4. What state must survive a restart for the run to continue correctly?

This chapter develops that accounting from a single-device training step to a multi-dimensional device mesh. The goal is not a universal recipe. Model shape, sequence length, numerical format, and network topology jointly determine the best layout.

Start with one logical optimizer step

Data parallelism is the simplest place to begin. Each of $D$ replicas receives a different part of the batch, runs the same model, and contributes to one synchronous gradient update. If every replica processes $m$ micro-batches of $B_\mu$ sequences before the optimizer runs, then

$$B_{\mathrm{global}} = DmB_\mu.$$

Here, $B_{\mathrm{global}}$ is the number of sequences in one optimizer step, $D$ is the data-parallel degree, $m$ is the number of micro-batches accumulated per replica, and $B_\mu$ is the number of sequences in one micro-batch. Token counts can differ between sequences, so production runs often track tokens per step as well as sequences per step.

With equal local batch sizes, an all-reduce produces the mean gradient

$$g = \frac{1}{D}\sum_{r=1}^{D}g_r, \qquad g_r = \frac{1}{B_{\mathrm{global}}/D} \sum_{i\in\mathcal B_r}\nabla_\theta \ell_i(\theta).$$

Here, $r$ indexes replicas, $\mathcal B_r$ is replica $r$'s local batch, $\theta$ is the parameter vector, $\ell_i$ is the loss for example $i$, and $g_r$ and $g$ are the local and global mean gradients. If local batch sizes differ, averaging the replica means equally is wrong; the reduction must weight each sum by its example or token count.

This equation is the semantic contract. Different sharding schemes may change where tensors live and the order of floating-point reductions, but they should implement the same intended synchronous update. Bitwise identity is a stronger and usually topology-dependent requirement.

Account for model-state memory

Plain distributed data parallelism (DDP) replicates parameters, gradients, and optimizer state on every replica. It distributes compute and the batch, but it does not shard this persistent model state. Zero Redundancy Optimizer (ZeRO) removes the replication in stages (Rajbhandari et al. 2020):

Here, $P_\theta$ is the parameter count; $b_w$, $b_g$, and $b_o$ are the bytes per parameter used by working parameters, gradients, and all optimizer-related state; and $D$ is the data-parallel degree. For one common Adam recipe, $b_w=2$, $b_g=2$, and $b_o=12$: a two-byte working parameter and gradient, plus a four-byte master parameter and two four-byte moments. That recipe uses $16P_\theta$ bytes before activations and temporary storage, but these dtypes are choices rather than constants of Adam.

The table is a persistent-state estimate, not a peak-memory estimate. Under full sharding, a wrapped module's parameters are gathered before computation and gradients are reduced and scattered afterward. A useful lower bound is

$$M_{\mathrm{peak}} \gtrsim \frac{P_\theta(b_w+b_g+b_o)}{D} + P_u b_w + M_{\mathrm{act}} + M_{\mathrm{tmp}}.$$

Here, $P_u$ is the number of parameters simultaneously materialized by the current and prefetched units, $M_{\mathrm{act}}$ is saved activation memory, and $M_{\mathrm{tmp}}$ covers communication workspaces, allocator overhead, padding, and other temporary buffers. PyTorch fully sharded data parallel (FSDP) applies this full sharding idea to wrapped modules and can prefetch a later unit's parameters while the current one computes (Zhao et al. 2023). ZeRO-Infinity extends the memory hierarchy to CPU memory and NVMe storage (Rajbhandari et al. 2021).

More sharding does not make communication disappear. In the element-volume model used by the ZeRO paper, DDP and stages 1 and 2 move roughly two model copies per step, while stage 3 moves roughly three because parameters are gathered for both forward and backward computation (Rajbhandari et al. 2020). Per-device state falls as $1/D$; parameter-gather payload per device approaches one model copy as $D$ grows, while latency and topology can still make a larger group slower.

Four more ways to divide the work

Data parallelism is only one axis. A model that still does not fit, or whose layers do not run efficiently on one device, needs other cuts.

Tensor parallelism: split a layer

Tensor parallelism (TP) divides matrix multiplications within a layer. For a feed-forward block with input $X$, the paired column and row partition is

$$H_r=\phi(XW_1^{(r)}), \qquad Z=\sum_{r=1}^{T}H_rW_2^{(r)}.$$

Here, $T$ is the tensor-parallel degree; $r$ indexes its ranks; $W_1^{(r)}$ is a column shard of the first projection; $W_2^{(r)}$ is the matching row shard of the second projection; $\phi$ is the activation function; and the sum that forms $Z$ requires a reduction. Attention can similarly shard query, key, and value heads, then row-shard the output projection.

Megatron-LM's original layout arranges the attention and feed-forward blocks so that each Transformer layer has two all-reduces in its forward pass and the corresponding reductions in backward (Shoeybi et al. 2019). TP lowers per-rank parameters and arithmetic, but places communication on the critical path of every layer. It is therefore usually mapped to the fastest fabric domain available. That is a placement rule, not a law: whether a TP group may cross nodes depends on message size, topology, and the performance target.

Pipeline parallelism: split the layer stack

Pipeline parallelism (PP) assigns consecutive groups of layers to stages and streams micro-batches through them. Stage boundaries carry activations forward and their gradients backward. The stages need not coincide with nodes.

At the start and end of a synchronous flush schedule, some stages are idle. For balanced stages with negligible communication, GPipe's idealized idle fraction is (Huang et al. 2019)

$$f_{\mathrm{bubble}}=\frac{p-1}{m+p-1}.$$

Here, $p$ is the number of pipeline stages, $m$ is the number of micro-batches per replica per optimizer step, and $f_{\mathrm{bubble}}$ is the fraction of aggregate stage slots left idle by fill and drain. This denominator matters: the often-quoted $(p-1)/m$ is bubble time relative to ideal compute time, not the fraction of total stage capacity.

The formula is useful for sizing, not prediction. Unequal stage times, communication, optimizer work, and alternative schedules change the result. A 1F1B schedule can reduce saved activation memory without changing this ideal bubble; interleaved schedules can reduce the bubble but send more point-to-point messages (Narayanan et al. 2021). PipeDream explores asynchronous schedules and weight versions, a different optimizer-semantics tradeoff from GPipe's synchronous flush (Narayanan et al. 2019).

The calculation below makes the scale of the bubble concrete. For 32 stages and 32 micro-batches, nearly half of aggregate stage capacity is still idle in this idealized schedule.

def bubble_fraction(stages, microbatches):
    """Balanced GPipe flush schedule; communication is ignored."""
    if stages < 1 or microbatches < 1:
        raise ValueError("stages and microbatches must be positive")
    return (stages - 1) / (microbatches + stages - 1)


for stages in (4, 8, 16, 32):
    values = {
        microbatches: round(bubble_fraction(stages, microbatches), 3)
        for microbatches in (1, 4, 8, 16, 32, 64)
    }
    print(f"p={stages}: {values}")

print("p=32, m=32:", f"{bubble_fraction(32, 32):.1%}")

Sequence parallelism and context parallelism: split tokens

These names describe related but distinct mechanisms. Megatron-style sequence parallelism (SP) shards the layer-normalization, dropout, and other regions that TP otherwise replicates along the sequence dimension. It reuses the TP group and replaces selected all-reduces with reduce-scatter and all-gather pairs; it is not usually another multiplier in the device mesh (Korthikanti et al. 2022). Selective activation recomputation complements SP by recomputing operations that are cheap relative to the memory they save.

Context parallelism (CP) divides the sequence through attention itself. Each rank owns local queries and passes or gathers key/value blocks needed to compute exact attention. Ring Attention streams those blocks around a ring while using an online softmax (Liu et al. 2023); all-to-all designs are another option. CP reduces local sequence storage, but dense attention still performs $O(L^2)$ global arithmetic for context length $L$. Maximum context remains bounded by aggregate memory, communication, device count, and the degree to which communication can overlap computation.

Expert parallelism: split conditional parameters

Mixture-of-experts models add expert parallelism (EP). Tokens are routed to the ranks that hold their selected experts and the outputs are returned to the original token owners, commonly with all-to-all communication (Lepikhin et al. 2020). A forward pass often has a dispatch and a combine exchange; backward sends the corresponding gradient traffic. Hot experts cause stragglers, so routing locality, load balance, optional capacity padding, and dropless kernels are systems concerns as well as modeling choices. The routing objective itself is covered in Chapter 9.

Figure 10.4 summarizes what each axis moves. The collective names are common implementations, not promises: frameworks can fuse, decompose, or schedule them differently.

Map the cuts onto the network

For independent data, tensor, pipeline, and context dimensions, a dense-model mesh satisfies

$$N = D T p C.$$

Here, $N$ is the world size, $D$ the data-parallel degree, $T$ the tensor-parallel degree, $p$ the number of pipeline stages, and $C$ the context-parallel degree. SP normally shares the $T$ group. EP is model-dependent: its group may be carved from or overlap another mesh dimension, so blindly multiplying by an expert degree can double-count devices.

The placement principle is simple: put the most frequent and latency-sensitive exchanges inside the fastest fabric domain. TP often gets that domain because it communicates inside every layer. EP needs strong bisection bandwidth for all-to-all traffic. PP message size depends on the boundary activation, while ZeRO-3/FSDP repeatedly gathers parameters. Which axis crosses a node or rack boundary is therefore a measurement-driven choice, not a fixed hierarchy (Narayanan et al. 2021).

Collectives have different semantics. An all-reduce reduces values and gives every rank the result. A reduce-scatter reduces and leaves each rank one shard. An all-gather concatenates shards on every rank. An all-to-all sends a distinct shard to each peer. PP mostly uses point-to-point send and receive. Libraries such as NCCL implement several algorithms for these operations and select among them according to the hardware and message shape (NVIDIA 2024).

Communication is necessary work, but only its exposed part extends the step's critical path. FSDP can gather the next unit while the current unit computes; gradient buckets can reduce during backward; PP and EP messages can overlap nearby kernels. Overlap is not free: communication may contend with compute for memory bandwidth, execution resources, or the network. A step trace, rather than a list of enabled flags, shows whether the overlap worked.

Choose numerical formats per operation

Mixed precision is not one global dtype. A common recipe performs large matrix multiplications in a low-precision format, accumulates reductions more carefully, and keeps sensitive operations or optimizer state in a wider format. The original mixed-FP16 recipe used FP32 master parameters, FP32 accumulation, and loss scaling (Micikevicius et al. 2017). BF16 has eight exponent bits and therefore roughly FP32's normal range, so it usually avoids FP16's loss scaling; its seven fraction bits still make accumulation and small updates precision problems.

FP8 offers two common layouts: E4M3 spends more bits on precision, while E5M2 spends more on range (Micikevicius et al. 2022). Values must be scaled into the representable interval, and the scaling granularity, whether tensor, tile, or block, is part of the recipe. Hopper-class hardware can have twice BF16's peak dense matmul rate for FP8 and FP8 operands use half as many payload bytes, but casts, scale metadata, non-matmul work, and retained high-precision state reduce the end-to-end gain.

DeepSeek-V3 demonstrated a full pre-training run whose core matrix multiplications used FP8, while several sensitive operations and outputs stayed in BF16 or FP32 (DeepSeek-AI 2024). A later NVIDIA experiment trained a 12B model for ten trillion tokens with NVFP4 and reported results comparable to an FP8 baseline, using fine-grained scaling, Hadamard transforms, stochastic rounding, and selective higher precision (NVIDIA 2025). These are measured recipes, not evidence that every model or operation can be lowered safely.

A robust update has an explicit order: accumulate all micro-batches; unscale FP16 gradients if loss scaling is active; check finite values across all ranks; compute any global gradient norm over the logical, fully sharded gradient; clip if configured; then update every shard or skip the update everywhere. Scaling state, including FP8 amax history or a dynamic loss scaler, belongs in the checkpoint. Narrow formats can cause underflow, saturation, Inf/NaN values, skipped steps, divergence, or quiet quality drift. Validation must look for all of them.

Measure throughput without hiding assumptions

Model FLOPs utilization (MFU) compares observed model work with a stated hardware peak (Chowdhery et al. 2023):

$$\mathrm{MFU}= \frac{r_{\mathrm{tok}}F_{\mathrm{model,tok}}} {N P_{\mathrm{peak}}}.$$

Here, $r_{\mathrm{tok}}$ is aggregate observed tokens per second, $F_{\mathrm{model,tok}}$ is the forward-plus-backward FLOP count required by the model per token, $N$ is the accelerator count, and $P_{\mathrm{peak}}$ is the matching-precision peak FLOP rate of one accelerator. The report must state the FLOP convention, how active MoE parameters and padding are treated, the precision peak used, and the measurement window.

Hardware FLOPs utilization (HFU) instead counts executed FLOPs, including activation recomputation. HFU can therefore exceed MFU even when useful-token throughput is unchanged. Neither metric proves training quality or reliability. Steady-state MFU commonly excludes restart time; end-to-end goodput should also charge failed steps, checkpoints, evaluation pauses, and input stalls.

When MFU is low, inspect a trace before changing the mesh. The usual causes are small or poorly shaped kernels, stage imbalance, exposed communication, excessive recomputation, input stalls, and synchronization behind a straggler. Peak memory, collective duration by group, tokens per second, step-time tails, and time lost to recovery complete the operational picture.

Make restart a correctness property

In an ordinary synchronous job, one failed rank invalidates its communicators. Elastic launchers generally stop and rebuild the worker group, then load a committed checkpoint; they do not hot-swap one device into a running collective. The control sequence is: detect the lack of progress, abort communication, fence or quarantine the suspect worker, replace capacity, rendezvous, rebuild groups, load a validated checkpoint, verify the data cursor, and resume.

A recoverable checkpoint is an atomic set, not merely a parameter file. It contains model and any master parameters; optimizer tensors and step; learning-rate state; loss-scaler and low-precision scale history; per-rank RNG state; sampler, shuffle, data-mixture counters, and cursor; plus a shard manifest, checksums, and a completion marker. Saving at an optimizer-step boundary avoids storing partially accumulated gradients. Keep more than the latest checkpoint because a numerically damaged state may already have been saved.

There are two legitimate resume contracts. Deterministic replay attempts the same batches and random draws and usually requires the same topology plus deterministic kernels. Coverage-equivalent resume may reshard onto a different world size while preserving which data is consumed overall. It need not be bitwise identical. State the contract and test it.

Checkpoint cadence trades write cost against repeated work. A first-order fail-stop model gives

$$W(\tau)\approx\frac{C}{\tau}+\frac{\tau}{2M}, \qquad \tau^\star\approx\sqrt{2CM}.$$

Here, $W(\tau)$ is the expected fraction of time lost to checkpointing and recomputation, $\tau$ is useful compute time between checkpoints, $C$ is the exposed checkpoint cost, $M$ is the synchronized job's mean time between interruptions, and $\tau^\star$ is the approximate minimizing interval (Daly 2006). The approximation assumes independent fail-stop events and omits correlated outages, recovery delay, and storage contention. Asynchronous checkpointing lowers exposed $C$, but still consumes device links, host memory, network capacity, and storage bandwidth. CheckFreq demonstrates adaptive cadence and pipelined snapshot/persist phases rather than a free write (Mohan et al. 2021).

Verify the layout before a long run

A distributed configuration is ready only after it passes small, deliberate tests:

• Update equivalence: on a tiny deterministic batch, compare the unsharded and sharded loss, gradients, and one optimizer update within a stated tolerance.
• Memory accounting: measure persistent, peak, and temporary memory while sweeping micro-batch and sequence length. Confirm that the wrapped-unit and prefetch assumptions match the trace.
• Communication accounting: profile each group and identify exposed rather than merely total communication. Test the planned cross-node boundaries.
• Numerical stability: compare a higher-precision control with the proposed per-operation recipe; record nonfinite tensors, skipped steps, pre-clip gradient norms, clipped-step frequency, and evaluation drift.
• Restart correctness: kill a worker during compute and during checkpoint persistence. Reject incomplete manifests, restore the declared data contract, and compare the resumed trajectory with the control.
• Goodput: run long enough to include checkpoints, input loading, and at least one recovery. Report useful tokens per wall-clock second alongside MFU.

FlashAttention illustrates why every layer of this accounting matters. It computes exact attention without materializing the full score matrix in HBM (Dao et al. 2022). For batch size $B$, head count $H$, context length $L$, and $b$ bytes per score element, materializing the scores alone would use

$$M_{\mathrm{score}} = BHL^2b.$$

Here, $M_{\mathrm{score}}$ excludes softmax workspace, inputs, outputs, saved activations, and allocator overhead. FlashAttention removes this quadratic HBM tensor, but it does not remove the dense attention arithmetic or the need to shard a context that exceeds aggregate memory.

The same discipline applies to the whole run: name the tensor, count its bytes, identify when it exists, and state which communication or recovery step follows. That turns a collection of parallelism acronyms into an engineering design that can be checked.

Further reading