Back to Learning
AI Next Labs / Learning
5D Parallelism: how LLM training is split across GPUs
Why a large model needs many GPUs, and the six ways to split one training step across them: data, tensor, sequence, context, pipeline and expert parallelism.
Published
Written by Jaydeep Raijada.
A large language model doesn't fit on one GPU, and even if it did, training it on one GPU would take years. This series shows the five ways to split the work across many GPUs, plus expert parallelism for mixture-of-experts models: why each exists, what it costs in communication, and how they fit together. Every lesson has step-through diagrams and worked numbers.
The series
- Part 0: Foundations. Transformer shapes, the 16 bytes per parameter, activation memory, FLOPs, gradient accumulation and recomputation, with step-through diagrams.
- Part 1: Data parallelism. Copy the model, split the batch: ring all-reduce, overlapping communication, the batch-size recipe and ZeRO stages 1 to 3.
- Part 2: Tensor parallelism. Split every weight matrix across GPUs: column and row splits with real numbers, a TP transformer block, and what it costs in communication.
- Part 3: Sequence parallelism. Remove the activations tensor parallelism leaves replicated, using reduce-scatter and all-gather at no extra communication cost.
- Part 4: Context parallelism. Split long sequences across GPUs in every operation: ring attention, online softmax with real numbers, and zig-zag load balancing.
- Part 5: Pipeline parallelism. Split the layers into stages and keep every GPU busy: GPipe, 1F1B, interleaving, Zero Bubble and DualPipe, with a schedule simulator.
- Part 6: Expert parallelism. Put different experts on different GPUs: the router, top-k routing with real numbers, all-to-all dispatch and combine, and the communication bottleneck.
- Part 7: Combining the dimensions. Map each kind of parallelism to the right network link, walk through a 24-GPU layout, and plan a cluster.
The interactive overview below has a memory calculator, the six cuts drawn side by side, and a cheat sheet.
Long-form guides
- The Ultra-Scale Playbook: Training LLMs on GPU Clusters (Guide). Hugging Face (nanotron). Follows nearly the same path as these notes, with interactive memory estimates.
- How to Scale Your Model (Guide). Google DeepMind's scaling book. Rooflines, sharding and collective costs from first principles.
- How to Train Really Large Models on Many GPUs? (Blog). Lilian Weng, 2021. A concise survey of parallelism and memory-saving techniques.
- Machine Learning Engineering Open Book (Book). Stas Bekman. Practical notes on training at scale, networking and debugging.