5D parallelism: data, tensor, context, pipeline and expert

When data parallelism and ZeRO run out, there are four more ways to cut a training job across GPUs. Each one splits a different axis of the work and pays with a different kind of network traffic. This guide names the axis, the traffic and the original paper for each, and ends with how the pieces are stacked on a real cluster.

Video by Jia-Bin Huang · Watch on YouTube

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. Watch

Chapter guide

Timestamps come from the video's description. The right-hand column is this course's reading guide, written from the chapter titles; it doesn't describe what the video shows frame by frame.

2. The ideas, explained

The work cube

One training step processes a batch of sequences, each made of tokens, through a stack of layers. Picture that as a cube with three axes: sequences × tokens × layers. Each parallelism slices a different part of it:

data parallelism     splits sequences        traffic: gradients, once per step
context parallelism  splits tokens           traffic: keys/values, every attention layer
pipeline parallelism splits layers           traffic: activations, at stage boundaries
tensor parallelism   splits inside a layer   traffic: activations, twice per layer
expert parallelism   splits the MoE experts  traffic: tokens, all-to-all per MoE layer

total GPUs = DP × TP × CP × PP   (EP usually reuses GPUs from the DP/CP groups)

Recap: data parallelism and its limits

Data parallelism and ZeRO are explained in Video 01 and Module 08 · Distributed training. The limits that push you to the other four methods: ZeRO shards model states but not activations, which grow with sequence length; the global batch is at least DP × (one sequence), and very large batches stop helping optimisation; and gradient synchronisation over thousands of GPUs eventually costs more than it saves.

Tensor parallelism (Megatron-LM)

Shoeybi et al. (2019) split each transformer block's matrices across GPUs. For the MLP Z = GeLU(X·A)·B:

column-linear:  A = [A1 | A2]       GPU i computes Y_i = GeLU(X·A_i)
                (GeLU is elementwise, so no communication is needed here)
row-linear:     B = [B1 ; B2]       GPU i computes Z_i = Y_i·B_i
                Z = Z1 + Z2         ← one all-reduce

attention: split the heads (Q, K, V projections column-wise),
           output projection row-wise  ← one all-reduce
per layer: 2 all-reduces in the forward pass, 2 in the backward pass

Every layer waits on these all-reduces, so tensor parallelism is kept inside one server, where GPUs share fast NVLink-class links. The communication volume scales with batch × sequence × hidden size, so it grows with the activations, not the weights.

Sequence parallelism

Korthikanti et al. (2022) noticed that tensor parallelism leaves LayerNorm and dropout replicated on every GPU, with full-size activations. They split those regions along the sequence dimension instead, and replace each all-reduce with a reduce-scatter (into the sequence-split region) and an all-gather (back into the tensor-split region). That's the same total communication as the all-reduce, and much less activation memory. The same paper introduced selective activation recomputation.

Context parallelism and ring attention

Context parallelism splits every layer's tokens across GPUs, so one long sequence spreads over several servers. Everything except attention is per-token and needs no communication. Attention needs each query to see every earlier key and value. Ring attention (Liu, Zaharia and Abbeel, 2023) arranges the GPUs in a ring: each holds a block of queries, and key/value blocks are passed around the ring while each GPU computes attention against the block it currently holds, overlapping the transfer with compute.

That only works because softmax can be computed in pieces. Online softmax (Milakov and Gimelshein, 2018) keeps a running maximum and a running sum, and rescales when a new block raises the maximum:

for each incoming K,V block:
    s     = q·Kᵀ / √d
    m_new = max(m, max(s))
    l     = l·exp(m − m_new) + Σ exp(s − m_new)
    o     = o·exp(m − m_new) + exp(s − m_new)·V
    m     = m_new
output = o / l          (exactly equal to softmax over all blocks)

With a causal mask, a GPU holding early tokens has little work and a GPU holding late tokens has a lot. The zig-zag fix cuts the sequence into 2N chunks and gives GPU i chunks i and 2N−1−i, so every GPU gets one early and one late chunk. Ring attention is covered in more depth in Module 08 · Beyond transformers.

Pipeline parallelism: GPipe, 1F1B, interleaving

Pipeline parallelism gives consecutive layers to different GPUs (stages) and sends activations between neighbours. A naive schedule leaves all but one stage idle. GPipe (Huang et al., 2018) splits the batch into m micro-batches so stages overlap; with p stages the idle fraction is

bubble ≈ (p − 1) / (m + p − 1)       p = 8, m = 32  →  7/39 ≈ 18%
interleaved (v chunks per GPU):  bubble time shrinks by about v

GPipe runs all forwards, then all backwards, so it stores activations for all m micro-batches. The 1F1B schedule, from PipeDream (Harlap, Narayanan et al., 2018), alternates one forward and one backward once the pipeline is full, so at most about p micro-batches are in flight. Narayanan et al. (2021) added the synchronous 1F1B schedule with interleaved stages to Megatron-LM: each GPU holds several non-adjacent chunks of layers, which shrinks the bubble at the cost of more point-to-point messages. The GPipe schedule and its bubble are animated in 30u30 Day 24 · GPipe.

Zero-bubble pipelines and DualPipe

Qi et al. (2023, "Zero Bubble Pipeline Parallelism") split the backward pass into two parts: B, the gradient with respect to the layer's input (which the previous stage is waiting for), and W, the gradient with respect to the weights (which nobody is waiting for). Doing B first and slotting W into the gaps removes the bubble. The paper reports the first zero-bubble schedules under synchronous training and up to 23% higher throughput than 1F1B at a similar memory limit. DeepSeek-V3 (2024) describes DualPipe, a bidirectional schedule that feeds micro-batches from both ends of the pipeline and overlaps forward and backward computation with the expert-parallel communication of another micro-batch.

Expert parallelism

In a mixture-of-experts layer, a router sends each token to its top-k experts. Expert parallelism places different experts on different GPUs (GShard, Lepikhin et al., 2020; Switch Transformer, Fedus et al., 2021). Each MoE layer then needs two all-to-alls: dispatch tokens to their experts, then combine the outputs back. Experts' weights aren't duplicated, so it scales parameters cheaply, but the all-to-all is sensitive to load imbalance and network topology. Routing, load balancing and capacity are covered in Module 08 · Beyond transformers (with a load-balancing demo) and Module 05 · Trillion-scale architecture.

Pipeline vs ZeRO-3

Both spread the layers' weights so no GPU holds the whole model. ZeRO-3 shards every layer across all GPUs and moves weights: each layer is all-gathered before use. Pipeline parallelism gives whole layers to one stage and moves activations, point-to-point, at stage boundaries. ZeRO-3 has no bubble and no scheduling to tune, but its weight traffic only hides behind compute if each GPU has enough tokens per step. Pipeline traffic is small, but you must feed enough micro-batches to hide the bubble, and the stages must be balanced.

Network hierarchy and a real layout

A GPU server typically has 8 GPUs joined by very fast links; servers are joined by a slower network. So the chattiest traffic goes innermost:

innermost (inside a server):  TP      all-reduce every layer
                              CP, EP  per attention / MoE layer
                              PP      point-to-point at stage edges
outermost (across the cluster): DP    one gradient sync per step

The Llama 3 paper (Grattafiori, Dubey et al., 2024) trained the 405B model with this four-dimensional layout, ordered TP, CP, PP, DP from innermost to outermost. Its scaling table lists, for 16,384 GPUs at 8K context, TP 8 × PP 16 × DP 128 = 16,384; for the 128K-context stage it adds CP 16 and reduces DP to 8. Read the paper's parallelism section for the exact batch sizes and utilisation figures.

Rules of thumb

The video ends with its own three-step cheatsheet; watch it for the author's version. The common rules in the literature (for instance Hugging Face's Ultra-Scale Playbook) are: first make the model fit (TP inside a server, then PP or ZeRO-3 across servers, CP if long sequences blow up activations, EP for MoE); then use DP for the remaining GPUs to reach the target global batch; then tune micro-batch size, recomputation and overlap for throughput.

3. Where the course goes deeper

4. Why it matters for robotics

Self-check

References

The video's description lists no references. Primary papers cited in this guide:

• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019)
• Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022)
• Ring Attention with Blockwise Transformers for Near-Infinite Context (Liu, Zaharia, Abbeel, 2023)
• Online normalizer calculation for softmax (Milakov, Gimelshein, 2018)
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (Huang et al., 2018)
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training (Harlap, Narayanan et al., 2018)
• Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al., 2021)
• Zero Bubble Pipeline Parallelism (Qi et al., 2023)
• DeepSeek-V3 Technical Report (DeepSeek-AI, 2024)
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (Lepikhin et al., 2020)
• Switch Transformers: Scaling to Trillion Parameter Models (Fedus, Zoph, Shazeer, 2021)
• The Llama 3 Herd of Models (Grattafiori, Dubey et al., 2024)
• The Ultra-Scale Playbook: Training LLMs on GPU Clusters (Hugging Face)