Where the memory actually goes
Before you can scale anything, you have to know what is filling the device. A training step for a model with P parameters holds four large tensors: the weights, the gradients, the optimizer states, and the activations saved for the backward pass. With Adam in mixed precision the first three alone cost roughly 16 bytes per parameter — 2 for the bf16 weight, 2 for the bf16 gradient, and 12 for the fp32 master copy plus the two Adam moments. A 7-billion-parameter model therefore needs about 112 GB just to exist, before a single token of activation memory. That number is why a single 80 GB accelerator cannot train it, and it is the starting point for every decision in this track. See mixed-precision training for why we keep both low- and high-precision copies.
Mixed-precision Adam holds roughly 16 bytes per parameter — weights, gradients, and optimizer states together.
Data parallelism: replicate and average
The simplest way to use N devices is data parallelism: put a full copy of the model on each one, feed each a different slice of the global batch, and after the backward pass average everyone's gradients so every replica takes the same step. Mathematically you are computing one large-batch SGD update, just with the per-example gradients spread across machines. This is the backbone of distributed training and the default starting configuration even for the largest runs.
The averaging step is a collective communication called all-reduce: every device ends up holding the summed gradient. Doing it naively (send everything to one master) makes that master a bottleneck whose traffic grows with N. The fix is ring all-reduce, which arranges devices in a logical ring and passes gradient chunks around it, so each link carries a constant volume regardless of how many devices you add — bandwidth-optimal and the reason data parallelism scales to thousands of GPUs.
All-reduce sums every replica's gradient, then dividing by N gives the averaged update.
for micro in range(accum_steps):
loss = model(next_batch()) / accum_steps
loss.backward() # grads accumulate in .grad
all_reduce(model.grads, op=SUM) # ring all-reduce across N ranks
model.grads /= world_size # turn the sum into an average
optimizer.step(); optimizer.zero_grad()Bigger batches without bigger memory
Often the batch you want (for stable optimization) is larger than the batch that fits. gradient accumulation solves this: run several micro-batches in sequence, let their gradients pile up in the `.grad` buffers, and only call the optimizer after the last one. The effective batch becomes `micro_batch × accum_steps × world_size` while peak activation memory stays at a single micro-batch. The right target batch size is not arbitrary — the gradient noise scale tells you roughly how large a batch your problem can usefully absorb before extra examples stop improving the gradient estimate.
Accumulating K micro-batches across N replicas yields a large effective batch without extra memory.
Hiding the network behind the math
Every all-reduce costs time on the wire, and if the GPUs sit idle waiting for it, your expensive cluster is mostly heating the room. The decisive optimization is communication–computation overlap: start reducing the gradients of the last layers the instant they are computed, while the backward pass is still grinding through earlier layers. Because backprop produces gradients layer by layer, almost the entire all-reduce can be hidden under compute that was going to happen anyway.
- Bucket parameters so several layers' gradients form one all-reduce, amortizing per-call launch overhead.
- Fire each bucket's all-reduce from a backward hook the moment that bucket is full.
- Run the reduction on a separate CUDA stream so it overlaps the ongoing backward compute.
- Synchronize only once, just before the optimizer step, when every bucket must be done.
The blunt diagnostic: achieved throughput over theoretical peak measures how well overlap hides the network.