Tensor Parallelism › Column Parallelism

Column parallelism splits a weight matrix along its output (column) dimension. Given $N$ GPUs with indices $0, 1, \ldots, N-1$, we partition the weight matrix $A \in \mathbb{R}^{d \times h}$ into $N$ column slices:

$$A = [A_0 \mid A_1 \mid \cdots \mid A_{N-1}]$$

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

Each GPU $i$ stores its slice $A_i$ and computes:

$$Y_i = X A_i$$

where:

• $X \in \mathbb{R}^{b \times d}$ is the full input, replicated on all GPUs
• $A_i \in \mathbb{R}^{d \times (h/N)}$ is the column shard on GPU $i$
• $Y_i \in \mathbb{R}^{b \times (h/N)}$ is the partial output on GPU $i$

The partial outputs $Y_i$ together form the full output: concatenating $[Y_0, Y_1, \ldots, Y_{N-1}]$ along the column dimension recovers the full $Y = XA \in \mathbb{R}^{b \times h}$. This is correct because matrix multiplication with a partitioned weight matrix produces a partitioned result:

$$XA = X[A_0 \mid A_1 \mid \cdots \mid A_{N-1}] = [XA_0 \mid XA_1 \mid \cdots \mid XA_{N-1}]$$

The critical observation is that each GPU's computation $Y_i = XA_i$ is completely independent. No communication is needed during the forward pass to produce $Y_i$; each device can compute its shard without knowledge of what the other devices are doing. The price is that the input $X$ must be fully available on every GPU before computation begins. In practice, this is handled by an all-gather operation at the start of the layer or by keeping the input in the replicated state (which happens automatically when tensor parallelism is paired with data parallelism).

Why is replicating the input acceptable? Because the input to each transformer layer is the residual stream, which after an all-reduce from the previous layer is identical on all GPUs. There is no communication cost to have a replicated input in steady-state operation, since the previous layer's row-parallel all-reduce already distributed it.

Column parallelism is natural for the query, key, and value projections in multi-head attention. Recall that multi-head attention splits the model's hidden dimension into $N_h$ heads, where each head operates on a $d_k = d / N_h$ dimensional subspace. When you have $N$ GPUs and split the weight matrix along the head dimension, each GPU naturally handles $N_h / N$ attention heads. The column-parallel split corresponds exactly to distributing heads across devices.

This alignment between tensor parallelism and multi-head attention is not a coincidence. The original Megatron-LM paper was designed with this correspondence in mind. Because attention heads are completely independent of one another during the self-attention computation (they attend to the full sequence in parallel subspaces without interacting), you can assign entire heads to individual GPUs and have each GPU compute attention for its assigned heads without ever communicating with other GPUs during the attention computation itself. The interaction between heads happens only when their outputs are projected back to the model's hidden dimension, and that projection is handled by the row-parallel output projection.

For an attention layer with hidden size $d = 4096$ and 32 heads on 4 GPUs, each GPU holds the Q, K, V projection weights for 8 heads: matrices of size $4096 \times 1024$ rather than $4096 \times 4096$. Each GPU attends over its 8 heads independently and produces attention outputs of size $b \times 1024$. The output projection then acts on these 1024-dimensional features row-parallel, accumulating the full 4096-dimensional output across GPUs via all-reduce.

The feed-forward network (FFN) also benefits from column parallelism. A standard transformer FFN applies two linear transformations with a nonlinearity between them:

$$\text{FFN}(X) = \sigma(XW_1 + b_1) W_2 + b_2$$

where:

• $X \in \mathbb{R}^{b \times d}$ is the input batch with $b$ tokens and hidden dimension $d$
• $W_1 \in \mathbb{R}^{d \times 4d}$ is the expansion weight matrix, increasing the dimension by a factor of 4
• $b_1 \in \mathbb{R}^{4d}$ is the expansion bias
• $\sigma(\cdot)$ is the activation function (GeLU in most modern transformers, ReLU in older ones)
• $W_2 \in \mathbb{R}^{4d \times d}$ is the contraction weight matrix, reducing back to $d$
• $b_2 \in \mathbb{R}^{d}$ is the output bias

Column parallelism applied to $W_1$ splits the expansion across GPUs. GPU $i$ receives the full input $X$ and computes $XW_{1,i}$ using its column shard $W_{1,i} \in \mathbb{R}^{d \times d}$ (for a 4-way split of a $4d$-dimensional FFN). The activation function is then applied locally because GeLU and ReLU operate element-wise on each dimension independently. This is a subtle but important point: element-wise nonlinearities are trivially parallelizable because each output element depends on exactly one input element, which lies entirely within the local shard.

After the nonlinearity, each GPU holds $\hat{Z}_i = \sigma(XW_{1,i}) \in \mathbb{R}^{b \times d}$. These shards together form the full intermediate activation $\hat{Z} = [\hat{Z}_0 \mid \hat{Z}_1 \mid \cdots \mid \hat{Z}_{N-1}] \in \mathbb{R}^{b \times 4d}$, ready to be consumed by the row-parallel second layer.