PidokuInfra
03Tensor Parallelism
On this page

Tensor Parallelism

Advanced 2h Difficulty 4/5 Topic 03 of 12

Prerequisites III.09, 01, 02

★ The dominant form of parallelism in LLM inference. Understand the layout and you understand why it needs NVLink.

1. What is it?#

Split each weight matrix across GPUs so every GPU computes part of every layer.

One GPU:                    TP=2:
  y = x @ W                   GPU 0: y0 = x @ W[:, :h/2]
  W is (d, h)                 GPU 1: y1 = x @ W[:, h/2:]
                              then concatenate or reduce

Every GPU processes every token, holding a slice of every weight. This is different from pipeline parallelism (which splits by layer) and data parallelism (which replicates).

Diagram — One layer, sharded across two GPUs#

flowchart TB
  X["Layer input x<br/>replicated on every GPU"] --> G0
  X --> G1
  subgraph G0["GPU 0"]
    A0["Half the attention heads<br/>half of the MLP"]
  end
  subgraph G1["GPU 1"]
    A1["Other half of the heads<br/>other half of the MLP"]
  end
  G0 --> AR(("all-reduce<br/>sum"))
  G1 --> AR
  AR --> Y["Full layer output on every GPU"]
  Y -.->|"two all-reduces per layer, for every token"| X

  class A0,A1 compute
  class AR warn
  class X,Y neutral

2. Problem → Why → Optimization#

PROBLEM   A 70B model at FP16 is 141 GB — doesn't fit on an 80 GB GPU.
          And even if it did, 141 GB / 3.35 TB/s = 42 ms per token is too slow.
WHY       Both memory capacity and memory bandwidth are per-GPU resources.
OPTIMIZE  Split the weights across N GPUs. Each reads 1/N of the bytes,
          so decode is N× faster, and each holds 1/N of the weights.

TRADE-OFFS
  ✓ latency improves nearly N×
  ✓ capacity: the model and its KV split across GPUs
  ✗ an AllReduce after every attention block and every FFN block
  ✗ scaling efficiency 70-90% per doubling, and it degrades
  ✗ requires a fast interconnect (NVLink) to be worthwhile
  ✗ all GPUs must be healthy; a failure takes down the whole group

WHEN TO USE   The model doesn't fit, or you need lower latency, AND the GPUs
              are NVLink-connected.
WHEN NOT TO   You only need throughput (use DP), or the GPUs are PCIe-only.

3. Simple analogy#

Four accountants splitting one ledger by column.

Each takes columns 1-25, 26-50, 51-75, 76-100. Each processes every transaction, but only their columns. Four times faster per transaction — if they can combine their partial results quickly.

The combining step is the AllReduce. If they’re in the same room (NVLink), it’s a moment. If they’re in different buildings (PCIe or network), the combining dominates and four accountants are slower than one.

4. The layout — this is the whole technique#

Megatron-style TP splits each layer so that exactly one AllReduce is needed per block.

The FFN#

FFN: y = down( act( gate(x) ) * up(x) )

  gate: (d, d_ff)   split by COLUMN  → each GPU gets (d, d_ff/N)
  up:   (d, d_ff)   split by COLUMN  → each GPU gets (d, d_ff/N)
  down: (d_ff, d)   split by ROW     → each GPU gets (d_ff/N, d)

Per GPU:
  gate_i = x @ W_gate[:, i-th slice]       (d, d_ff/N)  — no communication
  up_i   = x @ W_up[:, i-th slice]                       — no communication
  h_i    = silu(gate_i) * up_i                           — elementwise, local
  y_i    = h_i @ W_down[i-th slice, :]      (d_ff/N, d)  — PARTIAL result
  ────────────────────────────────────────────────────────────────
  y = AllReduce(y_i)                        ← ONE AllReduce for the whole FFN

Why column-then-row is the magic: splitting the first matrix by column means each GPU produces a slice of the intermediate that it can process locally through the activation. Splitting the second by row means each GPU produces a partial sum of the full output, and summing partial sums is exactly what AllReduce does.

If you split both by column, you’d need an AllGather in the middle and handle the elementwise activation across GPUs. The column-then-row pattern avoids that entirely.

Attention#

  W_q, W_k, W_v: split by COLUMN (i.e. by HEAD)
      GPU i gets heads [i·h/N : (i+1)·h/N]
  W_o:           split by ROW

Per GPU:
  q_i, k_i, v_i for its own heads          — no communication
  attention over its own heads              — no communication
                                              (the KV cache is also split by head!)
  o_i = attn_out_i @ W_o[i-th rows, :]      — partial result
  ────────────────────────────────────────────────────────────────
  out = AllReduce(o_i)                      ← ONE AllReduce for attention

Attention parallelizes beautifully across heads because heads are independent. The KV cache splits with them, which is why TP also multiplies your KV capacity.

The full picture#

Per transformer layer:
  RMSNorm (replicated — cheap, all GPUs compute it)
  attention with split heads → AllReduce (size: batch × seq × d)
  residual add (local)
  RMSNorm (replicated)
  FFN with column/row split → AllReduce (size: batch × seq × d)
  residual add (local)

→ 2 AllReduces per layer × L layers.
  Llama-3-70B: 80 layers → 160 AllReduces per forward pass.

160 collectives per token. That’s why the interconnect matters so much.

The embedding and LM head#

Embedding:  split by VOCABULARY. Each GPU holds V/N rows.
            Lookup: each GPU gathers what it has, then AllReduce
            (or each GPU has all embeddings — they're relatively small).

LM head:    split by VOCABULARY (column split of (d, V)).
            Each GPU produces logits for V/N tokens.
            For sampling: either AllGather the logits (expensive, V is large)
            or do a distributed argmax/sample (each GPU finds its local best,
            then a small AllReduce of (value, index) pairs).
            → distributed sampling is much cheaper. Good engines do it.

5. The communication cost#

Per AllReduce: 2 × (N-1)/N × message_size bytes moved per GPU
               (ring algorithm: reduce-scatter + all-gather)

message_size = batch × seq_len × d × bytes

Llama-3-70B decode, batch 32, TP=8, BF16:
  message_size = 32 × 1 × 8192 × 2 = 512 KB
  per AllReduce per GPU: 2 × (7/8) × 512 KB = 896 KB
  × 160 AllReduces = 143 MB per token

  Over NVLink 4 (900 GB/s bidirectional, ~450 GB/s effective per direction):
      143 MB / 450 GB/s = 0.32 ms
  Over PCIe Gen5 (~50 GB/s):
      143 MB / 50 GB/s = 2.9 ms
  
  Compare to the compute+weight-read time per token at TP=8:
      17.6 GB / 3.35 TB/s = 5.3 ms

  NVLink: 0.32/5.3 = 6% overhead     ✓ excellent
  PCIe:   2.9/5.3 = 55% overhead     ✗ terrible

That calculation is the entire argument for NVLink. Do it for your hardware before choosing a TP degree.

Note also: latency, not just bandwidth, matters. Each AllReduce has a fixed latency (~5-10 µs on NVLink, ~20-50 µs over PCIe or network). At 160 AllReduces per token:

NVLink:  160 × 7 µs = 1.1 ms of pure latency
PCIe:    160 × 30 µs = 4.8 ms

For small batches, the fixed latency dominates the bandwidth cost. This is why custom all-reduce implementations (which use one-shot algorithms optimized for small messages) matter at low batch — vLLM and TRT-LLM both ship them.

6. Under the hood — what a TP=4 deployment looks like#

4 processes, one per GPU, in lockstep:

  rank 0 (also runs the scheduler)
  rank 1, 2, 3 (executors)

Every forward pass:
  1. rank 0's scheduler decides the batch
  2. broadcast the batch metadata to all ranks
  3. all ranks execute the same code path on their own slice
  4. 160 AllReduces, all ranks participating
  5. rank 0 collects the sampled tokens and returns them

If ANY rank hangs or crashes, all ranks hang.
NCCL has no fault tolerance: a dead rank means a dead group.

This lockstep property has operational consequences: you cannot drain one GPU of a TP group, you cannot roll them independently, and a single GPU fault kills the group. Set NCCL timeouts and have supervision that restarts the whole group.

7. Performance#

Llama-3-70B, H100 with NVLink, batch 32:

TP   Weights/GPU   ITL (theory)  ITL (real)  Throughput   Efficiency
1    141 GB        —             OOM         —            —
2    70.6 GB       21.1 ms       26 ms       1,230 tok/s  —
4    35.3 GB       10.5 ms       14 ms       2,290 tok/s  93%
8    17.6 GB        5.3 ms        8.5 ms     3,760 tok/s  82%
16   8.8 GB         2.6 ms        6.2 ms     5,160 tok/s  56%  ← degrading

Efficiency = (throughput at TP=N) / (N/2 × throughput at TP=2)

Efficiency degrades because the communication cost is fixed per AllReduce regardless of N (the message size shrinks as (N-1)/N but the latency doesn’t), while the compute per GPU keeps shrinking. Past TP=8, you’re paying more in communication than you gain in bandwidth.

Practical guidance: TP ≤ 8, within one NVLink domain. For larger models, combine TP=8 with pipeline parallelism across nodes (file 04) rather than TP=16.

8. Production implications#

  • TP degree must divide num_key_value_heads (or the implementation replicates KV heads, wasting memory). A model with 8 KV heads cleanly supports TP ∈ {1,2,4,8}.
  • Keep TP groups within an NVLink domain. Check nvidia-smi topo -m.
  • Use the engine’s custom all-reduce if available — it beats NCCL for the small messages typical of decode. vLLM: enabled by default when possible; check the logs.
  • Set NCCL timeouts (NCCL_TIMEOUT, default is long). A hung rank should fail fast, not hang forever.
  • Supervise the whole group. If one rank dies, restart all of them.
  • TP=8 is the practical ceiling for a single node. Beyond that, add PP or DP.
  • Distributed sampling matters. Verify your engine doesn’t AllGather full logits every step — at V=128k and batch 64 that’s 16 MB per step per GPU.

9. Common mistakes#

TP across PCIe. 50%+ overhead. Check the topology.

TP degree not dividing the KV head count. Silent memory waste or failure.

TP=16 or higher. Efficiency collapses; use TP=8 × PP or DP.

Not setting NCCL timeouts. A hung rank hangs forever, holding GPUs.

Assuming linear scaling. It’s 80-93% per doubling up to TP=8, then worse.

Forgetting the fixed latency component. At small batch, AllReduce latency (not bandwidth) dominates.

Draining one GPU of a TP group. Not possible; the group is atomic.

10. Hands-on exercise#

A. Derive the layout. On paper, work through the FFN split for TP=2 with d=4, d_ff=8. Show that column-split-then-row-split requires exactly one AllReduce, and that column-then-column would require two collectives.

B. Measure the scaling. Run a model at TP=1, 2, 4, 8 (as far as your hardware allows). Measure ITL and throughput. Compute the efficiency at each degree. Compare to section 7.

C. Measure the AllReduce cost. Use nccl-tests (all_reduce_perf) to measure AllReduce bandwidth and latency for message sizes from 4 KB to 64 MB on your hardware. Plot both. Where’s the latency-bound regime?

D. Compute the overhead. Using your measured numbers, compute the predicted communication overhead for your model at TP=4 and TP=8, at batch 1 and batch 64. Compare to the measured difference between TP degrees.

E. Topology experiment. If you have both NVLink-connected and PCIe-only GPU pairs, run TP=2 on each and compare. Quantify the interconnect’s effect.

11. Interview questions#

  1. Explain the Megatron column/row split for the FFN and why it needs only one AllReduce.
  2. How is attention split for TP, and what happens to the KV cache?
  3. How many collectives per token for an 80-layer model at TP=8?
  4. Compute the communication overhead for TP=8 on NVLink vs PCIe. Show your work.
  5. Why does TP scaling efficiency degrade past TP=8?
  6. Why must the TP degree divide the number of KV heads?
  7. What are the operational consequences of TP’s lockstep execution?
  8. Why does the fixed AllReduce latency matter more at small batch?

12. Further reading#

  • [ESTABLISHED] Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism” (2019) — the layout comes from here
  • [ESTABLISHED] Pope et al., “Efficiently Scaling Transformer Inference” (2022) — the inference partitioning analysis
  • [REFERENCE] NCCL documentation and nccl-tests
  • [REFERENCE] vLLM distributed/ and its custom all-reduce implementation
  • Next: 04 — Pipeline parallelism
esc

↑↓ navigate↵ openesc close