Collectives and the numbers that decide a design
By the end: You can estimate communication time, memory per GPU and model FLOPs utilisation, and use them to choose a parallelism strategy.
- All-reduce, all-gather, reduce-scatter
- How long a gradient sync takes
- Memory per GPU, worked
- Model FLOPs utilisation
New to distributed training? A GPU has a fixed amount of memory, and a large model plus everything training needs beside it does not fit in one. Distributed training spreads the work across many GPUs, often across many machines. There are two basic moves: give every GPU a complete copy of the model and a different slice of the data, then average what each of them learned; or cut the model itself into pieces so that each GPU holds only part of it. Why it matters
Four numbers decide a distributed training design, and all four are arithmetic you can do before requesting a single GPU: how much memory each device needs, how long a gradient sync takes, what fraction of the step that sync is, and how much of the hardware you are actually using. This section works all four for a concrete model — seven billion parameters on eight GPUs — because the numbers, not the vocabulary, are what an interviewer is probing and what stops you renting a cluster that cannot help.
The three collectives#
Almost all distributed training is built from three operations, and knowing which one a framework is doing tells you what it costs. On NVIDIA hardware they are implemented by NCCL, the library PyTorch calls when you select the nccl backend, and it is NCCL that decides the ring or tree the data actually travels. All-reduce combines a value across every device and leaves every device with the combined result — this is the gradient sync in plain data parallelism. Reduce-scatter combines and then gives each device only its slice. All-gather is the reverse: each device contributes its slice and everyone ends up with the whole. The identity worth remembering is that an all-reduce is exactly a reduce-scatter followed by an all-gather, which is why sharded training can replace one with the other two and pay roughly the same total.
operation each device starts with each device ends with
all-reduce its own full gradient the summed full gradient
reduce-scatter its own full gradient the summed slice it owns
all-gather the slice it owns every slice, assembledHow long a gradient sync takes#
A ring all-reduce does not send everything to everyone. Each device passes slices around a ring, and the traffic per device works out at twice the payload times N minus one over N — approaching twice the payload however many devices you add, which is the useful property. Seven billion bf16 gradients is 14 GB, so on eight devices each one moves 24.5 GB per step.
params, N = 7e9, 8
grad_bytes = params * 2 # bf16 gradients
payload = grad_bytes * 2 * (N - 1) / N # ring all-reduce traffic per device
print(f"{grad_bytes/1e9:.0f} GB of gradients -> {payload/1e9:.1f} GB moved per device")
for link, bw in [("NVLink", 900e9), ("PCIe 4.0 x16", 25e9),
("100 Gb Ethernet", 12.5e9), ("10 Gb Ethernet", 1.25e9)]:
print(f" {link:18} {payload/bw*1000:9.1f} ms per step")
# 14 GB of gradients -> 24.5 GB moved per device
# NVLink 27.2 ms per step
# PCIe 4.0 x16 980.0 ms per step
# 100 Gb Ethernet 1960.0 ms per step
# 10 Gb Ethernet 19600.0 ms per stepA time in milliseconds means nothing until you compare it with the step it sits inside. A transformer forward and backward pass costs about six FLOPs per parameter per token, so at 4,096 tokens per device that is 1,376 TFLOPs a step, and eight H100s sustaining a realistic forty per cent of their 989 TFLOP/s bf16 peak take 435 ms of compute.
interconnect sync step sync as a share
NVLink 27 ms 462 ms 5.9%
PCIe 4.0 980 ms 1415 ms 69.3%
10 Gb Ethernet 19600 ms 20035 ms 97.8%This is the calculation that saves the money
Do it before provisioning, not after a week of poor scaling. If the sync is a large share of the step there are only three moves, and all of them are cheaper than more GPUs: accumulate gradients over several micro-batches so the sync happens once per several forward passes, shard so that reduce-scatter replaces all-reduce, or get a better interconnect. Buying more devices on a slow network makes the ratio worse, because the payload per device barely falls while the number of hops grows.
Memory per GPU, worked#
The memory question is the one people get wrong by a factor of five, because they count the parameters and stop. Mixed-precision training with Adam holds five things per parameter, and the parameters are the smallest of them.
params = 7e9
GiB = 1024 ** 3 # memory is sold and reported in binary units
costs = [("bf16 weights", 2), ("bf16 gradients", 2), ("fp32 master copy", 4),
("Adam momentum", 4), ("Adam variance", 4)]
for name, b in costs:
print(f"{name:18} {b:2d} bytes/param {params*b/GiB:6.1f} GiB")
total = params * sum(b for _, b in costs) / GiB
print(f"{'TOTAL':18} {sum(b for _,b in costs):2d} bytes/param {total:6.1f} GiB")
# bf16 weights 2 bytes/param 13.0 GiB
# bf16 gradients 2 bytes/param 13.0 GiB
# fp32 master copy 4 bytes/param 26.1 GiB
# Adam momentum 4 bytes/param 26.1 GiB
# Adam variance 4 bytes/param 26.1 GiB
# TOTAL 16 bytes/param 104.3 GiBThat total is also the rule of thumb worth memorising, because it answers the interview question directly: mixed-precision Adam costs roughly sixteen bytes per parameter, so multiply your parameter count by sixteen and compare it with your device. Plain SGD with momentum costs eight and full fp32 Adam costs twenty. Activations sit on top and scale with batch size rather than parameter count, which is why they are the part you can trade.
weights = params * 2 / GiB # bf16 copy every device keeps resident
grads = params * 2 / GiB
opt = params * 12 / GiB # fp32 master + Adam momentum + Adam variance
for name, per_gpu in [
("none (plain DDP)", weights + grads + opt),
("stage 1 (optimiser)", weights + grads + opt / N),
("stage 2 (+ gradients)", weights + grads / N + opt / N),
("stage 3 (+ parameters)", (weights + grads + opt) / N),
]:
print(f"{name:23} {per_gpu:6.1f} GiB fits an 80 GiB card: {per_gpu < 80}")
# none (plain DDP) 104.3 GiB fits an 80 GiB card: False
# stage 1 (optimiser) 35.9 GiB fits an 80 GiB card: True
# stage 2 (+ gradients) 24.4 GiB fits an 80 GiB card: True
# stage 3 (+ parameters) 13.0 GiB fits an 80 GiB card: TrueModel FLOPs utilisation#
MFU is the single number that says whether a training run is healthy: the FLOPs your model actually needed, divided by the FLOPs the hardware could have delivered in the same wall-clock. It ignores everything you cannot bill for — recomputation from activation checkpointing, padding, waiting on the network — which is exactly what makes it honest.
tokens_per_step = N * 4096
flops = 6 * params * tokens_per_step # 6 FLOPs per parameter per token
compute_s = flops / (989e12 * 0.4 * N) # 40% of bf16 peak, sustained
comm_s = payload / 900e9 # NVLink
step = compute_s + comm_s
achieved = flops / step / N
print(f"step {step*1000:.0f} ms achieved {achieved/1e12:.0f} TFLOP/s per GPU")
print(f"MFU {achieved/989e12*100:.1f}%")
# step 462 ms achieved 372 TFLOP/s per GPU
# MFU 37.6%MFU you measure what it usually means
40-50% healthy; tune elsewhere
25-35% communication or activation checkpointing overhead
10-20% the data loader, or a batch too small to fill the GPU
under 10% something is on the CPU, or synchronising every stepWhere people go wrong
Planning a multi-node run from GPU FLOPs alone, without estimating how long communication takes over the actual network.
If this comes up in an interview
Estimate how much memory it takes to train a 7B model with Adam in mixed precision.
About 16 bytes per parameter: 2 for 16-bit weights, 2 for gradients, 4 for 32-bit master weights and 8 for Adam's two moments. That is about 112 GB before activations, so it needs sharding across several GPUs.
Then they push further
How long would a gradient all-reduce take across 8 GPUs over a 100 Gbit/s network?
Each GPU moves about 2 times 7/8 of 14 GB, around 24 GB. At about 12.5 GB/s that is roughly two seconds per step, which is why this is done over fast in-machine links or overlapped with computation.
What is MFU and what would a low value tell you?
The share of peak hardware FLOPs spent on model computation. A low value means time goes to communication, data loading, idle pipeline stages or inefficient kernels, and says where to look.
What gets marked up. Strong candidates do the arithmetic out loud with the 16 bytes per parameter figure and relate communication time to interconnect bandwidth.
Build this: For a model size of your choice, write a small calculator for memory per GPU under DDP and full sharding, all-reduce time over two link speeds, and the tokens per second needed for 40% MFU. Check it against a real run if you have one.
Watch it explained
This is one section of Distributed training, 22 hours across 5 sections. Read it in full to get the whole topic in order.