6.21 Scaling training: precision, accumulation, distribution

What this is and why it exists

At some point the model no longer fits, and the problem stops being a modelling problem and becomes an engineering one. There are three levers — use smaller numbers, fake a bigger batch, split the work across devices — and they compose. There is also one assumption that wastes more money than any of them save: that adding accelerators makes training proportionally faster. It does not, unless you have measured what the ones you already have are doing.

The vocabulary

• Mixed precision — running most operations in a lower-precision number format.
• Loss scaling — multiplying the loss so small gradients survive the lower format.
• Gradient accumulation — summing gradients over several batches before stepping.
• Effective batch size — the batch the optimiser behaves as though it saw.
• Data parallel — every device holds the whole model and processes different data.
• Sharding — splitting the model's own state across devices rather than replicating it.
• Activation checkpointing — discarding intermediate values and recomputing them in the backward pass.
• Utilisation — the fraction of time an accelerator is actually computing.

The mental model

Start by understanding what occupies the memory, because each lever targets a different part. Four things: the parameters, the gradients (one number per parameter), the optimiser state (for an adaptive optimiser, two more numbers per parameter), and the activations saved during the forward pass for use in the backward pass. The first three scale with the model; the last scales with the model and the batch size together, and for a large batch it dominates. Knowing which one is filling your memory tells you which lever to pull.

Mixed precision is the first and largest win. PyTorch describes it as an arrangement "where some operations use the torch.float32 (float) datatype and other operations use lower precision floating point datatype", and the automatic system "tries to match each op to its appropriate datatype" — matrix multiplications and convolutions in the lower format, reductions and anything numerically delicate kept in full. Memory for activations roughly halves and the arithmetic runs considerably faster on hardware built for it.

The half-precision format needs one safeguard, and the reason is precise. The documentation states it plainly: "gradient values with small magnitudes may not be representable in float16. These values will flush to zero ('underflow'), so the update for the corresponding parameters will be lost." The remedy is loss scaling — "'gradient scaling' multiplies the network's loss(es) by a scale factor and invokes a backward pass on the scaled loss(es). Gradients flowing backward through the network are then scaled by the same factor", moving them into a range the format can represent, with the factor removed before the optimiser steps. The brain-float format has the same range as full precision and less precision, so it needs no scaling at all, which makes it the simpler choice on hardware that supports it. One usage rule from the documentation is worth quoting because people get it wrong: autocast "should wrap only the forward pass(es) of your network, including the loss computation(s). Backward passes under autocast are not recommended."

Gradient accumulation fakes a large batch on small hardware. Run several small batches, add their gradients together, and step once. The optimiser sees a gradient equivalent to the total, so the effective batch size is the small batch times the number accumulated. It is the simplest way to train something that will not otherwise fit, and it costs nothing but time — the same computation, arranged differently. Two details: average rather than sum, or the gradient magnitude changes with the accumulation count and your learning rate no longer means what it did; and take care with normalisation layers, which still see only the small batch, so their statistics are those of the small batch regardless of the effective size.

Activation checkpointing trades computation for memory and belongs alongside these two. Discard the saved intermediate values for a segment of the network and recompute them during the backward pass. The cost is roughly one extra forward pass; the saving on the activation memory that dominates large-batch training is substantial, and it is frequently what makes a configuration fit at all.

Then distribution, in three levels. The simplest replicates the model on every device, splits each batch between them, and averages the gradients — every device holds the complete model and its complete optimiser state, so this only helps when the model already fits and you want to go faster. The next level shards the model's own state: parameters, gradients and optimiser state are divided across devices and gathered only when needed. Since the optimiser state for an adaptive optimiser is the largest of the three, sharding it first buys the most for the least communication, which is why these schemes are staged. Sharding is what makes very large models trainable at all, because no single device ever has to hold the whole state.

Everything distributed is paid for in communication. Gradients must be combined across devices every step, and that traffic can exceed the computation it enables — which is why doubling the devices reliably fails to halve the time. Larger batches per device improve the ratio; overlapping communication with computation helps and the frameworks do it; and beyond a certain point the interconnect between devices decides your throughput rather than the accelerators themselves.

Finally, the mistake this topic exists to prevent. Before scaling out, measure utilisation. If your accelerator is busy 40 percent of the time, the other 60 is spent waiting — usually for data, sometimes for synchronisation, sometimes for small operations that never fill the device — and adding a second accelerator gives you two devices idle 60 percent of the time and a communication cost on top. Profile first, and fix the input pipeline before buying anything. The dataloader topic's remedies apply directly: more workers, cheaper transformations, pinned memory, a storage format that reads in one operation. Then check the batch size, since a batch too small to fill the device leaves it partly idle by construction. Only when utilisation is high is scaling out the right purchase, and then the question becomes the communication ratio rather than the device count.

What you should now be able to explain or do

Name the four things occupying memory and say which scales with batch size. Explain what mixed precision changes and why the half-precision format needs loss scaling. State the rule about what autocast should wrap. Use gradient accumulation correctly, including the averaging and normalisation caveats. Say what activation checkpointing trades. Distinguish replication from sharding and say why the optimiser state is sharded first. Explain why device count does not translate into speed. Measure utilisation and act on it in the right order.

Check yourself