Sixteen Copies of the Same Numbers, in Sixteen Places
Last timeGiving Each Machine Some Layers
The data split has every device storing an identical copy of everything. Stop doing that, hand each device a slice, and fetch the rest just before it is needed.
The two previous lessons cut the model up. This one does not cut the model at
all. The data split stays exactly as it was, and the only thing that changes is
what each device bothers to store.
What every device is holding
Go back to the first lesson. Training a seven billion parameter model needs
about fourteen gigabytes for the half-precision weights, fourteen for the
gradients, and eighty-four for the master copy and the two running averages the
optimiser keeps. A hundred and twelve gigabytes of state.
Under the plain data split, every device holds all of it. And because every
device took an identical step from an identical starting point using an
identical averaged gradient, the copies are not merely similar, they are bit
for bit the same numbers.
Splitting the three pieces
The saving comes in three stages, and they are usually implemented as switches
you turn on in order, because each one saves more and costs more.
The optimiser state goes first. It is the largest of the three and the easiest,
because it is only touched once per step, at the moment the weights are
updated. Give each device one sixty-fourth of the parameters to be responsible
for, let it hold the master copy and the running averages for only that slice,
and let it update only that slice. At the end of the step, every device has a
fresh copy of its own slice of the weights, and a short exchange gives everybody
everybody else's slice.
The gradients go second. Rather than every device ending up with the full
averaged gradient, each device ends up with the averaged gradient for only the
slice it is responsible for. This is a different collective from the one in the
third lesson: the same addition, but the result scattered rather than given to
everybody. It costs exactly half the traffic of the full exchange.
The weights go third, and this is the one that changes the shape of the
computation rather than just the bookkeeping.
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