Tensor Parallelism › Row Parallelism

Row parallelism splits a weight matrix along its input (row) dimension. Given the weight matrix $A \in \mathbb{R}^{d \times h}$, we partition along the row dimension into $N$ slices:

$$A = \begin{bmatrix} A_0 \\ A_1 \\ \vdots \\ A_{N-1} \end{bmatrix}$$

where each $A_i \in \mathbb{R}^{(d/N) \times h}$.

Each GPU $i$ stores its row shard $A_i$ and receives a corresponding slice $X_i \in \mathbb{R}^{b \times (d/N)}$ of the input, then computes:

$$Z_i = X_i A_i$$

where:

• $X_i \in \mathbb{R}^{b \times (d/N)}$ is the input shard on GPU $i$
• $A_i \in \mathbb{R}^{(d/N) \times h}$ is the row shard on GPU $i$
• $Z_i \in \mathbb{R}^{b \times h}$ is the partial result on GPU $i$

The full output requires summing the partial results:

$$Y = XA = \sum_{i=0}^{N-1} X_i A_i$$

This is an all-reduce operation: each device produces a partial sum over the full output dimension $h$, and summing these partial results gives the correct final output. Unlike column parallelism, where the output shards can be consumed independently (each GPU's shard feeds into a corresponding downstream shard), row parallelism requires synchronization to assemble the full output.

The correctness of this decomposition follows directly from the block form of matrix multiplication. If we write the full input as $X = [X_0 \mid X_1 \mid \cdots \mid X_{N-1}]$ (conceptually partitioned along the column dimension) and the weight matrix as stacked row blocks, then:

$$\begin{aligned} XA &= [X_0 \mid X_1 \mid \cdots \mid X_{N-1}] \begin{bmatrix} A_0 \\ A_1 \\ \vdots \\ A_{N-1} \end{bmatrix} \\ &= X_0 A_0 + X_1 A_1 + \cdots + X_{N-1} A_{N-1} \end{aligned}$$

Each term $X_i A_i$ is a partial contribution to the full output, and summing them gives the correct result. Note that each partial result $Z_i = X_i A_i$ has the full output shape $\mathbb{R}^{b \times h}$ rather than a shard of it. This is the defining feature of row parallelism: the partial results are over the full output space, and correctness requires their sum, not their concatenation.

The design of row parallelism makes most sense when you see it as the second half of a pair. The output of a column-parallel layer is a feature-sharded tensor, where GPU $i$ holds the columns $[i \cdot h/N, (i+1) \cdot h/N)$ of the full hidden state. This is precisely the input format that row parallelism expects: GPU $i$ should receive input shard $X_i$ corresponding to its row shard $A_i$. When you place a row-parallel layer directly after a column-parallel layer, the output of the first layer is already distributed in exactly the right way for the second layer to consume without any additional communication.

This chaining property is what makes tensor parallelism efficient. Without it, you would need an all-gather after the column-parallel layer to reassemble the full activation before the row-parallel layer, adding an extra communication step. With careful layer pairing, the communication is deferred to the very end, where a single all-reduce assembles the final output.