Chapter 4.4 β€” Distributed Training

Contents

  1. From "what to train on" to "how to physically train it"
  2. Why one accelerator is never enough at frontier scale
  3. Data parallelism and its limit
  4. Tensor parallelism: sharding the layers themselves
  5. Pipeline parallelism: sharding the layers sequentially
  6. ZeRO and FSDP: sharding the optimizer state instead of the model
  7. Activation memory: checkpointing, sequence parallelism, and overlap
  8. Combining strategies: 3D and 4D parallelism
  9. Interview angle
  10. Self-check questions
  11. Sources

1. From "what to train on" to "how to physically train it"

Chapter 4.3 gave you a way to decide, quantitatively, how large a model to train and how many tokens to feed it for a given compute budget, and closed by pointing out that it said nothing about how you actually get such a model onto hardware and trained in finite time. That gap is not a minor implementation detail β€” at the parameter and token counts scaling laws now routinely recommend, the naive approach of "put the whole model on one device and run gradient descent" is not slow, it is simply impossible, for reasons worth stating precisely rather than waving at. This chapter is about the systems engineering that makes frontier-scale training runs physically realizable at all, independent of any of the algorithmic choices covered so far.

2. Why one accelerator is never enough at frontier scale

There are two separate reasons a single accelerator cannot train a frontier-scale model, and it's worth keeping them distinct because they call for different solutions. The first is memory: training requires storing not just the model's parameters, but also their gradients (the same size as the parameters) and, for adaptive optimizers like Adam, per-parameter optimizer state (typically two additional tensors the same size as the parameters, for the running first and second moments), plus activations retained for the backward pass. Even before accounting for activations, parameters plus gradients plus optimizer state for a model with tens or hundreds of billions of parameters, stored in the precisions typically used for stable training, simply exceeds the memory of any single accelerator available today by a wide margin β€” this is a hard constraint, not a performance inconvenience, and no amount of patience fixes it. Concretely, for mixed-precision training with Adam, Rajbhandari et al.'s ZeRO paper works out that these "model states" cost roughly $16\,\Psi$ bytes per parameter, where $\Psi$ is the parameter count: 2 bytes for the fp16 parameter plus 2 bytes for the fp16 gradient plus 4+4+4 bytes for the fp32 master copy and the two Adam moment estimates. For a 70-billion-parameter model that's over 1 TB of model states alone, before a single activation is stored β€” which is exactly why no single accelerator comes close to fitting a frontier-scale model. The second reason is throughput: even in the hypothetical world where everything fit in memory, training on the token counts scaling laws now recommend on a single device would take a length of time measured in years or decades, which is operationally useless regardless of memory. Distributed training exists to solve both problems simultaneously β€” spreading memory usage across many devices, and spreading compute across many devices to bring wall-clock time down to something practical.

3. Data parallelism and its limit

The simplest and oldest distributed training strategy is data parallelism: replicate the entire model on every device, split each training batch into disjoint shards, have each device compute a forward and backward pass on its own shard independently, and then average the resulting gradients across all devices before applying the optimizer step, so that every replica stays synchronized and effectively behaves as one large-batch training run distributed over many devices. This is attractive because it requires essentially no change to the model itself β€” every device holds an identical, complete copy of the model β€” and the only new engineering burden is the gradient-averaging communication step (an all-reduce across devices) after each backward pass.

Its limit is exactly the memory problem described above: pure data parallelism requires every single device to hold the full model, its full gradients, and its full optimizer state, which means data parallelism alone does nothing to help if the model doesn't fit on one device in the first place. It scales throughput β€” more devices process more of the batch in parallel β€” but it does not scale the amount of model you can train, which means at frontier scale it has to be combined with a strategy that actually shards the model itself across devices, not merely the data.

4. Tensor parallelism: sharding the layers themselves

Tensor parallelism, sometimes called model parallelism, addresses the memory problem head-on by splitting individual weight matrices β€” and the layers built from them β€” across devices, so that no single device ever needs to hold a complete copy of any given layer. Megatron-LM, Shoeybi et al.'s approach to this problem, works out concretely how to split the two dominant compute-heavy components of a transformer layer, the attention block and the feedforward block, across devices in a way that requires only a small, well-defined amount of communication per layer rather than communication after every individual matrix multiplication. For the feedforward block, this typically means splitting the first linear layer's weight matrix by columns across devices (so each device computes a different slice of the intermediate activation independently, with no communication needed for that step) and splitting the second linear layer's weight matrix by rows (so that summing the partial outputs from all devices, a single communication step, reconstructs the correct final output). Attention is split analogously by dividing the attention heads themselves across devices, since each head's computation is independent of the others until the final output projection.

The central tradeoff here is communication cost: tensor parallelism requires devices to synchronize (via an all-reduce, typically) at specific points within every single layer, not just once per batch as data parallelism does, which means it demands very high-bandwidth, low-latency interconnects between the devices involved β€” it works well within a single high-bandwidth server or a small cluster of them, but degrades quickly if you try to stretch it across devices connected by slower networking, because the frequent per-layer communication becomes the bottleneck rather than the compute itself. This is why tensor parallelism is typically used at a relatively small "width" β€” splitting across a handful to a few dozen devices within a tightly-networked node or rack β€” rather than as the sole strategy for spreading a model across an entire large cluster.

5. Pipeline parallelism: sharding the layers sequentially

Pipeline parallelism takes a different cut at the same memory problem: instead of splitting individual layers across devices, it splits the model's layers into contiguous groups (stages) and assigns each stage to a different device, so that data flows through the devices sequentially, the way an assembly line does β€” device one computes the first several layers and passes its output to device two, which computes the next several layers, and so on. This requires far less communication per step than tensor parallelism, since only the activations at stage boundaries need to be transmitted, but it introduces a different problem: naively, only one device is doing useful work at any given moment while the others sit idle waiting for their input to arrive, which defeats the purpose of parallelizing at all.

The standard fix is micro-batching: split each training batch into several smaller micro-batches and feed them into the pipeline one after another in a staggered schedule, so that while device two is processing the first micro-batch's later stage, device one is already processing the second micro-batch's earlier stage, keeping every device busy simultaneously with different micro-batches at different points in the pipeline. There remains an unavoidable startup and drain period β€” the "bubble" β€” during which the pipeline is filling up or emptying out and not every device is fully utilized, and a meaningful part of pipeline-parallel engineering is choosing the number and size of micro-batches, and the specific scheduling of forward and backward passes across them, to minimize the fraction of total training time lost to this bubble.

6. ZeRO and FSDP: sharding the optimizer state instead of the model

ZeRO, the Zero Redundancy Optimizer described by Rajbhandari et al., takes a genuinely orthogonal approach to the ones above: rather than splitting the model's layers or weight matrices across devices at all, it keeps data parallelism's simple structure β€” every device processes a different data shard β€” but eliminates the redundancy of every device storing an identical, complete copy of the optimizer states, gradients, and (in its most aggressive stage) the parameters themselves. Ordinary data parallelism duplicates all of this on every replica even though, at any given moment, each device only strictly needs the portion of it relevant to the specific computation it's currently doing; ZeRO instead shards these tensors across the data-parallel group, with each device holding only a fraction of the optimizer state, gradients, or parameters, and dynamically gathering the pieces it needs from other devices via communication exactly when a given computation requires the full tensor, then releasing them again afterward. This dramatically cuts per-device memory relative to plain data parallelism β€” often by close to a factor equal to the number of data-parallel devices, depending on which ZeRO stage is used: at ZeRO's most aggressive stage, per-device memory for model states falls to roughly $\frac{16\,\Psi}{N_{dp}}$ bytes, where $N_{dp}$ is the number of data-parallel devices, down from the full $16\,\Psi$ every device pays under plain data parallelism β€” without requiring the model itself to be architecturally split the way tensor or pipeline parallelism demands, which makes it comparatively simple to apply to an existing model implementation.

FSDP, Fully Sharded Data Parallel, is the implementation of essentially this same idea that became the standard, widely-used realization of ZeRO-style sharding in practice, integrated into mainstream training frameworks rather than living only in a separate specialized library. The conceptual content is the same as ZeRO's most aggressive stage β€” shard parameters, gradients, and optimizer state across data-parallel workers, gathering full parameter shards just-in-time for each layer's computation and freeing them immediately afterward β€” but it's worth knowing FSDP by name specifically, because it is the term most commonly used in current training-framework documentation and in interview conversations about this exact technique.

7. Activation memory: checkpointing, sequence parallelism, and overlap

Everything in sections 3 through 6 shrinks the memory taken up by parameters, gradients, and optimizer state. There is a fourth consumer of device memory that none of them touch: activations, the intermediate outputs of every layer computed during the forward pass and kept around because the backward pass needs them to compute gradients. At the batch sizes and sequence lengths frontier training actually uses, activation memory is not a rounding error β€” for a sufficiently deep model on long sequences, it can exceed the memory taken up by the parameters themselves, and it is the reason a configuration that comfortably fits its weights and optimizer state under ZeRO/FSDP can still run out of memory the moment you increase sequence length or batch size.

Activation checkpointing (also called gradient checkpointing, following Chen et al.'s formulation) trades compute for this memory directly: instead of retaining every layer's activations through the whole forward pass, keep only a sparse set of "checkpoint" activations at chosen layers, discard the rest, and during the backward pass, recompute the discarded activations on demand by re-running the forward computation from the nearest retained checkpoint. Chen et al. showed this lets memory scale as roughly $O(\sqrt{L})$ instead of $O(L)$, where $L$ is the number of layers, at the cost of one extra forward pass through the recomputed segments β€” commonly quoted at around 30% additional compute for a substantial reduction in peak activation memory. This is a textbook example of a tradeoff this book keeps returning to: spend more FLOPs, which are comparatively abundant, to save memory, which at frontier scale is the tighter constraint.

Sequence parallelism, described by Korthikanti et al., attacks a specific gap that tensor parallelism (section 4) leaves open: operations like LayerNorm and dropout aren't split by Megatron-style tensor parallelism, because they operate independently per token rather than mixing information across the hidden dimension the way attention and the feedforward block do β€” so every tensor-parallel device ends up redundantly storing the full activation for those layers regardless of how many ways the matrix multiplications themselves are split. Sequence parallelism shards those specific activations along the sequence dimension instead, across the same tensor-parallel group that's already communicating for the split matrix multiplications, so the redundant copies disappear without introducing any new communication pattern beyond what tensor parallelism already pays for. Korthikanti et al. show that combining sequence parallelism with a more selective form of activation checkpointing β€” recomputing only the specific operations that are cheap to redo rather than checkpointing whole transformer blocks β€” cuts activation memory substantially further than either technique alone, while adding markedly less recomputation overhead than checkpointing every block.

A final, purely systems-level lever is worth naming because it costs nothing algorithmically: overlapping communication with compute. Every strategy in this chapter introduces some communication β€” gradient all-reduces in data parallelism, activation and gradient exchange at tensor- and pipeline-parallelism boundaries, parameter gathering in FSDP β€” and a naive implementation simply waits for each communication step to finish before continuing, leaving the accelerator idle while data moves over the network. Because the communication for one part of the model (say, a layer whose backward pass just finished) doesn't depend on the compute for another part (the next layer's backward pass, already ready to start), a framework can issue the communication asynchronously and let it proceed in the background while the accelerator keeps computing, effectively hiding communication latency behind useful work whenever the two don't have a genuine data dependency. Modern training frameworks do this by default wherever the dependency structure allows it, and it is one of the reasons well-engineered 3D/4D parallelism achieves throughput much closer to the theoretical compute-bound limit than a naive, unoverlapped implementation of the same strategies would.

8. Combining strategies: 3D and 4D parallelism

In practice, no single one of these strategies is used alone at frontier scale β€” each addresses a different bottleneck, and real large training runs combine several simultaneously, an approach commonly called 3D parallelism (data, tensor, and pipeline parallelism composed together) or 4D parallelism when ZeRO-style sharding is layered in as a fourth, orthogonal axis on top of the other three. A typical large-scale layout might use tensor parallelism within a tightly-networked node to split individual layers across a handful of devices, pipeline parallelism across groups of nodes to split the model's depth across stages, data parallelism across many such pipeline replicas to process more of the batch simultaneously, and ZeRO/FSDP-style sharding within the data-parallel dimension to cut redundant memory usage further still. The total device count such a layout consumes is simply the product of the sizes chosen for each axis, $N_{\text{total}} = N_{tp} \times N_{pp} \times N_{dp}$ (with an additional factor for the ZeRO/FSDP sharding degree when composing 4D parallelism). Getting this combination right β€” choosing how many devices to allocate to each axis, and in what topology relative to the actual physical network β€” is itself a substantial and consequential engineering discipline, and it is the concrete, unglamorous reason that training a frontier model is as much a distributed-systems problem as it is a machine-learning problem.

With a model that both fits in aggregate device memory and trains in reasonable wall-clock time, the remaining gap is behavioral rather than infrastructural: a model trained purely on the next-token prediction objective from Chapters 1.1 and 4.2 is very good at continuing text plausibly, but that is a different thing from being a helpful, instruction-following conversational assistant, and closing that gap is where the next chapter turns.

9. Interview angle

Q: Why can't you just train a 70-billion-parameter model on a single accelerator, even given unlimited time? A strong answer distinguishes the memory constraint from the throughput constraint: parameters, gradients, and (for Adam-style optimizers) per-parameter optimizer state together exceed any single device's memory at that scale regardless of how long you're willing to wait, so this is a hard capacity limit, not merely a speed problem β€” though the speed problem (years of wall-clock time on one device) is real too and would need solving separately even if memory weren't an issue.

Q: What's the actual difference between tensor parallelism and pipeline parallelism, and when would you reach for one over the other? A strong answer states that tensor parallelism splits individual weight matrices/layers across devices so each device holds a shard of every layer, requiring frequent, high-bandwidth communication within each layer's computation β€” suited to tightly-networked devices, typically within a node. Pipeline parallelism splits the model's layers into sequential stages across devices, needing communication only at stage boundaries, but risking idle "bubble" time, mitigated by micro-batching β€” better suited to less tightly coupled groups of devices, e.g. across nodes.

Q: Explain what ZeRO/FSDP do differently from tensor or pipeline parallelism. A strong answer clarifies that ZeRO/FSDP don't split the model architecturally at all β€” they keep the data-parallel structure (every device processes different data) but eliminate the redundant full copies of optimizer state, gradients, and parameters that plain data parallelism would otherwise store on every device, sharding those tensors across the data-parallel group and gathering them just-in-time. This is an orthogonal memory-saving axis, not a replacement for tensor/pipeline parallelism.

Q: What is the "bubble" in pipeline parallelism, and how is it addressed? A strong answer describes the bubble as the idle time at the start and end of processing a batch when not all pipeline stages have work to do, because the pipeline needs to fill up and drain, and explains that micro-batching β€” splitting a batch into smaller chunks fed through the pipeline in a staggered, overlapping schedule β€” keeps more devices busy simultaneously and reduces (without fully eliminating) the fraction of time lost to the bubble.

Q: Why do real frontier training runs combine multiple parallelism strategies rather than picking just one? A strong answer explains that each strategy addresses a different resource constraint (memory duplication, layer size, model depth, communication bandwidth) and none of them alone is sufficient at frontier scale β€” tensor parallelism alone doesn't scale across low-bandwidth links, pipeline parallelism alone leaves memory redundancy and bubble inefficiency, data parallelism alone doesn't reduce per-device model memory β€” so composing them (3D/4D parallelism) lets each axis be sized to the constraint it's best suited for, e.g., tensor parallelism within a fast-networked node and pipeline/data parallelism across nodes.

Q: You've sharded parameters and optimizer state with FSDP, but you're still running out of memory when you increase sequence length. What's going on, and what would you do about it? A strong answer identifies that FSDP-style sharding addresses parameters, gradients, and optimizer state, but not activation memory, which grows with sequence length and batch size independently of those. The fix is activation checkpointing (recompute discarded activations during the backward pass instead of storing all of them, trading compute for memory) and, if using tensor parallelism, sequence parallelism to remove the redundant per-device copies of LayerNorm/dropout activations that tensor parallelism alone leaves duplicated.

Q: What does "overlapping communication with compute" actually buy you, and why isn't it automatic? A strong answer explains that every parallelism strategy introduces communication (gradient all-reduces, activation exchange, parameter gathering), and a naive implementation blocks on each communication step, leaving the accelerator idle. Where a piece of communication has no genuine data dependency on the compute that could run concurrently (e.g. one layer's gradient all-reduce versus the next layer's backward pass), a framework can issue it asynchronously and hide its latency behind useful compute β€” real throughput gains, but only where the dependency graph actually permits reordering.

10. Self-check questions

  1. Name the two independent reasons a single accelerator cannot train a frontier-scale model, and explain why they require different kinds of solutions.
  2. Why does data parallelism scale throughput but not the maximum model size you can train?
  3. Walk through how Megatron-LM splits a feedforward block's two linear layers across devices, and explain why splitting the first by columns and the second by rows requires only one communication step rather than one after every matrix multiplication.
  4. What specific problem does pipeline parallelism introduce that tensor parallelism does not, and what is the standard mitigation?
  5. In what precise sense is ZeRO/FSDP "orthogonal" to tensor and pipeline parallelism rather than a competing alternative to them?
  6. Which of ZeRO's stages would you expect to save the most per-device memory, and what would you expect it to cost in exchange?
  7. Sketch a plausible 3D or 4D parallelism layout for training a very large model across many nodes, and justify which parallelism strategy you'd assign to the fastest-networked dimension of your cluster and why.
  8. Why can activation memory exceed parameter memory even after ZeRO/FSDP-style sharding, and what does activation checkpointing trade to reduce it?
  9. Explain specifically what redundancy sequence parallelism removes that tensor parallelism alone leaves in place, and why removing it doesn't require any new communication pattern.

11. Sources

  • Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., & Catanzaro, B. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053. https://arxiv.org/abs/1909.08053
  • Rajbhandari, S., Rasley, J., Ruwase, O., & He, Y. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC20. arXiv:1910.02054. https://arxiv.org/abs/1910.02054
  • Rasley, J., Rajbhandari, S., Ruwase, O., & He, Y. (2020). DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters. KDD 2020. Not on arXiv. https://www.deepspeed.ai/ (DOI: 10.1145/3394486.3406703)
  • Chen, T., Xu, B., Zhang, C., & Guestrin, C. (2016). Training Deep Nets with Sublinear Memory Cost. arXiv:1604.06174. https://arxiv.org/abs/1604.06174
  • Korthikanti, V., Casper, J., Lym, S., et al. (2022). Reducing Activation Recomputation in Large Transformer Models. arXiv:2205.05198. https://arxiv.org/abs/2205.05198

Self-check quiz

Test your understanding of this chapter with a short quiz β€” a mix of concept questions and quick calculations.

What are the two independent reasons a single accelerator cannot train a frontier-scale model?

Explanation: Memory is a hard capacity limit; throughput is a separate, also-real constraint β€” both need solving, and neither is fixed by simply waiting longer.

Why does data parallelism scale throughput but not the maximum model size you can train?

Explanation: Pure data parallelism replicates the entire model, gradients, and optimizer state on every device β€” it helps nothing if the model doesn't fit on one device to begin with.

What is the key structural difference between tensor parallelism and pipeline parallelism?

Explanation: Tensor parallelism needs high-bandwidth, low-latency links since it communicates within every layer; pipeline parallelism communicates only at stage boundaries, tolerating slower interconnects better.

How does ZeRO/FSDP differ from tensor and pipeline parallelism?

Explanation: ZeRO/FSDP is an orthogonal memory-saving axis: same data-parallel structure, but each device holds only a shard of optimizer state/gradients/parameters, gathered just-in-time.

A model has 7 billion parameters. Storing parameters, gradients, and Adam's two optimizer moment tensors β€” four tensors total, each the same size as the parameters β€” all in 4-byte (fp32) precision, what is the total memory required, in GB (using $1\text{ GB} = 10^9$ bytes)?

Explanation: 4 tensors Γ— 7 billion params Γ— 4 bytes = 112 Γ— 10⁹ bytes = 112 GB.

Using the 112 GB total from the previous question, if ZeRO's most aggressive stage shards parameters, gradients, and optimizer state evenly across 8 data-parallel devices, what is the approximate per-device memory required for these tensors, in GB?

Explanation: 112 GB / 8 devices = 14 GB per device β€” a roughly 8x reduction compared to plain data parallelism's full duplication.

Using the approximation that pipeline-parallel bubble overhead β‰ˆ $(p-1)/m$, where $p$ is the number of pipeline stages and $m$ is the number of micro-batches, what is the bubble fraction for $p = 4$ stages and $m = 12$ micro-batches?

Explanation: (4-1)/12 = 3/12 = 0.25 β€” roughly a quarter of total time lost to pipeline fill/drain, before accounting for other overheads.

A transformer layer has hidden dimension $d_{\text{model}} = 4096$. Using tensor parallelism to split the feedforward block's intermediate activation evenly across $t = 8$ devices, how large is the slice of that activation each device computes?

Explanation: 4096 / 8 = 512 β€” each device handles an equal column-wise slice of the intermediate activation, requiring no communication for that step.

What does activation checkpointing trade away to reduce peak activation memory during training?

Explanation: Instead of retaining every layer's activations through the whole forward pass, checkpointing keeps only a sparse set and recomputes the rest on demand during the backward pass β€” roughly 30% extra compute for a substantial memory reduction.

In a 4D parallelism layout combining tensor, pipeline, and data parallelism with ZeRO/FSDP sharding, how is the total device count determined?

Explanation: Each axis splits a different dimension of the problem, so the total device count consumed is the product of the per-axis degrees, $N_{\text{total}} = N_{tp} \times N_{pp} \times N_{dp}$, with an additional factor for ZeRO/FSDP sharding when composing 4D parallelism.