Tensor Parallelism: Block Products Across Devices

A layer computes y equals W x, with a matrix too large for one device. Tensor parallelism cuts W into blocks and gives each device one block. Here there are two devices, blue and orange.

Cut W by rows. Each device holds half of the rows, multiplies them by the whole of x, and gets its own slice of y. An all gather puts the two slices together.

Cut W by columns instead. Each device holds half of the columns and the matching half of x. Each one gets a full length vector, but only a partial sum. An all reduce adds the two partial sums, and that is y.

Now a small network: W one, then a ReLU, then W two. Cut W one by rows. Each device gets its own slice of the hidden layer, and the ReLU acts on each entry by itself, so it runs on the slices with no talking at all.

Then cut W two by columns. Each device multiplies its slice of the hidden layer by its own columns, and a single all reduce adds the two partial outputs. One communication for the whole block.

Other cuts cost more. Cut W one by columns, and each device holds only a partial sum of the hidden layer. The ReLU of a partial sum is not part of the ReLU of the total, so the devices must add their sums before the ReLU.

Your turn. Choose how to cut each matrix so that the network gives the right output with a single communication.

Tensor parallelism is block matrix multiplication spread over devices. Rows split the output, columns split the sum, and pairing them the right way keeps the communication to one all reduce.