Half a Layer Each, and a Phone Call in the Middle of Every One
Last timeAdding Up the Gradients
Splitting a single layer across devices means exchanging numbers partway through the forward pass. That traffic scales with tokens, not parameters, which puts it on the fastest wire you own.
The splits so far have kept the model whole and divided the data. This one
divides the model itself, and it changes the character of the problem
completely, because the devices now have to talk to each other in the middle of
computing an answer rather than at the end of it.
Two ways to cut a matrix
Nearly everything in a transformer is a matrix multiply, so start there. An
input of some width is multiplied by a weight matrix to produce an output of
some width.
Cut the weight matrix by columns and give each device a slice. Every device
needs the whole input, and each produces a slice of the output. The outputs are
different parts of the answer and they go side by side: no arithmetic is needed
to combine them, only concatenation.
Cut the weight matrix by rows instead. Each device needs only the matching
slice of the input, and each produces something the full width of the output.
But none of them is the answer. Each is a partial sum, and the true output is
the sum of all of them. Combining now requires adding numbers across devices.
That asymmetry is the whole lesson. One cut is free to recombine and the other
is not.
Arranging a block so there is only one exchange
The useful trick is to chain the two cuts so the expensive recombination
happens once rather than twice.
The lesson stops here
4 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