Tensor Parallelism (TP)

Definition

Tensor parallelism (Megatron-LM, Shoeybi et al., arXiv:1909.08053) is intra-layer model parallelism: it partitions the weight matrix of a single GEMM across $t$ GPUs so each holds a slice, computes a partial result, and the group combines them with one all-reduce per block. Unlike ZeRO/FSDP (which shard state but every GPU runs the same compute), TP splits the computation itself — putting a collective on the critical path inside each layer.

Key math / How it works

For $Y=XA$ with $A\in {\mathbb{R}}^{k\times n}$, two ways to split:

• Column-parallel: $A=[A_1|\cdots |A_t]$, each GPU computes $Y_i=X{A}_i$ (owns output columns; $X$ replicated; no comms to produce $Y_i$).
• Row-parallel: $B=[B_1;\cdots ;B_t]$ with input split to match; each GPU computes a partial sum $Z_i=Y_i{B}_i$, and $Z={\sum }_i{Z}_i$ → all-reduce.

The column→row pairing (the key trick). The MLP block $Z=\text{GeLU}(XA)B$ makes $A$ column-parallel, $B$ row-parallel. The column-split first GEMM leaves the intermediate sharded along the $4h$ dimension — exactly the layout the row-parallel second GEMM consumes. GeLU is elementwise, so $\text{GeLU}([Y_1|Y_2])=[\text{GeLU}(Y_1)|\text{GeLU}(Y_2)]$ commutes with the sharding and applies locally with no comms. Only the final output is a partial sum → one all-reduce. Swapping the order (row-first) would force an all-reduce before the nonlinearity (to form the correct pre-activation) plus one after — two collectives.

Megatron formalizes this with conjugate operators: $f$ (identity fwd / all-reduce bwd) and $g$ (all-reduce fwd / identity bwd). Per transformer layer: 2 all-reduces forward + 2 backward (one per block, MLP and attention).

Attention is parallel over heads. Q/K/V projections are column-parallel (each GPU gets whole heads), per-head attention+softmax runs locally (FlashAttention composes cleanly), the output projection $W_O$ is row-parallel closing with the all-reduce. Constraints: num_attention_heads and hidden_size must be divisible by $t$.

Comms cost. Each all-reduce is over a $bsh$ activation tensor. Per layer (bf16, ring all-reduce, 4 collectives):
$V_{\text{TP/layer}}=4\times 2\frac{t-1}{t}(bsh)\times 2\text{B}=16\frac{t-1}{t}bsh\text{bytes}$
This is exposed on every layer’s critical path with no compute to hide it, and roughly independent of $t$ for large $t$. Example ($h=s=8192$, $b=1$, $t=8$): ~896 MB/GPU/layer → ~1.0 ms over NVLink (900 GB/s) but ~18 ms over IB (50 GB/s).

TP does not scale across nodes. Keep $\text{TP}\le 8$ (one NVLink island). Crossing the node boundary onto ~18×-slower InfiniBand tanks throughput.

Why it matters

TP shards both parameter and activation memory and keeps GEMMs large, but its bandwidth appetite pins it inside a node. It is the innermost, chattiest axis of N-D parallelism, almost always paired with sequence parallelism for free activation savings. Notably DeepSeek-V3 dropped TP entirely on bandwidth-limited H800 NVLink, leaning on EP+PP+FP8 instead.

Taught in

• 10-model-parallelism — §3 the column→row pairing, head-split attention, comms cost.

See also

• megatron-lm — the framework that introduced TP
• sequence-parallelism — the free activation-memory companion
• pipeline-parallelism — the across-layers alternative
• gemm — the matmul being split
• collective-communication — the per-layer all-reduce
• 3d-parallelism — TP as the innermost axis
• gpu-systems-for-llms