Training at Scale

tensor parallelism

Some weight matrices are simply too large to fit, or too slow to apply, on one device. Tensor parallelism splits an individual operator across devices: instead of replicating a matrix multiply, it partitions the weight matrix itself and lets each device compute a piece of the result. The classic transformer recipe splits the first feed-forward matrix column-wise so each device produces a slice of the hidden activation, then splits the second matrix row-wise so a single all-reduce sums the partial outputs back into the full result.

Attention is sharded by heads — each device owns a subset of the query, key, and value projections and computes those heads independently, with an all-reduce after the output projection. The arrangement is chosen so that only two collectives per block are needed in the forward pass (and two in the backward), keeping the cut surfaces small. Because every micro-step waits on an all-reduce of full activation tensors, tensor parallelism demands very high bandwidth and low latency, which is why it is normally confined within a single node's NVLink domain.

Tensor parallelism is the lever you reach for when a layer is the bottleneck — huge hidden dimensions, or activations that must shrink per device. It is rarely used alone past eight-way: beyond the NVLink boundary the all-reduce cost dominates, so larger runs nest it inside pipeline and data parallelism.

Y = \text{Dropout}\big(\,[\,XA_1,\,XA_2\,]\,B\,\big),\quad B=[B_1;B_2]

Megatron's MLP split: column-parallel A (no communication), then row-parallel B whose partial products are summed by one all-reduce.

Tensor parallelism partitions the math of each operator, so the result is bit-equivalent to the unsharded layer — unlike pipeline parallelism, it does not change the dependency structure, only where the multiplies happen.

Also called
TPintra-layer parallelism張量平行層內平行