Tensor Parallelism

ZeRO-3 from the last chapter splits the parameters across GPUs, but every GPU still executes the full computation on full-sized tensors — it just gathers the weights it's missing first. Tensor parallelism does something fundamentally different: it splits the math itself, so each GPU only ever computes part of every matrix multiplication and no single GPU ever holds a full activation.

That's a more powerful kind of splitting, and a more expensive one. It synchronizes twice per transformer block, in the critical path, with no way to hide the communication behind compute. Which is why — spoiler for the punch line — tensor parallelism lives and dies by the bandwidth hierarchy from chapter 1, and essentially never crosses a node boundary.

Splitting a matmul, two ways#

Everything in a transformer is $Y = XA$ for some weight $A$. There are exactly two ways to shard $A$ across devices, and the entire art is in choosing which one for each matmul so their communication cancels.

Column parallelism. Split $A$ by columns: $A = [A_1 | A_2 | \cdots | A_p]$. Each device holds a column-slice and the full input $X$, and computes $Y_i = X A_i$ — a column-slice of the output. No communication to compute; a final all-gather if you need the full $Y$ reassembled.

Row parallelism. Split $A$ by rows, which requires $X$ to be split by columns to match: $A = [A_1; A_2; \ldots; A_p]$ and $X = [X_1 | \cdots | X_p]$. Each device computes a partial product $Y_i = X_i A_i$ of the full output shape, and the true output is their sum — an all-reduce.

Each of these needs one collective to produce a usable full-sized output: column needs an all-gather, row needs an all-reduce. So a single matmul costs one collective. But watch what happens when you chain them.

The Megatron trick: column then row#

Here is the insight the source notes skip, and it's the whole reason tensor parallelism is practical. The transformer MLP is two matmuls back to back with a nonlinearity between:

Suppose you shard the first matmul column-wise and the second row-wise. Then:

• First matmul: each device computes $\text{GeLU}(X A_i)$, a column-slice of the hidden activation. GeLU is elementwise, so it commutes with the column split — no communication needed, and this is why the column split has to come first (a row split would need the full pre-activation summed before the nonlinearity).
• Second matmul: that column-slice of the hidden state is exactly the column-slice of input that row parallelism wants. Each device computes $\text{GeLU}(XA_i)\,B_i$, a partial output, and one all-reduce sums them.

The column split's all-gather and the row split's input-scatter annihilate. Two matmuls, and only one all-reduce at the very end.Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019. The column-then-row pairing is the paper's core contribution and it's genuinely elegant — the intermediate never needs to be reassembled because the second layer's sharding was chosen to consume the first layer's sharding directly.

Attention gets the same treatment and it's even more natural: split the Q, K, V projections column-wise by attention head, so each device owns a subset of whole heads and computes their attention entirely locally, then split the output projection row-wise. One all-reduce for the whole attention block.

So a full transformer layer costs two all-reduces in the forward pass (one for attention, one for the MLP) and two more in the backward pass. Four per layer, per step, on the critical path.

class ColumnParallelLinear(nn.Module):
    # Weight A is split by columns across the TP group; each rank holds A_i.
    def forward(self, x):                 # x is replicated across the group
        return x @ self.A_i               # local column-slice of the output
 
class RowParallelLinear(nn.Module):
    # Weight B is split by rows; input arrives already column-split.
    def forward(self, x_i):               # x_i is this rank's slice
        y_partial = x_i @ self.B_i        # full-shape partial sum
        return all_reduce(y_partial)      # <-- the block's only communication
 
# The MLP: column feeds row, and the GeLU in between commutes with the split,
# so nothing is gathered until the final all-reduce.
def mlp(x):
    return RowParallelLinear()(gelu(ColumnParallelLinear()(x)))

Why it never leaves the node#

Now the bandwidth argument, which is short and decisive.

Those four all-reduces per layer are synchronous and unhideable. Unlike DDP's gradient reduction — which overlaps with the backward pass because the gradients trickle out layer by layer — a tensor-parallel all-reduce sits directly between two matmuls that depend on its result. The GPUs compute, stop, communicate, and only then continue. Every byte of that communication is exposed latency.

Each all-reduce moves (from chapter 7) about $2 \cdot \frac{p-1}{p} \cdot X$ bytes where $X$ is the activation size. Over 32 layers, twice per layer, that's a lot of traffic in the critical path. On NVLink at ~450 GB/s it's tolerable and hides partially behind the matmuls' own duration. On InfiniBand at ~50 GB/s — nine times slower — it would dominate, and your expensive tensor cores would spend most of their time idle waiting on the network.

So the rule, which you can now derive rather than memorize: tensor parallelism degree stays ≤ the number of GPUs on one NVLink island — 8 on a standard DGX/HGX node. Beyond that you switch to pipeline parallelism (chapter 9), which communicates far less and tolerates the slow inter-node links. This single constraint shapes every large-scale training layout, and it falls straight out of the chapter 1 hierarchy.

Sequence parallelism: closing the last gap#

Tensor parallelism splits the attention and MLP blocks beautifully, but it leaves two regions replicated: the LayerNorms and the dropouts between blocks, which operate on the full activation and can't be split by the head/column scheme. Every TP rank redundantly stores and computes those full-sized activations — a real chunk of the chapter 2 activation budget, un-sharded.

Sequence parallelism fixes it by splitting those regions along the sequence dimension instead.Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, MLSys 2023 — the same paper as chapter 3's selective recompute. Sequence parallelism and selective recompute are its two halves, and they compose. LayerNorm is independent per token, so token $t$'s normalization needs only token $t$ — splitting the sequence across the TP ranks is free of any cross-token dependency. The elegant part is the handoff at the block boundaries: the transition from sequence-split (norm) to tensor-split (attention/MLP) and back is exactly an all-gather one way and a reduce-scatter the other.

And here's the payoff — from chapter 7, all-reduce = reduce-scatter + all-gather. The all-reduce that TP already does at each block boundary gets replaced by a reduce-scatter + all-gather pair that does double duty: it both performs the reduction TP needed and effects the sequence↔tensor resharding SP needs. Same total communication volume as plain TP, and now the norm/dropout activations are sharded $p$ ways too. It's close to a free reduction in activation memory, which is why it's standard in every serious Megatron-style setup.

Where the memory actually lands#

Tensor parallelism divides three things by the TP degree $p$: the parameters (each rank holds $1/p$ of every weight matrix), the optimizer state on those parameters, and most of the activations. Add sequence parallelism and the last holdouts — norms and dropouts — divide by $p$ as well.

What it does not divide is anything replicated across the TP group, and it does not reduce communication — it converts a memory problem into a bandwidth problem, spending the fastest links you have. That's the trade: TP buys you the ability to hold a layer that no single GPU could, at the cost of the most expensive, least-hideable communication in the whole stack. You use exactly as much of it as one node's NVLink can absorb, and not one GPU more.