Tensor parallelism is the one form of model parallelism that cuts inside a single matrix multiply. When a weight matrix — the 16,384-wide feed-forward projection, the fused QKV projection — is too large for one GPU, Megatron-style tensor parallelism shards that matrix across several GPUs so each holds and computes only its slice. The elegant part is the algebra: split the first linear layer by columns and the second by rows, and the partial results line up so the whole block needs exactly one all-reduce to become correct again. This piece derives the column-parallel and row-parallel identities, composes them into the MLP and attention blocks, counts the all-reduces (two per layer each way), works a numeric example, and explains why all that communication chains you to NVLink inside a single node.
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.
A worked example
Take d = 4096, d_ff = 16384, tensor-parallel degree p = 4, and a micro-batch of T = 4096 tokens in fp16. The column-parallel weight A: [4096, 16384] splits into four [4096, 4096] shards; each GPU produces a hidden shard, applies GeLU locally, then row-parallel B: [16384, 4096] emits a partial Z_i: [4096, 4096] that the all-reduce sums.
activation to all-reduce: [T, d] = [4096, 4096]
bytes (fp16): 4096 * 4096 * 2 = 33.5 MB
ring all-reduce traffic: ~2 * (p-1)/p * size
per GPU, p=4: ~1.5 * 33.5 MB = ~50 MB
per layer (attn + MLP): ~2 * 50 MB = ~100 MB forward
+ another ~100 MB in the backward passSo one layer moves on the order of 200 MB across the interconnect every step. On a 600 GB/s NVLink fabric that is roughly 0.3 ms of pure communication per layer — tolerable. Push the same tensors over a ~25 GB/s PCIe or inter-node link and it balloons past 8 ms per layer, and the blocking collectives swamp the compute. The numbers, not the algebra, are why tensor parallelism lives and dies by interconnect bandwidth.
Why it needs NVLink and stays within a node
The example makes the constraint concrete: activation volume this large, fired twice per layer on the blocking critical path where it cannot hide behind compute, is only cheap on a very fast, very low-latency fabric.
That fabric is NVLink (and NVSwitch), offering hundreds of gigabytes per second of all-to-all bandwidth within a server — often an order of magnitude more than PCIe, and far more than the InfiniBand or Ethernet links between servers. So tensor parallelism is almost always confined to a single node: p is typically 2, 4, or 8, matching the NVLink-connected GPUs in one box. To scale beyond a node you do not raise the tensor-parallel degree; you compose it with pipeline parallelism across nodes and data parallelism across replicas — the ‘3D parallelism’ layout — keeping the bandwidth-hungry tensor-parallel collectives on the fast intra-node wires and the sparser collectives on the slower links.
The compute versus communication balance
Whether tensor parallelism pays off is a ratio. Each block does O(T · d · d_ff) floating-point work but communicates only O(T · d) bytes in its all-reduce — compute grows with the hidden width while communication does not, so wider models and larger token batches amortize the collective better. Sharding a tiny layer is all overhead; sharding a fat one can be nearly free. Break-even also shifts with p: ring all-reduce traffic scales as 2(p-1)/p, climbing toward 2× as p grows, so higher degrees cost proportionally more communication for the same compute per GPU.
The pitfalls follow directly. Do not shard the nonlinearity across a partial sum — keep column-then-row order so GeLU stays local. Do not push p past the node’s NVLink domain, where the barriers fall onto slow links and dominate. And remember it reduces per-GPU memory and compute but adds synchronization; reach for it when a layer genuinely will not fit, not as a first lever for speed. Used inside its budget — a wide model, a fast fabric, a handful of GPUs — it turns an impossibly large matmul into p comfortable ones joined by a single all-reduce.