Tensor Parallelism › Worked Example: Column-Parallel FFN

Let's trace through a specific numerical example to make the computation concrete. Suppose we have a feed-forward network with input dimension $d = 4$, hidden dimension $h = 8$, and we are using $N = 2$ GPUs.

The full weight matrix $W_1 \in \mathbb{R}^{4 \times 8}$ is split column-wise:

$$W_1 = [W_{1,0} \mid W_{1,1}]$$

where $W_{1,0} \in \mathbb{R}^{4 \times 4}$ lives on GPU 0 and $W_{1,1} \in \mathbb{R}^{4 \times 4}$ lives on GPU 1.

Given an input $X \in \mathbb{R}^{2 \times 4}$ (batch of 2 tokens, each with 4 features), both GPUs have the full input. GPU 0 computes:

$$Z_0 = X W_{1,0} \in \mathbb{R}^{2 \times 4}$$

and GPU 1 computes:

$$Z_1 = X W_{1,1} \in \mathbb{R}^{2 \times 4}$$

After the GeLU nonlinearity applied locally on each GPU, the activations $\hat{Z}_0$ and $\hat{Z}_1$ together form the full intermediate representation $[\hat{Z}_0 \mid \hat{Z}_1] \in \mathbb{R}^{2 \times 8}$. Neither GPU needed to communicate with the other to reach this point.

The second weight matrix $W_2 \in \mathbb{R}^{8 \times 4}$ is split row-wise:

$$W_2 = \begin{bmatrix} W_{2,0} \\ W_{2,1} \end{bmatrix}$$

where $W_{2,0} \in \mathbb{R}^{4 \times 4}$ on GPU 0 and $W_{2,1} \in \mathbb{R}^{4 \times 4}$ on GPU 1.

GPU 0 multiplies its activation shard by its weight shard: $Y_0 = \hat{Z}_0 W_{2,0} \in \mathbb{R}^{2 \times 4}$. GPU 1 multiplies its activation shard by its weight shard: $Y_1 = \hat{Z}_1 W_{2,1} \in \mathbb{R}^{2 \times 4}$.

The final output is obtained by all-reducing (summing) across GPUs:

$$Y = Y_0 + Y_1 \in \mathbb{R}^{2 \times 4}$$

This is identical to the full computation $Y = \text{GeLU}(XW_1) W_2$, but computed with weight matrices that fit on individual GPUs. The only communication is the final all-reduce that sums $Y_0$ and $Y_1$.

What makes this example particularly illuminating is the absence of intermediate communication. Between the column-parallel step and the row-parallel step, GPUs 0 and 1 proceed completely independently. GPU 0 never needs to know what $Z_1$ looks like. GPU 1 never needs to know what $Z_0$ looks like. They only need to coordinate once, at the very end, to produce the final output. This one communication per two linear layers is the core efficiency of Megatron-style tensor parallelism.