A 70-billion-parameter model does not fit on one GPU, and a model that fits would take years to train on one. So large models train on hundreds or thousands of GPUs at once. This module covers the four basic ways to split the work and memory, and how real runs combine them.

11.1 Why One GPU Is Not EnoughRun this section

There are two limits, memory and time. With the usual mixed-precision setup (section 10.4), training needs per parameter:

what	precision	bytes per parameter
weights used in the forward and backward pass	bf16	2
gradients	bf16	2
master copy of the weights	fp32	4
Adam's two running averages (section 10.2)	fp32	8
total		16

For $\Psi$ parameters that is $16\Psi$ bytes of training state: 1.12 terabytes at 70 billion, against the 80 to 192 GB of the largest common accelerators. On top come the the backward pass needs (section 2.5), which grow with batch size, sequence length, width and depth, and can exceed the training state for long sequences.

Then time: (section 10.5) gives about $7 \times 10^{23}$ FLOPs for an 8-billion-parameter model on 15 trillion tokens, about 55 years on one GPU sustaining $4 \times 10^{14}$ FLOPs per second.

The four techniques below split different things: the batch, the training state, the matrices inside a layer, and the layers themselves.

11.2 Data Parallelism

The simplest split is the batch. Each of $N_d$ GPUs holds a full copy of the model and gets a different slice of each batch:

Every GPU runs the forward and backward pass on its own slice, producing its own gradient.
The GPUs average their gradients, so that every GPU holds the gradient of the whole batch.
Every GPU applies the same update to its own copy, so the copies stay identical.

Step 2 is an : every GPU contributes a tensor and receives the sum (or average). The ring all-reduce sends about twice the gradient's size per GPU regardless of GPU count, so this scales to thousands of GPUs, and it overlaps with the backward pass, sending a layer's gradients while earlier layers are still computing.

Data parallelism multiplies throughput but does nothing for memory: every GPU still holds the full $16\Psi$ bytes, so the model must already fit on one.

11.3 ZeRO and FSDP: Sharding the Training StateRun this section

Keeping $N_d$ identical copies is wasteful: each GPU only needs a layer's full weights while computing it. ZeRO (Rajbhandari et al., 2020) removes the redundancy in three stages, each sharding (splitting into $N_d$ pieces, one per GPU) one more part of the state:

Stage 1 shards the optimizer state: each GPU keeps Adam's averages and master weights for its $1/N_d$ share, and updates only that share.
Stage 2 also shards the gradients.
Stage 3 also shards the weights. Before each layer, the GPUs all-gather its weights (each sends its piece to everyone), use them, and discard them.

With all three stages, the training state per GPU is

M_{\text{ZeRO-3}} = \frac{16\Psi}{N_d}

PyTorch's Fully Sharded Data Parallel (FSDP) implements stage 3. The price is about 1.5 times the traffic of plain data parallelism, since the weights cross the network in both passes; on fast interconnects that is a good trade.

Figure 11.3

How much training state each GPU holds

Choose a model size, a GPU count, a ZeRO stage, and degrees of tensor and pipeline parallelism. The bar is one GPU’s share of the 16 bytes per parameter; the dashed line is an 80 GB GPU.

Activations are not included. Data-parallel degree = GPUs ÷ (tensor × pipeline). ZeRO stages shard across the data-parallel GPUs only.

Predict, then try

Set a 70-billion-parameter model on 8 GPUs with plain data parallelism, then switch to ZeRO stage 3. Before you look, predict whether the per-GPU training state at 8 GPUs fits in 80 GB in each case, and how many GPUs stage 3 needs for it to fit.

In words. Given model size, GPU count, ZeRO stage and tensor and pipeline degrees, the figure stacks each GPU's memory (weights, gradients, master weights, optimizer state) against its capacity, and reports per-step traffic. Raising the stage or degrees shrinks the bar and raises the traffic.

11.4 Tensor ParallelismRun this section

ZeRO still computes each layer on one GPU. Tensor parallelism (Shoeybi et al., 2019, Megatron-LM) splits the matrix multiplications themselves. For on $t$ GPUs, split $W_1$ by columns and $W_2$ by rows to match. GPU $i$ computes $\phi(\mathbf{x}W_1^{(i)})W_2^{(i)}$ alone, since $\phi$ acts entry by entry, and one all-reduce sums the partial results:

\operatorname{FFN}(\mathbf{x}) = \sum_{i=1}^{t} \phi\left(\mathbf{x} W_1^{(i)}\right) W_2^{(i)}

introduces 2 notations

Attention splits by heads, with its share of and one all-reduce. Each GPU holds and computes $1/t$ of every matrix. But the all-reduces happen twice per block and block what follows, so tensor parallelism stays within one server's very fast links (NVLink), typically 8 GPUs at most.

11.5 Pipeline ParallelismRun this section

Pipeline parallelism splits by depth: each of $p$ GPUs holds a stage of $L/p$ consecutive layers, and only activations at stage boundaries cross between GPUs, which suits slower links between servers.

One batch at a time would keep only one stage busy, so each batch is split into $m$ micro-batches streamed through: stage 2 works on micro-batch 1 while stage 1 starts micro-batch 2. Idle time remains while the pipeline fills and drains. In the GPipe schedule (Huang et al., 2019), this pipeline bubble is

\text{bubble} = \frac{p - 1}{m + p - 1}

More micro-batches shrink it. The 1F1B schedule, interleaving forward and backward passes, reduces memory for in-flight micro-batches, and interleaved and zero-bubble schedules shrink the bubble further.

Figure 11.5

Micro-batches fill the pipeline; the grey is the bubble

Each row is a GPU holding one stage. Forward passes run down the stages, backward passes run back up. Grey is idle time.

GPipe schedule; a backward pass is drawn twice as long as a forward pass. Bubble = (p − 1)/(m + p − 1).

Predict, then try

With 4 stages, how many micro-batches are needed before the bubble falls below 10%? Work it out from the formula before trying values.

In words. A timeline with one row per stage and one block per micro-batch, forward and backward in different colours; the grey gaps at start and end are the bubble. Sliders set stages and micro-batches, and the readout gives the bubble fraction.

11.6 Activation CheckpointingRun this section

Storing every layer's activations can take more memory than the model. Activation checkpointing (or gradient checkpointing) keeps only some, the checkpoints, and reruns the forward pass from the nearest one when a discarded activation is needed.

Checkpointing at block boundaries costs roughly one extra forward pass, about a third more compute per step, and can cut activation memory several-fold. Chen et al. (2016) showed that a checkpoint every layers makes memory grow like $\sqrt{L}$ instead of $L$ . Modern systems recompute selectively, keeping tensors that are cheap to store but expensive to recompute, such as attention outputs.

Browser lab

Checkpointing trades memory for recomputation

A simulator walks the backward pass of an L-layer chain, keeping only every k-th layer’s activation and recomputing each segment from its checkpoint when the gradient reaches it. It counts the most activations held at once and the layers recomputed. Section 11.6 predicts that k ≈ √L makes the peak grow like √L instead of L, for about one extra forward pass.

Counts layer activations of equal size and ignores the gradient buffers; real frameworks also choose which tensors inside a layer to keep.

11.7 Putting It Together, and Expert Parallelism

Real runs combine them, often called 3D parallelism:

Tensor parallelism inside a server, where links are fastest, since it communicates in every layer.
Pipeline parallelism across servers, since it only passes activations at stage boundaries.
Data parallelism (usually with ZeRO or FSDP) across copies of the arrangement, since its traffic overlaps with computation.

For example, 1,024 GPUs as tensor parallelism of 8 per server, pipeline of 4 across servers, and 32 data-parallel copies: $8 \times 4 \times 32 = 1{,}024$ , with activation checkpointing on top as needed. Long sequences add sequence or context parallelism, dividing one sequence's tokens across GPUs.

Mixture-of-experts models (Module 17) contain many feed-forward networks, the experts, of which each token uses a few. Expert parallelism places experts on different GPUs and sends each token to its chosen experts' GPUs, an all-to-all exchange. Section 17.5 shows how this shapes expert design.

Exercises

Check

Which of the four techniques reduces the memory each GPU needs for the training state, and which only increases throughput?

Apply

A 70-billion-parameter model is trained with ZeRO stage 3 on 64 GPUs. How many gigabytes of training state does each GPU hold, at 16 bytes per parameter?

Apply

What is the pipeline bubble for 4 stages and 8 micro-batches?

Apply

With ZeRO stage 1 on $N_d$ GPUs, each GPU keeps full bf16 weights and gradients but only its share of the 12 bytes per parameter of master weights and Adam state. Write the per-GPU memory, and evaluate it for a 7-billion-parameter model on 8 GPUs, in GB.

Extend

Explain why tensor parallelism is kept within a server while data parallelism spans the whole cluster.

Core Academic References

Rajbhandari et al. (2020): ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC 2020. arXiv:1910.02054
Zhao et al. (2023): PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. VLDB 2023. arXiv:2304.11277
Shoeybi et al. (2019): Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053
Narayanan et al. (2021): Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. SC 2021. arXiv:2104.04473
Huang et al. (2019): GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. NeurIPS 2019. arXiv:1811.06965
Chen et al. (2016): Training Deep Nets with Sublinear Memory Cost (activation checkpointing). arXiv:1604.06174
Korthikanti et al. (2022): Reducing Activation Recomputation in Large Transformer Models. arXiv:2205.05198