Home / Fundamentals / Distributed Training

Tensor, pipeline, sequence and expert parallelism

Megatron-style tensor parallelism, pipeline schedules and the bubble, sequence and context parallelism for long inputs, expert parallelism for MoE, and how to map them onto the cluster topology.

15 min read

Data parallelism and FSDP shard state, but each GPU still computes whole layers on whole sequences. When a layer's activations or a model's per-step communication become too large, split the model itself.

Tensor parallelism (TP)

Split each weight matrix across t GPUs. In the MLP (Megatron-style):

  • The first matmul is column-parallel: each GPU holds h × (4h/t) and computes its slice of the hidden activations; the nonlinearity applies independently.
  • The second matmul is row-parallel: each GPU holds (4h/t) × h, producing a partial sum; one all-reduce combines them.

Attention splits naturally by heads. Result: two all-reduces per layer in forward and two in backward, each of b·s·h activations. That's frequent and latency-sensitive, so TP is kept within a node on NVLink (t ≤ 8).

Sequence parallelism splits the norm and dropout regions along the sequence dimension, replacing each all-reduce with a reduce-scatter plus all-gather (same volume) and dividing those activations by t.

Pipeline parallelism (PP)

Split layers into p stages on different GPUs; send activations forward and gradients backward point-to-point. With one batch, only one stage works at a time, so the batch is split into m micro-batches.

bubble fraction = (p − 1) / (m + p − 1)

Scroll sideways to see the whole diagram.

Pipeline bubble vs micro-batchesChart: p = 4 stages; p = 8 stages; p = 16 stages. 1248163264128020406080100micro-batches per step (m)idle fraction (%)p = 4 stagesp = 8 stagesp = 16 stagesp=4, m=16: 16%
Pipeline bubble vs micro-batches. Fraction of time each GPU idles in a GPipe or 1F1B schedule, (p − 1) / (m + p − 1), for p pipeline stages and m micro-batches per step. Deeper pipelines need many more micro-batches to keep the bubble small; interleaving divides it by roughly the number of chunks per GPU.
  • GPipe: all forwards, then all backwards; must store activations for all m micro-batches.
  • 1F1B: after warmup, alternate one forward and one backward; at most p micro-batches in flight, same bubble.
  • Interleaved 1F1B: each GPU holds v non-contiguous chunks of layers, cutting the bubble by about v at the cost of more communication.
  • Zero-bubble schedules split the backward into input-gradient and weight-gradient parts to fill idle time.

PP's communication is small (activations at stage boundaries), so it works across nodes, but it needs many micro-batches, which constrains the global batch size.

Context (sequence) parallelism for long inputs

For 128k-token sequences, activations of even one sequence exceed a GPU's memory. Context parallelism splits the sequence across GPUs; attention needs keys and values from all chunks, exchanged in a ring (ring attention) or via all-to-all over heads (Ulysses-style). Causal masking makes chunks unequal in work, so chunks are assigned in a load-balanced order.

Expert parallelism (EP)

In mixture-of-experts layers, experts are placed on different GPUs; each token is routed with an all-to-all to the GPUs holding its chosen experts and back. Load imbalance and all-to-all bandwidth dominate, so expert placement and capacity factors matter.

Putting it together: topology-aware layouts

Scroll sideways to see the whole diagram.

3D parallel layout: TP 8 × PP 4 × DP 64Components: Stage 1 · node (TP 8 · emb, 1–19); Stage 2 · node (TP 8 · layers 20–40); Stage 3 · node (TP 8 · layers 41–61); Stage 4 · node (TP 8 · 62–80 + LM head); Stage 1 × 63 (same weights shard); Stage 2 × 63; Stage 3 × 63; Stage 4 × 63; ZeRO-1 shards (fp32 + Adam ÷ 64). Flow: Stage 1 · node to Stage 2 · node (acts); Stage 2 · node to Stage 3 · node (acts); Stage 3 · node to Stage 4 · node (acts); Stage 1 · node to Stage 1 × 63 (DP reduce-scatter); Stage 2 · node to Stage 2 × 63; Stage 3 · node to Stage 3 × 63; Stage 4 · node to Stage 4 × 63; ZeRO-1 shards to Stage 4 · node (all stages).REPLICA 1 OF 64 · 4 NODES · 32 GPUSREPLICAS 2 … 64 (SAME LAYOUT, DIFFERENT DATA)Stage 1 · nodeTP 8 · emb, 1–19Stage 2 · nodeTP 8 · layers 20–40Stage 3 · nodeTP 8 · layers 41–61Stage 4 · nodeTP 8 · 62–80 + LM headStage 1 × 63same weights shardStage 2 × 63Stage 3 × 63Stage 4 × 63ZeRO-1 shardsfp32 + Adam ÷ 64actsactsactsDP reduce-scatterall stages
3D parallel layout: TP 8 × PP 4 × DP 64. Each node's 8 GPUs form one tensor-parallel group over NVLink and hold one pipeline stage of about 20 layers. Four nodes make one model replica; activations flow stage to stage point-to-point. 64 replicas process different data, and each stage reduce-scatters its gradients with the same stage in every other replica over InfiniBand, with optimizer state sharded by ZeRO-1.

A common recipe for a large dense model:

  1. TP = 8 within each node (NVLink).
  2. PP across a few nodes, just enough that each stage's weights, optimizer shards and activations fit.
  3. DP (with ZeRO-1, or FSDP/HSDP) across the rest, on the inter-node network.
  4. Add CP for long-context phases and EP for MoE.

total GPUs = TP × PP × DP (× CP). Choose the smallest model-parallel degree that fits, because each kind of model parallelism costs efficiency.

What to remember

  • TP: within a node, 2 all-reduces per layer per direction.
  • PP: bubble (p−1)/(m+p−1); needs many micro-batches.
  • CP for long sequences, EP for MoE, DP for everything else.