# Module 11: Training Across Many GPUs

> Where GPU memory goes during training and how it is split: data parallelism, ZeRO and FSDP, tensor and pipeline parallelism, activation checkpointing.

- Part III: Building an LLM · core track · Advanced · 4 hours
- By the end you can explain how a model too large for one GPU is trained, and estimate memory per GPU.
- Builds on: [Module 10: Pretraining](https://transformers-and-math.drmhse.com/md/10-pretraining.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/11-training-across-many-gpus
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com

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 Enough

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 *activations* 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: $C \approx 6ND$ (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:

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

Step 2 is an *all-reduce*: 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 State

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.

> **Interactive figure** `parallelism-planner`, on the web page: https://transformers-and-math.drmhse.com/modules/11-training-across-many-gpus#figure-parallelism-planner
>
> **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.
>
> **What happens:** Plain data parallelism needs the full 1,120 GB on every GPU, whatever the GPU count: it never fits. Stage 3 divides it by the GPU count: 140 GB at 8 GPUs, still too much, but 70 GB at 16. Activations come on top, so in practice more GPUs, or tensor and pipeline parallelism, are added as well.

> **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 Parallelism

ZeRO still computes each layer on one GPU. *Tensor parallelism* (Shoeybi et al., 2019, Megatron-LM) splits the matrix multiplications themselves. For $\operatorname{FFN}(\mathbf{x}) = \phi(\mathbf{x}W_1)W_2$ 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)}
$$

Attention splits by heads, with its share of $W^O$ 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 Parallelism

*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.

> **Interactive figure** `pipeline-bubble`, on the web page: https://transformers-and-math.drmhse.com/modules/11-training-across-many-gpus#figure-pipeline-bubble
>
> **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.
>
> **What happens:** You need $3/(m + 3) < 0.1$, so $m > 27$: 28 micro-batches. The bubble shrinks only slowly, which is why large runs use many micro-batches per step and schedules designed to fill the gaps.

> **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 Checkpointing

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 $\sqrt{L}$ 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** `checkpoint-memory`: an experiment that runs in the reader's browser on the web page: https://transformers-and-math.drmhse.com/modules/11-training-across-many-gpus#lab-checkpoint-memory

## 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

**Exercise (check).** Which of the four techniques reduces the memory each GPU needs for the training state, and which only increases throughput?

*Solution:* Plain data parallelism only increases throughput: every GPU still holds everything. ZeRO/FSDP, tensor parallelism and pipeline parallelism all reduce per-GPU memory: ZeRO by sharding the state, tensor parallelism by splitting every matrix, and pipeline parallelism by giving each GPU only some of the layers.

**Exercise (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?

*Answer:* 17.5 (within 0.1)

*Solution:* $16 \times 70 \times 10^9 / 64 = 17.5 \times 10^9$ bytes, 17.5 GB per GPU, leaving room for activations on an 80 GB device.

**Exercise (apply).** What is the pipeline bubble for 4 stages and 8 micro-batches?

*Answer:* 0.2727 (within 0.001)

*Solution:* $(p - 1)/(m + p - 1) = 3/11 \approx 0.273$: just over a quarter of the time, a typical GPU is idle.

**Exercise (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.

*Answer:* 38.5 (within 0.1)

*Solution:* $M = 4\Psi + 12\Psi / N_d = 7 \times 10^9 \times (4 + 12/8) = 7 \times 10^9 \times 5.5 = 38.5$ GB.

**Exercise (extend).** Explain why tensor parallelism is kept within a server while data parallelism spans the whole cluster.

*Solution:* Tensor parallelism needs an all-reduce inside every block, on the critical path: the next computation cannot start until it finishes, so it needs the fastest possible links, which only exist between GPUs in one server. Data parallelism communicates once per step, the gradients, and that transfer can overlap with the backward pass, so it tolerates the slower network between servers.

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