GPUs and distributed training

The big idea

Training a neural network is mostly one operation repeated billions of times: multiply big matrices. A CPU has a handful of powerful cores built for branching, unpredictable code. A GPU has thousands of simpler cores built to do the same arithmetic on many numbers at once, plus very fast memory to feed them. Matrix multiplication is exactly that kind of work.

Then models outgrew single GPUs. The limit is usually not speed but memory: the weights, their gradients, the optimizer's bookkeeping and the activations saved for the backward pass all have to sit in GPU memory at the same time. Distributed training is the set of tricks for splitting that work and that memory across many GPUs while keeping them busy. Think of a kitchen preparing a banquet: you can give every cook the full recipe and a share of the plates (data parallel), split each dish's steps between cooks standing side by side (tensor parallel), or set up a line where each cook does one course and passes it on (pipeline parallel).

How it works

Step 1: Where the memory goes

For a model with $\Psi$ parameters trained with Adam in mixed precision, the ZeRO paper counts the model statesEverything that scales with the parameter count during training: weights, gradients and optimizer states such as Adam's momentum and variance. like this:

• fp16 weights: $2\Psi$ bytes, and fp16 gradients: $2\Psi$ bytes.
• Optimizer states: an fp32 copy of the weights, Adam's momentum and its variance, $4\Psi$ bytes each, so $12\Psi$.

Total: $16\Psi$ bytes. For GPT-2 with 1.5 billion parameters that is at least 24 GB, against 3 GB for the fp16 weights alone. A 7.5-billion-parameter model needs 120 GB before a single activation is stored.

On top come activationsIntermediate outputs of each layer saved during the forward pass because the backward pass needs them to compute gradients. They grow with batch size and sequence length., which grow with batch size and sequence length. The same paper's example: GPT-2 at 1.5 billion parameters with sequence length 1K and batch size 32 needs about 60 GB of activations. Activation checkpointingSaving only some layers' activations in the forward pass and recomputing the rest during the backward pass, trading extra compute for much less memory. stores only some of them and recomputes the rest during the backward pass; Chen and colleagues (2016) showed memory can grow with the square root of the number of layers at the cost of one extra forward pass, and ZeRO cites about 33 percent recomputation overhead, bringing that example to about 8 GB.

Step 2: Mixed precision

Training in 16-bit floats halves the memory for weights, activations and gradients and uses the GPU's fastest arithmetic; Micikevicius and colleagues (2017) note half-precision math throughput on recent GPUs was 2x to 8x higher than single precision. Their recipe, still the standard one, has three parts: keep an fp32 master copy of the weights that accumulates the updates, use half-precision arithmetic that accumulates into fp32, and use loss scalingMultiplying the loss by a large constant before the backward pass so small gradients do not underflow to zero in 16-bit floats, then dividing the gradients by the same constant before the update.: multiply the loss by a constant before backpropagation so tiny gradients do not round to zero in fp16, then divide them back before the update.

Step 3: Data parallelism and the all-reduce ring

Data parallelismTraining with a full copy of the model on every GPU, each processing a different slice of the batch, and averaging gradients across GPUs every step. is the default. Each GPU holds the whole model, takes a different slice of the batch, and the gradients are averaged before every update, which gives exactly the same result as one big batch on one device. The averaging is an all-reduceA collective operation where every worker contributes a buffer and every worker ends up with the sum (or average) of all of them.. The Horovod paper from Uber describes the ring version popularized by Baidu: each of N workers talks to two neighbours, 2(N - 1) times; in the first N - 1 rounds received values are added, in the last N - 1 they replace the local ones. It is bandwidth-optimal for large buffers, and each worker sends about $2(N-1)/N$ times its gradient size, nearly constant as N grows.

Data parallelism scales throughput, not memory: every GPU still needs all $16\Psi$ bytes.

Step 4: ZeRO and FSDP: stop storing copies

In plain data parallelism, all N GPUs hold identical optimizer states, gradients and weights. ZeRO (Rajbhandari and colleagues, Microsoft, 2019) removes that redundancy in three stages:

• Stage 1 shards the optimizer states across the N GPUs: about 4x less memory for large N, same communication volume.
• Stage 2 also shards the gradients: about 8x less, same communication volume.
• Stage 3 also shards the parameters, gathering each layer's weights just before it is used: memory falls linearly with N, for about 1.5 times the communication of plain data parallelism.

The paper's table for a 7.5-billion-parameter model on 64 GPUs: 120 GB per GPU without ZeRO, 31.4 GB with stage 1, 16.6 GB with stage 2 and 1.88 GB with stage 3. PyTorch's FSDPFully sharded data parallel: PyTorch's implementation of sharding parameters, gradients and optimizer states across data-parallel workers, the same idea as ZeRO stage 3. brings the same idea into PyTorch: its documentation describes sharding module parameters, gradients and optimizer states across data-parallel workers to save memory at the cost of communication.

Step 5: Tensor and pipeline parallelism

When one layer is too big, or when even sharded training is too slow, split the model itself.

• Tensor parallelismSplitting the matrices inside each layer across GPUs so every GPU computes part of every layer, with collective communication inside the layer. splits each layer's matrices. Megatron-LM (NVIDIA, 2019) splits the two matrix multiplies of a transformer's MLP block across GPUs so the block needs only one all-reduce in the forward pass and one in the backward. That communication happens inside every layer, so it is kept within a server: Megatron's GPUs had 300 GB/s between GPUs inside a server against 100 GB/s between servers. It trained an 8.3-billion-parameter model on 512 GPUs with 8-way model parallelism, sustaining 15.1 PetaFLOPs at 76 percent scaling efficiency.
• Pipeline parallelismSplitting a model's layers into consecutive stages on different GPUs and streaming micro-batches through them so the stages work at the same time. gives each GPU a run of consecutive layers. A naive pipeline leaves most GPUs idle; GPipe (Google, 2018) splits each mini-batch into M micro-batches that flow through the K stages one after another, accumulating gradients and applying them at the end. The idle time, the bubble, becomes small; GPipe found it negligible when M is at least 4 times K.

Large training runs combine all three: tensor parallel inside a server, pipeline across a few servers, data parallel (often sharded) across the rest.

Try it

• The default is 7 billion parameters in mixed precision with no sharding: 112 GB on every GPU, 1.4 times an 80 GB GPU, no matter how many GPUs you add. Now pick ZeRO-3 / FSDP with 8 GPUs: 14 GB each.
• Set 7.5B and 64 GPUs and step through the sharding stages: 120, 31.4, 16.6 and 1.88 GB, the ZeRO paper's table. Notice how much of the 120 is optimizer state.
• Set 1.5B in mixed precision with no sharding: 24 GB. Add 60 GB of activations (the paper's sequence-length-1K, batch-32 example) and it no longer fits 80 GB; at about 8 GB, the checkpointed figure, it does.
• Switch to inference: 70B with 16-bit weights needs 140 GB; at 4 bits, 35 GB fits a 40 GB GPU. That is why quantization ( Model compression ) matters most for serving.

Going deeper

Code: memory, a ring all-reduce, and loss scaling

import numpy as np
# 1. Memory for model states with mixed-precision Adam: 2 (fp16 weights) + 2 (fp16 grads) + 12 (fp32 copy, m, v) bytes.
def per_gpu_gb(params, gpus=1, stage=0):
    w, g, o = 2 * params, 2 * params, 12 * params
    return (w / (gpus if stage >= 3 else 1) + g / (gpus if stage >= 2 else 1) + o / (gpus if stage >= 1 else 1)) / 1e9
for name, p in [("GPT-2 size, 1.5B", 1.5e9), ("7.5B", 7.5e9)]:
    print(f"{name:17s} one GPU {per_gpu_gb(p):6.1f} GB | 64 GPUs: ZeRO-1 {per_gpu_gb(p, 64, 1):5.1f}"
          f"  ZeRO-2 {per_gpu_gb(p, 64, 2):5.1f}  ZeRO-3 {per_gpu_gb(p, 64, 3):5.2f} GB")

# 2. Ring all-reduce: N workers each hold a gradient; afterwards every worker holds the sum.
def ring_allreduce(grads):
    N = len(grads); chunks = [np.array_split(g.copy(), N) for g in grads]; sent = 0
    for step in range(N - 1):                                   # reduce-scatter: pass partial sums around
        for i in range(N):
            c = (i - step) % N; chunks[(i + 1) % N][c] += chunks[i][c]; sent += chunks[i][c].size
    for step in range(N - 1):                                   # all-gather: pass finished chunks around
        for i in range(N):
            c = (i + 1 - step) % N; chunks[(i + 1) % N][c] = chunks[i][c].copy(); sent += chunks[i][c].size
    return [np.concatenate(ch) for ch in chunks], sent / N
rng = np.random.default_rng(9)
for N in [4, 8, 64]:
    grads = [rng.normal(size=240_000) for _ in range(N)]
    out, per_worker = ring_allreduce(grads)
    ok = all(np.allclose(o, np.sum(grads, axis=0)) for o in out)
    print(f"{N:2d} GPUs: every copy equals the sum: {ok}; each GPU sent {per_worker / 240_000:.3f}x"
          f" its gradient size; 2(N-1)/N = {2 * (N - 1) / N:.3f}")

# 3. Mixed precision: small gradients underflow in fp16 unless the loss is scaled up first.
grad = np.array([3e-4, 2e-6, 5e-8, 1e-9])
print("fp16 without scaling:", grad.astype(np.float16).astype(np.float64))
print("fp16 with loss scale 1024, unscaled in fp32:", (grad * 1024).astype(np.float16).astype(np.float64) / 1024)

GPT-2 size, 1.5B  one GPU   24.0 GB | 64 GPUs: ZeRO-1   6.3  ZeRO-2   3.3  ZeRO-3  0.38 GB
7.5B              one GPU  120.0 GB | 64 GPUs: ZeRO-1  31.4  ZeRO-2  16.6  ZeRO-3  1.88 GB
 4 GPUs: every copy equals the sum: True; each GPU sent 1.500x its gradient size; 2(N-1)/N = 1.500
 8 GPUs: every copy equals the sum: True; each GPU sent 1.750x its gradient size; 2(N-1)/N = 1.750
64 GPUs: every copy equals the sum: True; each GPU sent 1.969x its gradient size; 2(N-1)/N = 1.969
fp16 without scaling: [2.99930573e-04 2.02655792e-06 5.96046448e-08 0.00000000e+00]
fp16 with loss scale 1024, unscaled in fp32: [2.99930573e-04 2.00048089e-06 5.00003807e-08 9.89530236e-10]

• The memory function reproduces the numbers above: 24.0 GB for 1.5 billion parameters on one GPU, 120.0 GB for 7.5 billion, and 31.4, 16.6 and 1.88 GB with ZeRO stages 1 to 3 on 64 GPUs.
• The ring all-reduce is written out chunk by chunk. For 4, 8 and 64 workers every copy ends up equal to the sum, and each worker sends 1.500, 1.750 and 1.969 times its gradient size, exactly $2(N-1)/N$.
• Loss scaling: cast to fp16 directly, a gradient of 1e-9 becomes 0 and 5e-8 is rounded to 5.96e-8. Multiplied by 1024 before the cast and divided back in fp32, they survive as 9.9e-10 and 5.0e-8.

Choosing a setup

Start with the simplest thing that fits. If the model fits on one GPU, use data parallelism to go faster. If model states do not fit, shard them (ZeRO or FSDP) before reaching for model parallelism. If a single layer or the activations still do not fit, add activation checkpointing, then tensor parallelism within a server, then pipeline stages across servers. At every step measure throughput per GPU, because communication can eat the gain.

Trade-offs

Common mistakes

• Estimating training memory as the size of the weights; with Adam in mixed precision it is 16 bytes per parameter before activations.
• Adding GPUs with plain data parallelism to fix an out-of-memory error; each GPU still needs the full model states.
• Using tensor parallelism across servers on slow interconnect, so communication dominates.
• Training in fp16 without loss scaling and wondering why small gradients vanish.
• Running a pipeline with too few micro-batches, so most GPUs sit in the bubble.

Where you will see this

Every large model today is trained with these techniques. DeepSpeed (Microsoft's library behind ZeRO) and PyTorch FSDP are the common sharding implementations, Megatron-LM's tensor parallelism is widely reused for transformers, and Horovod brought ring all-reduce to TensorFlow users. The same memory arithmetic drives serving decisions for large language models in the AI track, in GPU memory math .

Check yourself

Recap