Module 5 of 6
Multi-GPU training: share the load
Sometimes one bag isn't enough. The model doesn't fit, or training would take months. Then we split the work, and the art is choosing what to split.
If I give my friends part of the load, they have to talk to each other. So every strategy trades memory or time for communication.
| Strategy | What is split | Each GPU holds | Talks when | Use it when |
|---|---|---|---|---|
| Data | The batch | A full copy of the model | After every step (gradients) | The model fits; training is too slow |
| ZeRO / FSDP | Optimizer states, gradients, weights | A slice of the training state | Every step, gathering slices | Training state doesn't fit |
| Tensor | Each layer's matrices | Part of every layer | Inside every layer | Layers are huge; GPUs share fast NVLink |
| Pipeline | The layers | A stage of consecutive layers | Between stages (activations) | The model is deep and spans machines |
1. Data parallelism
Every GPU gets a full copy of the model and a different slice of the batch. After the backward pass, the GPUs average their gradients with an all-reduce, so every copy takes the same update.
The volume per GPU barely changes as you add GPUs: it approaches 2M. But every step must wait for it, which is why 8 GPUs rarely give 8× the speed.
2. ZeRO and FSDP: stop storing copies
Plain data parallelism stores the full training state on every GPU. With mixed-precision Adam that is about 16 bytes per parameter: 2 for bfloat16 weights, 2 for gradients and 12 for float32 master weights and Adam's two states. ZeRO shards this state across GPUs instead.
3. Tensor parallelism: split each layer
A single matrix multiplication is split across GPUs. In a transformer MLP, the first weight matrix is split by columns and the second by rows, so each GPU computes its half end to end and the GPUs only need one all-reduce at the end of the block.
The input X (grey) is copied to both GPUs. The first weight matrix is split by columns (green on GPU 1, orange on GPU 2), the second by rows, so each GPU finishes its half before one all-reduce.
Because GPUs talk inside every layer, tensor parallelism needs very fast links. It is usually kept within one machine, where GPUs share NVLink at around 900 GB/s, rather than across machines on a network that is often 10 to 20 times slower.
4. Pipeline parallelism: split the layers
Consecutive layers become stages on different GPUs, like an assembly line. To keep everyone busy, the batch is cut into micro-batches. The idle gaps at the start and end are called the bubble.
Putting it together: 3D parallelism
The largest models combine all three: tensor parallelism inside each machine, pipeline parallelism across machines, and data parallelism across copies of the whole setup. The design goal is always the same: keep the chattiest communication on the fastest links.