Tensor Parallelism › Combining Column and Row Parallelism

The real power of tensor parallelism emerges from chaining column-parallel and row-parallel layers such that the output of one feeds naturally into the other, eliminating unnecessary communication.

Consider the two-layer FFN with $W_1$ (expansion) followed by $W_2$ (contraction):

1. Column-parallel $W_1$: GPU $i$ receives the full input $X$ and computes $Z_i = XW_{1,i}$ where $W_{1,i}$ is the $i$-th column shard of $W_1$. After applying the nonlinearity, each GPU holds $\hat{Z}_i = \text{GeLU}(Z_i)$, a shard of the intermediate activations. No communication has occurred at this point.

2. Row-parallel $W_2$: GPU $i$ holds the $i$-th row shard $W_{2,i}$ of $W_2$, where this partitioning is along the input dimension of $W_2$ (which is the same dimension along which $\hat{Z}_i$ is partitioned). Each GPU computes $Y_i = \hat{Z}_i W_{2,i}$, a partial sum over the full output space.

3. All-reduce: The partial sums $Y_i$ are summed across GPUs to produce the full output $Y = \sum_i Y_i$.

The key observation is that the output of column parallelism (sharded along the feature dimension) is exactly the right input format for row parallelism (which expects sharded inputs along the same dimension). This means no communication is needed between the two layers. The only communication happens at the end of the row-parallel layer, in the all-reduce that assembles the final output.

For a two-layer block (FFN or attention), this yields a total of one all-reduce in the forward pass. Without this careful pairing, a naive implementation would require an all-gather after the column-parallel layer (to reassemble the full activations before the second layer), costing an extra communication step. This design choice, seeming minor in the abstract, has a substantial impact on training throughput: every saved all-gather at model scale represents gigabytes of data that do not need to traverse the interconnect on every forward pass.

A subtle but important aspect of the FFN chaining is where the nonlinearity fits. The GeLU (or ReLU) function is applied between the two linear layers. In column-parallel mode, each GPU applies GeLU locally to its activation shard $Z_i$. This works correctly because GeLU is an element-wise function: $\text{GeLU}(z)$ depends only on the scalar value $z$, not on any other element in the tensor. There are no cross-feature dependencies that would require communication.

This would not be the case for operations like softmax or layer normalization, which are reduction operations over features. Softmax requires the full feature vector to compute the normalizing denominator, and layer normalization requires statistics (mean and variance) over the feature dimension. These operations cannot be applied to feature shards independently. This is why tensor parallelism requires careful attention to which operations are tensor-parallelizable and which require full-feature access, a consideration that motivated the sequence parallelism extension discussed later.

Multi-head attention can be decomposed similarly to the FFN. The QKV projections are column-parallel: each GPU handles a subset of attention heads independently. The output projection maps from the concatenated head outputs back to the model dimension, and it is row-parallel.

Why does this work? Because attention heads are completely independent of one another in standard self-attention. Each head attends to the full sequence but produces its output in a separate subspace of the representation. When you distribute heads across GPUs, each GPU attends over its subset of heads without needing to know what other GPUs' heads are doing. The interaction between heads happens only at the output projection layer, which aggregates all head outputs and is naturally handled by the row-parallel reduce.

Concretely, the QKV weight matrices are split column-wise across GPUs, with each GPU holding the Q, K, V weights for its assigned heads. Each GPU computes attention scores and applies them to values independently, since heads do not communicate during the attention computation. The attention output for each GPU's heads is then passed through the row-parallel output projection, which accumulates partial sums across GPUs in an all-reduce.

This design has a pleasing symmetry with the FFN decomposition: both sub-layers require exactly one all-reduce in the forward pass, giving two all-reduces total per transformer layer. The all-reduces happen at the output projection of attention and at $W_2$ of the FFN. Between the two sub-layers, residual connections are added to the accumulated results. Because every GPU holds the full output after each all-reduce, the residual stream is synchronized and identical on all devices before the next sub-layer begins, making sure that the replicated-input assumption for the next column-parallel layer holds automatically.