Skip to content
LLM Inference
S17Distributed Serving·216 lines

Tensor & Pipeline Parallelism

Split each matmul column-wise then row-wise and one all-reduce per sublayer is all the communication you need. Pipeline parallelism trades latency for capacity.

  • column/row sharding
  • all-reduce
  • head splitting
  • pipeline stages
  • micro-batches

Motivation

The problem

A 70B model in fp16 is 140 GB. An H100 has 80 GB, and you still need room for a KV cache. Quantization gets you one size class further; past that, the model has to live on several GPUs.

Splitting is easy. Splitting without spending all the benefit on communication is the hard part, and what connects the GPUs decides the answer.

Core idea

The solution

Two strategies, with different cost structures.

  • Tensor parallelism splits each matrix across GPUs. Every GPU works on every layer, so the model gets faster as well as smaller. The price is a collective in every layer to reassemble the output.
  • Pipeline parallelism gives each GPU a contiguous range of layers. Communication drops to one activation per stage boundary, which is almost nothing. But a request must traverse the stages in order, so latency does not improve and idle bubbles appear.

Tensor parallelism is cheap because of the Megatron sharding pattern: shard the first matmul by column and the second by row, and the whole sublayer needs exactly one all-reduce.

the megatron patternpython
class ShardedFFN:
    def __init__(self, w1, w2, w3, rank, world):
        # W1 and W3: split the OUTPUT dimension -> each rank owns d_ff/world columns
        self.w1 = shard_columns(w1, rank, world)
        self.w3 = shard_columns(w3, rank, world)
        # W2: split the INPUT dimension -> each rank owns d_ff/world rows
        self.w2 = shard_rows(w2, rank, world)

    def forward(self, x):
        h = silu(x @ self.w1) * (x @ self.w3)   # sharded activation, no comms
        y = h @ self.w2                          # PARTIAL sum of the full output
        return all_reduce(y)                     # one collective per sublayer

Attention shards by head, and GQA caps the degree

Each GPU takes a subset of the attention heads and computes them independently; the output projection is row-sharded exactly like W2. The tensor-parallel degree, though, must divide the number of KV heads. A model with 8 KV heads cannot use TP16 without replicating KV heads across ranks, and that regularly catches out people building 16-GPU deployments of Llama-3-70B.

Mechanics

How it works

Diagramtensor sharding and pipeline stages
TENSOR PARALLELISM — COLUMN-SHARD, THEN ROW-SHARD, ONE ALL-REDUCEx[b, d]W1 · COLUMN-SHARDEDGPU0GPU1GPU2GPU3each GPU owns d_ff/4 columns → no communication neededactivation is sharded, and that is fineSwiGLU applies elementwise — still no communicationW2 · ROW-SHARDEDGPU0GPU1GPU2GPU3each GPU produces a PARTIAL sum over the full output→ all-reduce here, once per sublayerall-reduce (sum across 4 GPUs)blockingTwo matmuls, one collective.Attention shards the sameway, by head.TP degree must divide thenumber of KV heads — GQAwith 8 KV heads caps TP at 8.PIPELINE PARALLELISM — SPLIT BY LAYER, NOT BY TENSORstage 0stage 1stage 2stage 3Pale cells are bubbles: fill and drain. Bubble fraction = (stages−1)/(micro-batches+stages−1).Communication is one activation per stage boundary, which is tiny. Good across slow links; costs latency.Rule of thumb:TP within a node(NVLink, 900 GB/s).PP across nodes(Ethernet, 25 GB/s).
Tensor parallelism buys lower latency and costs bandwidth. Pipeline parallelism buys capacity and costs latency. The interconnect decides which you can afford.

The communication cost, precisely

A ring all-reduce over N GPUs moves 2(N−1)/N times the payload per GPU. The payload here is the activation: batch × d_model × 2 bytes. There are two all-reduces per layer.

For a 70B model with 80 layers at batch 16 and d=8192, that is about 70 MB per step. Over NVLink at 450 GB/s per direction (H100's headline 900 GB/s is the bidirectional total) it costs about 0.16 ms, negligible against several milliseconds of compute. At batch 512 the payload grows 32×, and over PCIe or Ethernet it dominates completely.

Hence the rule everyone converges on: tensor parallelism inside a node, pipeline parallelism across nodes.

Data parallelism is not the same thing

Replicating the whole model across GPUs and load-balancing requests between replicas is data parallelism. It multiplies throughput with zero communication. It does nothing for a model that does not fit, and nothing for single-request latency.

Real deployments combine all three: TP within a node so the model fits and each token is fast, PP across nodes if it still does not fit, and DP replicas above that for throughput, with a router in front that ideally preserves prefix-cache locality.

Explore

Try it

Simulatorparallelism cost model
tensor parallel
pipeline parallel
interconnect
GPUs
4
weights / GPU
32.6 GB
step time
9.80 ms
scaling efficiency
99%
where a decode step goes
compute
  • reading weights (useful work)
  • all-reduce (tensor parallel tax)
  • pipeline bubble (idle stages)

Two all-reduces per layer, 60.0 MB per step in total. Every layer's output must be summed across all TP ranks before the next layer starts. It is a blocking, latency-sensitive collective rather than a background transfer.

scaling curve · step time vs TP degree
TP1 · 38.9ms
TP2 · 19.5ms
TP4 · 9.8ms
TP8 · 4.9ms

Compute halves with every doubling; communication does not. Past the crossover, adding GPUs makes each token slower. Find it by switching the interconnect.

pipeline bubbles
stage0

No pipeline parallelism: one stage, no bubbles. Raise PP to see the fill and drain cost appear.

Select a 70B model with TP8 on NVLink at batch 16: efficiency stays high and the step time drops nearly 8×. Now raise the batch to 512 and switch to PCIe. The all-reduce bar swells until TP1 beats TP8 outright, and the scaling curve panel draws the crossover for you. Raise PP after that: the bubble appears, then shrinks again as you add micro-batches.

Build it

Implementation

Sharding the attention block carries one constraint worth encoding as an assertion.

code/s17_parallelism.py (excerpt)python
class ShardedAttention:
    def __init__(self, w, rank, world, n_heads, n_kv_heads):
        assert n_heads % world == 0, "TP degree must divide query heads"
        assert n_kv_heads % world == 0, (
            f"TP{world} needs n_kv_heads divisible by {world}, got {n_kv_heads}. "
            "Either lower the TP degree or replicate KV heads across ranks."
        )
        self.local_heads = n_heads // world
        self.local_kv_heads = n_kv_heads // world

        h = slice(rank * self.local_heads, (rank + 1) * self.local_heads)
        self.wq = w.wq[:, h]                  # column shard
        self.wk, self.wv = shard_kv(w, rank, world)
        self.wo = w.wo[h, :]                  # row shard -> needs all-reduce

Every rank must sample the same token

In exact arithmetic, every rank holds identical logits after the final all-reduce. Floating-point reduction order can differ by a few ULPs, so ranks that sample independently can diverge on a near-tie. From then on they are generating different sequences against KV caches that disagree, and the output degrades over hundreds of tokens with no error anywhere. Sample on rank 0 and broadcast, or seed every rank identically and accept the tie-breaking risk.
Run it locally
Simulates tensor and pipeline parallelism in-process, asserts the sharded forward pass matches the single-GPU result exactly, and reports a communication-cost model across interconnect speeds and batch sizes.
$ python code/s17_parallelism.py
Expect: Exact numerical agreement between sharded and unsharded paths at TP1–TP11, an assertion failure demonstrating the GQA divisibility constraint, and a table where TP8 wins on NVLink but loses to TP1 on Ethernet at batch 512.

Only NumPy is required — setup instructions.

Production notes

In production

  • vLLM — tensor_parallel_size and pipeline_parallel_size; a MultiProcExecutor spawns one worker process per GPU coordinated over shared-memory message queues.
  • NCCL — the collective library underneath everything, with topology-aware ring and tree algorithms.
  • Sequence / ring attention — a third axis: split one sequence across GPUs and pass the FlashAttention running state around a ring. Needed only for extreme context lengths.
  • Expert parallelism — for MoE (S16), shard by expert rather than by tensor. The all-reduce becomes an all-to-all.

Exercises

  1. 1
    Compute the all-reduce volume for your model and find the interconnect bandwidth at which TP8 stops beating TP4.
  2. 2
    Implement the sample-on-rank-0-and-broadcast fix, then deliberately break it and measure how many tokens it takes before the ranks visibly diverge.
  3. 3
    Model a hybrid deployment: TP8 within nodes, PP2 across two nodes, DP4 above that. Compute total GPUs, per-GPU memory, and expected throughput, and compare with pure TP16.

Continue

What's next

You can now split a model across GPUs. The last structural idea splits the workload instead: S18 runs prefill and decode on entirely separate machines, because the two phases want opposite hardware.

Check yourself

Questions

Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.

Check yourselfQuestion 1 of 6

In the Megatron pattern, why is the first FFN matmul column-sharded and the second row-sharded?

Score 0/6