02 — Tensor Parallelism¶

The paper you must own: Megatron-LM 1 (Shoeybi et al., 2019, arxiv 2507.20534). Read it three times: once for intuition, once with a pen deriving the shapes, once with the code open. Then Megatron-LM 2 (Narayanan/Shoeybi 2021, arxiv 2310.01889) for composition with pipeline parallelism, and Megatron-LM 3 (Korthikanti et al., 2022 — “Reducing Activation Recomputation”) for sequence parallelism inside TP.

The core observation¶

A transformer block is dominated by four matmuls:

QKV projection: X (b, s, h) @ W_qkv (h, 3h) → QKV (b, s, 3h)
Attention output projection: A (b, s, h) @ W_o (h, h) → Y (b, s, h)
MLP up-projection: X (b, s, h) @ W_up (h, 4h) → U (b, s, 4h) (times 2 for SwiGLU)
MLP down-projection: U (b, s, 4h) @ W_down (4h, h) → Y (b, s, h)

Megatron’s trick: pair a column-parallel matmul with a row-parallel matmul so the sharded intermediate stays sharded and only one all-reduce reunites at the end.

Column-parallel vs row-parallel — the derivation¶

Column-parallel¶

Split W (h, d) = [W_1, W_2] along output dim (columns). Rank 0 holds W_1 (h, d/t), rank 1 holds W_2 (h, d/t).

Input X (b, s, h) is replicated on every rank.

Each rank computes Y_i = X @ W_i (b, s, d/t). Output is sharded along dim d.

Forward comm: none if downstream consumer is row-parallel; one all-gather if we need Y replicated.
Backward: gradient of X arrives as dX = dY_1 W_1^T + dY_2 W_2^T, which is a sum across ranks → all-reduce on dX.

Row-parallel¶

Split W (h, d) = [[W_1], [W_2]] along input dim (rows). Rank 0 holds W_1 (h/t, d), rank 1 holds W_2 (h/t, d).

Input must already be sharded along dim h: rank 0 holds X_1 (b, s, h/t).

Each rank computes Y_i = X_i @ W_i (b, s, d). Output is a partial sum on every rank.

Forward: all-reduce over Y to sum the partials.
Backward: dX comes out correctly sharded, no further comm.

The MLP pairing (this is where you feel the elegance)¶

X (replicated, h) 
  → [column-parallel W_up: h → 4h]  → U_sharded (b, s, 4h/t)   [no comm]
  → GeLU / SwiGLU (elementwise on shard)                       [no comm]
  → [row-parallel W_down: 4h → h]    → partial Y (b, s, h)      [no comm]
  → all-reduce                       → Y (replicated, h)         [ONE all-reduce]

That is one all-reduce for the entire MLP block, not two. GeLU/SwiGLU is elementwise so it commutes with the sharding along the output dim of W_up.

The attention pairing¶

Same trick, using the natural head-parallelism of attention:

X (replicated) 
  → [column-parallel W_qkv]  → QKV shards, each rank owns h_heads/t heads
  → attention on local heads (fully independent — heads don't talk!)
  → [row-parallel W_o]        → partial Y
  → all-reduce                 → Y (replicated)                    [ONE all-reduce]

Heads shard cleanly because attention within a head is self-contained. GQA/MQA change nothing structural — you shard by KV-group rather than by head, but the pattern is identical.

The count: two all-reduces per block per direction¶

Forward: 1 all-reduce (attention) + 1 all-reduce (MLP) = 2 all-reduces.
Backward: 2 more (one at each residual join, mirror of forward).
Total per training step per block: 4 all-reduces.
Total per decode token per block: 2 all-reduces.

Each all-reduce touches b × s × h × bytes payload (or b × 1 × h for a decode step).

Concrete numbers for Llama-3-70B, TP=4, decode batch 1, bf16¶

h = 8192, layers = 80.
Payload per all-reduce: 1 × 1 × 8192 × 2 = 16 KB.
All-reduces per token: 2 × 80 = 160.
Traffic per token: ~2.5 MB per rank pair, ring-adjusted ~5 MB per token across the ring.
At 50 tok/s: ~250 MB/s sustained NVLink traffic per rank. Trivial (H100 NVLink SXM = 900 GB/s). Latency, not bandwidth, is the killer at small batch because you pay ring-hop latency 160 times per token.

That last sentence is the whole reason TP at small batch shows sub-linear scaling. Even with fat NVLink, 160 all-reduce syncs per token add microseconds that compound. At batch 32 the bandwidth begins to matter and TP scales better.

Sequence parallelism (Megatron-3) — the free memory win¶

Between the two all-reduces, there are LayerNorm and dropout operations that are not sharded in vanilla TP — they operate on the full-hidden replicated activation. Megatron-3’s observation: shard those along the sequence dimension instead. The all-reduces become reduce-scatter → (SP region) → all-gather, no extra comm cost, but activation memory drops by t×. This is why Megatron’s modern default is TP+SP together.

Do not confuse this “SP inside TP” with Ring Attention / Context Parallelism (file 06), which shards attention across sequence for long-context training. Same word, different technique.

When to use TP¶

Yes:

Model doesn’t fit on one GPU and you have an NVLink domain.
Serving a 70B/405B where inter-GPU latency is fine (NVSwitch).
Latency-sensitive inference at small batch — TP splits the memory-bandwidth-bound decode work across HBMs.

No:

Across PCIe. Just no.
Across IB. All-reduce per token collapses.
Beyond NVLink domain size. On H100 SXM8, TP=16 spans two nodes and dies.
When you could just DP-replicate a smaller TP config for the same total memory but higher throughput.

The critical inference-time subtlety: KV cache also shards¶

Under TP, each rank owns heads/t attention heads, and therefore heads/t KV cache heads. The KV cache is naturally split — you don’t pay a memory tax for it. But if you use GQA with KV-group count g, then TP is capped at g (you can’t shard fewer than 1 KV head per rank). Llama-3-70B has g=8, so TP is capped at 8. This is a deployment-time gotcha.

When you’ll first hit this in code¶

Megatron-LM (github.com/NVIDIA/Megatron-LM): reference implementation. Read megatron/core/tensor_parallel/layers.py.
PyTorch tensor_parallel in torch.distributed (built into 2.4+): user-facing DTensor-based API.
vLLM and SGLang: TP is what you flip on with --tensor-parallel-size N. Under the hood they use custom all-reduce kernels (see FlashInfer / vllm’s custom AR) optimized for small-payload low-latency intra-node.
DeepSpeed AutoTP: automatic TP for inference — easier config, less control.

Exercises¶

On paper, derive the MLP pairing: start from X replicated on 4 ranks, follow every shape and comm operation, arrive at Y replicated. Verify only ONE all-reduce.
Derive the backward pass for column-parallel: given dY_i sharded, show why dX requires an all-reduce.
For a hypothetical model with h=4096, layers=32, seq=2048, batch=1, decode, compute total TP=4 all-reduce traffic per token in bytes.
Explain: why does GQA with 8 KV heads cap TP at 8 for Llama-3-70B?

Papers¶

Megatron-LM 1: arxiv:2507.20534 — the original.
Megatron-LM 2: arxiv:2310.01889 — 3D composition, interleaved pipeline.
Megatron-LM 3 / Reducing Activation Recomputation: Korthikanti et al. 2022 — sequence parallelism inside TP.
Efficient Large-Scale Language Model Training on GPU Clusters (companion paper) for perf accounting on 3072 GPUs.