Every Machine Does the Whole Thing, and Then They Agree
Last timeWhen One Is Not Enough
The simplest split gives every device the whole model and a different slice of the batch. What they have to exchange each step is fixed, and that fact decides the efficiency.
Start with the arrangement that needs the least machinery, because it is also
the one most runs use, often in combination with the others.
The arrangement
Give every device a complete copy of the model: weights, gradients, optimiser
state, all of it. Split the batch so each device gets a different slice.
Each device runs the forward pass on its own data, runs the backward pass, and
arrives at a gradient. Those gradients differ, because the data differed. If
each device now applied its own gradient, the copies would diverge and you
would have many different models rather than one trained on more data.
So before updating, the devices exchange gradients and each ends up holding the
average across all of them. Every device applies the same averaged gradient to
an identical copy of the weights, so every copy stays identical, and the run is
mathematically equivalent to one device processing the whole batch.
What goes over the wire
The exchanged quantity is the gradient: one number per parameter, at whatever
precision the arithmetic uses.
It is worth pausing on what that does not depend on. It does not depend on the
batch size, the sequence length, or how much work each device did. A device
that processed four sequences and a device that processed four hundred both
produce exactly one gradient per parameter, and both send exactly the same
number of bytes.
For a seven billion parameter model in half precision, that is fourteen
gigabytes per device per step. The exchange moves roughly twice that in
practice, for reasons the next lesson derives, so call it twenty-eight
gigabytes of traffic per step, every step, for weeks.
The lesson stops here
3 more paragraphs to go
You have read the opening. The rest of the argument, the problems that check whether it landed, and the lines worth keeping at the end all come with a plan.
The first lesson of every course in the library reads the whole way through, free, so you can see exactly what the rest of them are.
See the planThe contentsThis is the reading half
Starting the course gives you your own copy of it. Every idea on every page has problems standing under it, marked with a reason rather than a tick, and any sentence you do not believe can be opened and argued with. None of that can happen on a page nobody owns.
The contents