Distributed Training

A 7-billion-parameter model trained with mixed-precision Adam needs about 16 bytes per parameter just for its weights, gradients and optimizer states: 112 GB (gigabytes). That is more than the 80 GB of an NVIDIA H100 GPU (graphics processing unit), before a single activation is stored. And on trillions of tokens, the $6ND$ rule of section 8 gives years of compute for a single GPU.

Distributed training splits the work across hundreds or thousands of GPUs. There are four ways to split it, and this page explains each one with a small example:

• data parallelism: every GPU has the whole model and a different part of the batch;
• sharding: like data parallelism, but the memory of the model is divided among the GPUs;
• tensor parallelism: each matrix of a layer is cut into pieces, one per GPU;
• pipeline parallelism: each GPU holds a different group of layers.

All of them rely on a few communication operations, which come first. The real libraries cannot run in a browser, so the exercises simulate every GPU as a NumPy array and check that the distributed result equals the single-device one.

Collective communication

GPUs in a training job talk through a few collective operations, each called by all the workers at once (a worker is one GPU taking part in the job):

• All-reduce: every worker has an array; afterwards, every worker has the sum (or the mean) of all of them.
• Reduce-scatter: like all-reduce, but each worker ends up with only one slice of the sum.
• All-gather: each worker has one slice; afterwards, every worker has all the slices.

Example. Three workers hold [1, 2, 3], [10, 20, 30] and [100, 200, 300].

• After an all-reduce, all three hold [111, 222, 333].
• After a reduce-scatter, each holds one slice of that sum: one worker 111, another 222, the third 333.
• An all-gather of those three slices gives everyone [111, 222, 333].

So an all-reduce is exactly a reduce-scatter followed by an all-gather.

The ring algorithm

The obvious all-reduce, everyone sending everything to everyone, floods the network. The ring algorithm arranges the workers in a circle, where each one only sends to its right neighbour:

1. Each array is cut into $N$ chunks, for $N$ workers.
2. Reduce-scatter phase: in each of $N - 1$ steps, every worker sends one chunk to its right neighbour, which adds it to its own copy of that chunk. At the end, each worker holds one fully summed chunk.
3. All-gather phase: $N - 1$ more steps pass the finished chunks around the ring, until every worker has all of them.

Example. The same three workers, each array cut into three chunks of one number. Worker $i$ sends to worker $i + 1$, and worker 2 sends to worker 0.

after	worker 0	worker 1	worker 2
start	[1, 2, 3]	[10, 20, 30]	[100, 200, 300]
reduce-scatter, step 1	[1, 2, 303]	[11, 20, 30]	[100, 220, 300]
reduce-scatter, step 2	[1, 222, 303]	[11, 20, 333]	[111, 220, 300]
all-gather, step 1	[111, 222, 303]	[11, 222, 333]	[111, 220, 333]
all-gather, step 2	[111, 222, 333]	[111, 222, 333]	[111, 222, 333]

In step 1, worker 0 sends its first chunk (1) to worker 1, which adds it: $10 + 1 = 11$. After two steps each worker holds one complete sum (worker 0 has 222, worker 1 has 333, worker 2 has 111), and two more steps copy them around.

Each worker sends $2(N-1)/N$ times the size of the array in total, which is almost independent of $N$: with 4 workers it is 1.5 times the array, with 64 workers 1.97 times, and it never reaches 2. This is why the ring is bandwidth-optimal (Patarasuk and Yuan, 2009). On NVIDIA GPUs these operations are provided by NCCL (NVIDIA Collective Communications Library), which every framework below calls.

Data parallelism

The simplest split: every GPU holds a full copy of the model and processes a different slice of the batch. After the backward pass, the gradients are averaged with an all-reduce, and every GPU applies the same update, so the copies never diverge. It is gradient accumulation (section 8), with the micro-batches running on different GPUs at the same time.

Example. Two GPUs, a batch of 8 texts: each GPU runs 4 of them. For some weight, GPU 0 computes the gradient $0.2$ and GPU 1 computes $0.4$. The all-reduce (here a mean) gives both GPUs $0.3$, the gradient of the full batch of 8, and both apply the same update.

PyTorch's DistributedDataParallel, or DDP (distributed data parallel; Li et al., 2020), adds the engineering that makes it fast: gradients are grouped into buckets, and each bucket's all-reduce starts as soon as its gradients are ready, so that communication overlaps with the rest of the backward pass.

Sharding the model states: ZeRO and FSDP

Data parallelism keeps $N$ identical copies of everything. Rajbhandari et al. (2020) counted what "everything" is for mixed-precision Adam with $\Psi$ (psi) parameters:

• $2\Psi$ bytes of weights in FP16 (16-bit floating point);
• $2\Psi$ bytes of gradients in FP16;
• $12\Psi$ bytes of optimizer states: a master copy of the weights, the momentum and the variance, each in FP32 (32-bit floating point), 4 bytes each.

That is $16\Psi$ bytes in total, on every GPU. ZeRO (Zero Redundancy Optimizer) removes the redundancy in three stages, each one sharding one more piece: cutting it into $N$ parts, one per GPU, so that each GPU stores and updates only its own part.

How to read a row: the unsharded pieces cost their full size on every GPU, the sharded ones are divided by $N$. At stage 1, for example, the $4\Psi$ of weights and gradients stay whole, and the $12\Psi$ of optimizer states are split $N$ ways.

Example. A 7-billion-parameter model on $N = 8$ GPUs:

stage	bytes per GPU
none	$16 \cdot 7 = 112$ GB
1	$4 \cdot 7 + 12 \cdot 7 / 8 = 28 + 10.5 = 38.5$ GB
2	$2 \cdot 7 + 14 \cdot 7 / 8 = 14 + 12.25 = 26.25$ GB
3	$16 \cdot 7 / 8 = 14$ GB

Without ZeRO the model does not fit on an 80 GB GPU; from stage 1 on, it does.

With sharded optimizer states, a training step changes shape: the gradients are reduce-scattered (each GPU receives the averaged gradient of its own shard only), each GPU updates its shard, and the updated shards are all-gathered so that everyone has the full weights again. Since a reduce-scatter plus an all-gather is an all-reduce, stages 1 and 2 cost no more communication than plain data parallelism.

At stage 3 each GPU stores only its shard of the weights; before a layer runs, its full weights are rebuilt with an all-gather, and freed again afterwards. The communication volume grows to 1.5 times that of plain data parallelism, and memory per GPU falls by a factor of $N$. ZeRO is implemented in DeepSpeed; PyTorch's FSDP (Fully Sharded Data Parallel; Zhao et al., 2023) is the same idea built into PyTorch.

Tensor parallelism

When a single layer is too large, or too slow, for one GPU, its matrices can be split. Megatron-LM (Shoeybi et al., 2019) splits the MLP (multilayer perceptron) of a transformer, $Y = \mathrm{GeLU}(XA)\,B$, so that it needs only one all-reduce:

• The first matrix $A$ is split by columns, $A = [A_1, A_2]$. Each GPU computes $\mathrm{GeLU}(XA_i)$ on its own: the nonlinearity GeLU (Gaussian error linear unit) works entry by entry, so the columns never need to meet.
• The second matrix $B$ is split by rows, matching. Each GPU computes $\mathrm{GeLU}(XA_i)\,B_i$, a partial result with the full output shape.
• One all-reduce adds up the partial results: $Y = \sum_i \mathrm{GeLU}(XA_i)\,B_i$.

Example. One token $X = (1, 2)$, $A$ of shape $(2, 4)$, $B$ of shape $(4, 2)$, and 2 GPUs.

• Unsplit: $XA = (1, 2, 1, 0)$, $\mathrm{GeLU}$ of it $\approx (0.841, 1.955, 0.841, 0)$, and multiplied by $B$ the output is $(1.682, 2.796)$.
• GPU 1 has columns 0–1 of $A$ and rows 0–1 of $B$: its partial output is $(0.841, 1.955)$. GPU 2 has columns 2–3 and rows 2–3: $(0.841, 0.841)$.
• The all-reduce adds them: $(1.682, 2.796)$, the same output.

Why columns first? Split $A$ by rows instead, and each GPU would hold part of a sum that must be completed before the GeLU, because GeLU of a sum is not the sum of the GeLUs: $\mathrm{GeLU}(1 + 1) \approx 1.955$, while $\mathrm{GeLU}(1) + \mathrm{GeLU}(1) \approx 1.682$. That would cost a second communication.

Attention is split the same way, by heads. Tensor parallelism communicates inside every layer, so it is used within a node, where GPUs are linked by the fast NVLink connection.

Pipeline parallelism

Pipeline parallelism splits the model by depth: GPU 1 holds the first layers, GPU 2 the next ones, and so on, and the activations flow from stage to stage. With one batch at a time, only one GPU works while the others wait.

GPipe (Huang et al., 2019) splits the batch into $m$ micro-batches that follow each other through the $p$ stages, like cars on an assembly line.

Example. $p = 2$ stages and $m = 3$ micro-batches (0, 1, 2), in the forward pass. Each cell is one time slot; "·" means idle.

stage	slot 0	slot 1	slot 2	slot 3
stage 0	0	1	2	·
stage 1	·	0	1	2

Stage 1 waits in slot 0 for the first micro-batch, and stage 0 has nothing left in slot 3: 2 of the 8 cells are idle.

The pipeline has to fill and drain: with GPipe's schedule, each stage is idle for $p - 1$ of the $m + p - 1$ time slots of the forward pass, and the same in the backward pass. The idle time, called the bubble, is

where $p$ is the number of stages and $m$ the number of micro-batches. Narayanan et al. (2021) write it as $(p-1)/m$ of the ideal time. It shrinks only if $m$ is much larger than $p$.

Example. In the table, $(2 - 1) / (3 + 2 - 1) = 1/4$: the 2 idle cells out of 8. With $p = 4$ stages and $m = 8$ micro-batches the bubble is $3/11 \approx 27\%$ of the time; with $m = 32$ it drops to $3/35 \approx 8.6\%$.

The 1F1B (one forward, one backward) schedule alternates one forward and one backward pass, so that each stage stores the activations of at most $p$ micro-batches instead of $m$; its bubble is the same. Narayanan et al.'s interleaved schedule, where each GPU holds several smaller stages, shrinks the bubble further.

Putting it together

Large runs combine all of them. Narayanan et al. (2021) ran a trillion-parameter GPT (Generative Pre-trained Transformer) model at 502 petaFLOP/s, that is $5.02 \times 10^{17}$ floating-point operations per second, on 3072 GPUs. They combined tensor parallelism inside each node, pipeline parallelism across nodes, and data parallelism over the copies: 3D parallelism.

Activation checkpointing (Chen et al., 2016) adds a further trade: it stores only some activations and recomputes the rest during the backward pass, saving memory at the cost of about one extra forward pass.

The libraries

• NCCL: NVIDIA's collective-communication library; every framework below uses it on NVIDIA GPUs.
• PyTorch DDP and FSDP: data parallelism and sharding in the style of ZeRO stage 3, built into PyTorch.
• DeepSpeed: the original implementation of the ZeRO stages and their offloading variants.
• Megatron-LM / Megatron-Core (NVIDIA): tensor and pipeline parallelism for transformers.
• torchtitan (Liang et al., 2024): a training system native to PyTorch that composes these techniques in a modular way.
• nanotron (Hugging Face): a library for pre-training transformers, with the Ultrascale Playbook, a long practical guide to all of the above.

References

• P. Patarasuk, X. Yuan, Bandwidth optimal all-reduce algorithms for clusters of workstations, Journal of Parallel and Distributed Computing 69(2), 117–124, 2009.
• S. Li et al., PyTorch Distributed: Experiences on Accelerating Data Parallel Training, Proceedings of the VLDB Endowment 13, 3005–3018, 2020. (arXiv)
• S. Rajbhandari, J. Rasley, O. Ruwase, Y. He, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC20, 2020. (arXiv)
• Y. Zhao et al., PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel, Proceedings of the VLDB Endowment 16, 3848–3860, 2023. (arXiv)
• M. Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, arXiv, 2019. Section 3: the column/row split of the MLP.
• Y. Huang et al., GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS, 2019.
• D. Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21, 2021. Section 2.2: the pipeline bubble; 3D parallelism.
• T. Chen, B. Xu, C. Zhang, C. Guestrin, Training Deep Nets with Sublinear Memory Cost, arXiv, 2016. Activation checkpointing.
• W. Liang et al., TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training, arXiv, 2024.
• NCCL documentation, DeepSpeed, Megatron-LM, torchtitan, nanotron and the Ultrascale Playbook.