Tensor and Pipeline Parallelism
How frontier labs split a model across thousands of GPUs by sharding within layers (tensor parallel) and across layers (pipeline parallel), and how to pick the split.
Data parallelism replicates the whole model on every GPU. Once the model is too large for that to be viable, you have to actually split the model itself. Tensor parallelism (TP) shards within a layer - matrix multiplies are partitioned across GPUs. Pipeline parallelism (PP) shards across layers - different GPUs hold different parts of the depth. Both introduce their own pathologies, and the largest training runs in the world stack DP, TP and PP into "3D parallelism" to dodge each one's failure mode.
Tensor parallelism (Megatron style)
A transformer's dominant cost is two matmuls per block: the attention projection and the MLP. Megatron-LM shards them along complementary axes so the synchronisation can be deferred until the end of the block.
For the MLP Y = GeLU(X A) B:
- Shard
Acolumn-wise acrosstGPUs. Each GPU computesGeLU(X A_i)on its slice. No comms needed because the activation is element-wise. - Shard
Brow-wise across the sametGPUs. Each GPU computesGeLU(X A_i) B_i, which is its partial sum of the output. - AllReduce the partial sums at the end of the block to get
Y.
The pattern for attention is analogous (shard heads across GPUs, AllReduce after the output projection). Each transformer block costs one AllReduce in forward and one in backward, all of them inside the block.
Why TP is bandwidth-hungry
That AllReduce happens on the activations, not the gradients - so it scales with batch and sequence length, not parameter count. For long contexts and large batches, TP comms can exceed compute time unless the interconnect is NVLink-class. The unwritten rule across frontier labs: TP degree of at most 8, contained within a single node, so the AllReduces stay on NVLink / NVSwitch.
Pipeline parallelism
PP splits the model's layers across GPUs. GPU 0 holds layers 0-7, GPU 1 holds 8-15, and so on. A forward pass flows activations through the pipe; the backward pass flows gradients back.
The naive schedule (GPipe-style "fill and drain") leaves most GPUs idle most of the time. With p pipeline stages and one batch, the bubble (idle fraction) is roughly (p-1)/p. Micro-batching reclaims most of that.
Stages: 0 1 2 3
+---+---+---+---+
batch 0 | F | F | F | F | forward fills the pipe
batch 0 | | | | B | backward drains it
batch 0 | | | B | |
batch 0 | | B | | |
batch 0 | B | | | |
With m micro-batches per global batch and p stages, the bubble shrinks to (p-1)/(m + p - 1). The trick is making m large without blowing up activation memory (each in-flight micro-batch holds activations).
1F1B vs interleaved 1F1B
1F1B (one forward, one backward) scheduling, introduced in PipeDream, alternates forward and backward micro-batches once the pipe is full. Each stage holds at most p in-flight micro-batches' activations, instead of m in the GPipe schedule. The bubble fraction stays the same as GPipe but peak activation memory drops sharply.
Interleaved 1F1B (Megatron-LM's "virtual pipeline") splits each stage into v interleaved chunks of layers. Stage 0 owns layers 0, p, 2p, ...; stage 1 owns 1, p+1, 2p+1, ...; etc. The bubble shrinks by another factor of v at the cost of v times more pipeline-boundary communication. For frontier-scale training (p = 16, v = 4) the bubble drops from roughly 7% to under 2%.
3D parallelism: DP times TP times PP
When P parameters are too large for any one strategy alone, you compose:
- TP shrinks the per-GPU activation and weight footprint within a layer. Capped at the NVLink island size (typically 8).
- PP shrinks the per-GPU weight footprint across depth. Capped by the bubble: more stages, more bubble.
- DP soaks up the remaining GPUs by replicating the TP-PP unit on more batches. Capped by available global batch size (you cannot scale DP past what your effective batch tolerates without divergence).
So the world-size factorisation is:
world_size = DP x TP x PP
For a 540B-scale model on a 4096-GPU cluster, a typical split would be TP=8, PP=16, DP=32. TP runs within an 8-GPU node. PP runs across 16 nodes connected by InfiniBand. DP replicates that 128-GPU TP-PP unit 32 times.
How to choose your parallelism split
- Pick TP from interconnect, not from preference. TP degree = number of GPUs sharing the fastest interconnect (NVLink island). Almost always 8 today, sometimes 4.
- Pick PP from the largest model your TP island can hold. If your TP=8 island fits 24 layers of the model, then
PP = total_layers / 24. Round up; the bubble cost is acceptable. - Pick micro-batch count
mso the bubble is under 5%. Solve(p-1)/(m+p-1) < 0.05form. Typicallym = 4pto8p. - DP soaks up everything left. Once
TP x PPis fixed,DP = world_size / (TP x PP). - Check the global batch.
global_batch = per_micro_batch * m * DP. If that exceeds what your learning-rate schedule tolerates, you have over-parallelised; drop DP and add more PP or more accumulation.
When each stops scaling
- TP stops at the NVLink island boundary. Past that, the per-block AllReduce on activations crushes throughput.
- PP stops when the bubble dominates. Even interleaved 1F1B has limits; pushing past ~32 stages rarely pays.
- DP stops when the effective batch is too large for stable optimisation. The "critical batch size" for LLMs grows with model size but is not infinite.
This is why FSDP (sharded data parallel) keeps eating ground: it adds a memory-saving knob to the DP axis without requiring you to architect TP and PP. For many shapes, FSDP-only on a fat cluster matches 3D parallelism with a fraction of the engineering cost.
Further reading
- Megatron-LM: Training Multi-Billion Parameter Language Models - Shoeybi et al., the tensor-parallel sharding patterns described above.
- Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM - Narayanan et al., introduces interleaved 1F1B and the 3D-parallelism analysis.
- GPipe: Efficient Training of Giant Neural Networks - Huang et al., the original micro-batched pipeline formulation.
- NVIDIA/Megatron-LM - production reference implementation of 3D parallelism.
- Stas Bekman, Model Parallelism guide - field-engineering notes on choosing the split.
5 flashcards for this concept
Click a card to reveal the answer.