Training on several GPUs at once

Split the batch, average the gradients, take the same step

Each GPU is driven by its own copy of the training program, a separate process. The processes are numbered 0, 1, 2, …; the number is called the rank. Each process's data loader reads a different share of the examples, so no example is used twice in an epoch. The averaging step is called an all-reduce: every GPU sends in its gradient, and every GPU gets back the result. Only one process, usually rank 0, writes checkpoints and logs, since all copies are the same.

The averaged gradient matches the whole batch

This snippet checks the claim on the CPU. It computes one gradient on a batch of 256 examples, then four gradients on shares of 64 using four copies of the model, and averages them.

same_gradient.py

import copy
import torch

torch.manual_seed(0)                                  # same random draws on every run
model = torch.nn.Linear(4, 1)                         # a small model: 4 inputs, 1 output
x, y = torch.randn(256, 4), torch.randn(256, 1)       # one batch of 256 examples

def grad(m, xb, yb):                                  # gradient of the mean loss
    m.zero_grad()                                     # clear the old gradient
    ((m(xb) - yb) ** 2).mean().backward()
    return m.weight.grad.clone()                      # a copy of the gradient

whole = grad(model, x, y)                             # one GPU, the whole batch

copies = [copy.deepcopy(model) for _ in range(4)]     # four "GPUs", one copy each
# x.chunk(4) cuts x into 4 shares of 64 rows, one for each copy.
shares = [grad(m, xs, ys) for m, xs, ys in zip(copies, x.chunk(4), y.chunk(4))]
average = torch.stack(shares).mean(0)                 # the averaging step

print("share 0:", shares[0])
print("average:", average)
print("whole:  ", whole)
print("largest difference:", (average - whole).abs().max().item())

share 0: tensor([[-0.2048, -0.1499, -0.9062, -1.0957]])
average: tensor([[ 0.0028,  0.6184, -0.9061, -0.7804]])
whole:   tensor([[ 0.0028,  0.6184, -0.9061, -0.7804]])
largest difference: 3.5762786865234375e-07

One share alone gives a quite different gradient. The average of the four matches the whole batch. The tiny difference, about 4 × 10⁻⁷, comes from adding the same numbers in a different order with float32 rounding (see mixed precision for how floats round).

Each extra GPU helps less, because the averaging time does not shrink

With N GPUs, each GPU computes on 1/N of the batch, so computing takes 1/N as long. The all-reduce adds time that does not shrink, because every GPU must send and receive a full set of gradients. Suppose one GPU takes 400 ms for a batch, and the all-reduce takes 25 ms. (These are made-up round numbers.)

Try it. Predict the speed-up with 4 GPUs. Then move the slider to find the number of GPUs at which the GPUs spend as much time averaging as computing.