Tensor Parallelism › Why Split Tensors?

Before diving into how tensor parallelism works, it helps to be precise about the problem it solves and why the solutions take the forms they do.

A standard linear layer computes $Y = XA$, where $X \in \mathbb{R}^{b \times d}$ is a batch of $b$ tokens each with hidden dimension $d$, and $A \in \mathbb{R}^{d \times h}$ is a weight matrix mapping to some output dimension $h$. The output $Y \in \mathbb{R}^{b \times h}$ is dense: every output element depends on every input element. At first glance, this density suggests that any partitioning of the computation must involve data dependencies, and therefore communication, across devices.

But the density of the output does not mean the computation is inseparable. Matrix multiplication has a fundamental block structure that allows it to be decomposed in multiple ways. Understanding this structure is the conceptual key to tensor parallelism.

When $d$ and $h$ are large (e.g., $d = 12288$ for GPT-3's hidden size), the weight matrix $A$ can be enormous. A single linear layer in GPT-3's attention has four weight matrices for Q, K, V, and output projections, each of size $12288 \times 12288$, consuming roughly 576 MB each in fp16. Across 96 layers, the attention weights alone exceed 200 GB. The FFN layers add another 150 GB or more. No single device, regardless of generation, can hold this.

The key insight enabling tensor parallelism is that matrix multiplication is separable. Different columns of $A$ contribute to different columns of $Y$ without any cross-dependency, and different rows of $A$ can each be multiplied with the corresponding input slice independently, with results accumulated at the end. This separability is the mathematical foundation for both column-parallel and row-parallel splitting strategies. The question is not whether the computation can be split, but how to split it such that the communication required to assemble final results is minimized and falls at convenient points in the computation graph.

The choice of splitting direction is not arbitrary. Different layers and positions in the transformer have different mathematical structures, and the clever design of tensor parallelism is that it exploits these structures to pair column splits with row splits in a way that eliminates intermediate communication. The column-parallel output of one layer is exactly the right input format for a row-parallel next layer, creating a natural pipeline that requires only a single all-reduce per sub-layer rather than a communication step between every layer.