The problem: one matmul, too big for one GPU
Data parallelism replicates the whole model and splits the batch; it does nothing when a single layer’s weights — or the activations they produce — do not fit in one device’s memory. Pipeline parallelism splits the model by depth, but each layer still lives whole on one device. Tensor parallelism attacks the remaining axis: it splits an individual matrix multiply across GPUs.
Consider the core operation Y = X A, with input activations X: [T, d] (T tokens, model width d) and a weight A: [d, k]. A modern feed-forward weight is [d, 4d] — for d = 4096 that is [4096, 16384], ~134M parameters emitting a hidden activation [T, 16384]. Stack dozens of layers and neither weights nor activations fit on one accelerator. The question tensor parallelism answers: how do you cut A into p pieces, one per GPU, so each device does 1/p of the work — and how much must the pieces talk to reassemble the right answer?
Column-parallel: split the weight by output columns
The first way to shard Y = X A is to cut A along its columns (its output dimension). Write A = [A_1 | A_2 | … | A_p], where each A_i: [d, k/p] lives on GPU i. Because matrix multiplication distributes over a column partition:
Y = X A = X [A_1 | A_2 | ... | A_p]
= [X A_1 | X A_2 | ... | X A_p]
= [Y_1 | Y_2 | ... | Y_p] Y_i = X A_i : [T, k/p]Each GPU already holds the full input X (replicated), computes its own output slice Y_i locally, and no communication is needed in the forward pass — each device simply keeps its slice. The result is left sharded along the feature dimension, ready for a next operation that consumes a sharded input. The only synchronization hides in the backward pass: the gradient with respect to the shared input is ∇X = Σ_i ∇Y_i A_i^T, a sum across GPUs that requires an all-reduce.
Row-parallel: split the weight by input rows
The complementary shard cuts a weight B: [k, d] along its rows (its input dimension). This is the natural partner for a column-parallel output, because a feature-sharded activation Y = [Y_1 | … | Y_p] lines up row-for-row with B = [B_1; B_2; … ; B_p], where each B_i: [k/p, d] sits on GPU i. Block matrix multiplication then gives a sum, not a concatenation:
Z = Y B = [Y_1 | ... | Y_p] [B_1; ... ; B_p]
= Y_1 B_1 + Y_2 B_2 + ... + Y_p B_p
= Z_1 + Z_2 + ... + Z_p Z_i = Y_i B_i : [T, d]Now every GPU computes a partial output Z_i of the full shape [T, d], and the correct answer is their element-wise sum. That sum is the all-reduce: after the local matmuls, an all-reduce adds the p partials together and leaves the complete Z replicated on every device. So row-parallel is the mirror image of column-parallel — it needs the all-reduce in the forward pass and only an identity in the backward pass.
The f and g conjugate operators
Megatron formalizes the two communication points as a conjugate pair of operators inserted into the graph. f is identity in the forward pass and an all-reduce in the backward pass; g is an all-reduce in the forward pass and identity in the backward pass. A column-parallel layer is preceded by f (so its input gradient gets reduced); a row-parallel layer is followed by g (so its output gets reduced).
The payoff of pairing them is that the region between f and g needs no synchronization at all: a replicated X through f (no-op forward), a column-parallel weight producing a sharded hidden state, a per-element nonlinearity applied locally on the shard, a row-parallel weight producing partial outputs, then g to all-reduce them into a replicated result. One all-reduce forward (in g) and one backward (in f) bracket the entire sandwich — the structural trick that makes the scheme cheap: two large matmuls, one round of communication.
The MLP block: column then row, one all-reduce
The transformer feed-forward block is Z = GeLU(X A) B with A: [d, 4d] and B: [4d, d]. Make A column-parallel and B row-parallel and the pieces interlock: GPU i computes H_i = GeLU(X A_i) with A_i: [d, 4d/p], a hidden shard [T, 4d/p].
The nonlinearity is why the order matters. GeLU acts element-wise, so it applied to a column-shard equals the corresponding column-shard of GeLU(XA) — each GPU runs the activation on its own slice with no cross-talk. Had we sharded the other way, the nonlinearity would sit on a partial sum and GeLU(Z_1 + Z_2) ≠ GeLU(Z_1) + GeLU(Z_2), forcing an all-reduce before the activation. Column-then-row avoids that. Then Z = Σ_i H_i B_i with row-parallel B_i: [4d/p, d], and the trailing g all-reduces the partial Z_i into the final replicated output. The whole MLP: one all-reduce forward, one backward, with d_ff split p ways in between.
Attention: parallel across heads
Multi-head attention shards along an axis it already has — the heads. The QKV projection is a column-parallel linear that produces queries, keys, and values, and the natural cut assigns a disjoint subset of heads to each GPU. With h heads and p GPUs, device i owns h/p heads and computes their Q_i, K_i, V_i from the replicated input X.
The beauty is that attention is already independent per head: the softmax(Q K^T / √d_k) V for one head never touches another. So each GPU runs full self-attention for its heads with zero communication — the expensive N×N score matrices stay local — and holds an output shard [T, d/p]. The output projection W_O is made row-parallel over that same head partition, so its partials are summed by a trailing g. Structurally identical to the MLP: column-parallel in (split by heads), row-parallel out, one all-reduce forward, one backward.
Counting all-reduces: two per layer, each way
A transformer layer is an attention block plus an MLP block, each a column-then-row sandwich ending in a g all-reduce. So the forward pass costs two all-reduces per layer — one to recombine attention, one to recombine the MLP. The backward pass mirrors it: each block’s leading f fires an all-reduce on the input gradient, so the backward pass also costs two all-reduces per layer — four collective operations per step.
Multiply through: a 32-layer model runs 32 × 2 = 64 all-reduces in the forward pass alone, each a hard barrier the next matmul must wait on. And unlike data-parallel gradient syncs, which overlap with backward compute, these sit squarely on the critical path — frequent, blocking collectives over full activation tensors, which is what dictates the hardware tensor parallelism demands.