Volume II · Splitting the Work
Splitting the Sequence, Splitting the Experts
Lectures 3 and 4 gave a model three ways to be too big for one device — too many parameters per layer, too many layers, too much optimizer state — and a parallelism strategy for each. This lecture closes Volume II with two more ways a model can outgrow a single GPU that have nothing to do with parameter count at all: a sequence too long to fit, and a model that is sparse rather than merely large.
The mental model
There are four axes along which a model's work can be split — data, tensor, pipeline, and now sequence or expert — and every large training run in production is some composition of them, chosen to match where that particular model actually runs out of room.
A fourth axis: the sequence itself
Every splitting strategy so far assumed the sequence length was small enough that a single device could hold the activations and KV state for one layer, one micro-batch, at a time. Push context length far enough — hundreds of thousands or millions of tokens — and that assumption breaks even when the model's own parameters would fit comfortably. The attention computation over a sequence of length L costs memory and compute that scale with L², and at extreme lengths that quadratic term dominates everything else. Tensor parallelism splits inside a layer; pipeline parallelism splits across layers; neither one splits a single attention computation across the sequence dimension itself. That is the gap Ring Attention fills.
Ring Attention (Liu, Zaharia, and Abbeel — commonly cited as arXiv:2310.01889, "Ring Attention with Blockwise Transformers for Near-Infinite Context"; treat the exact arXiv identifier as worth re-verifying directly against the source rather than taking on faith) shards one long sequence into blocks, one block per device, and arranges the devices in a ring. Each device keeps its own query (Q) block stationary. The key/value (K/V) blocks, by contrast, travel: at each step, every device passes its current K/V block to its neighbor around the ring and receives a new one from the other side, computing a partial blockwise attention update against whatever K/V block currently sits in front of it. After enough steps to have passed every K/V block past every Q block once, each device has accumulated the full attention output for its own query block — without any single device ever having held the whole sequence's K/V cache at once.
The reason this scales the way it does is that the communication step — pass a K/V block to the next device, receive one from the previous device — can be overlapped with the compute step: while device A computes attention against the K/V block it currently holds, it is simultaneously sending its own block onward and receiving the next one. Overlap hides the communication cost behind useful compute, so throughput scales roughly linearly as more devices join the ring, rather than degrading the way a naive all-to-all exchange of every block with every other block would. Ring Attention is not the only way to parallelize across the sequence dimension — a survey of sequence-parallelism schemes (arXiv:2405.07719, "USP: A Unified Sequence Parallelism Approach") places it alongside alternatives such as DeepSpeed-Ulysses, noting that some of those alternatives hit tighter scalability limits than a ring-based approach as device count grows.
A fifth axis: which expert handles which token
Sequence parallelism answers "the input doesn't fit." Expert parallelism answers a different question: what if the model is enormous, but any given token only needs a small fraction of it? A mixture-of-experts (MoE) layer replaces one large feed-forward block with many smaller "expert" feed-forward blocks plus a router that, for each token, picks a small number of experts to send that token through. Total parameter count can be huge — most of the experts, most of the time, are not touched — while the compute per token stays close to what a much smaller dense model would cost.
GShard (Lepikhin et al., arXiv:2006.16668) demonstrated this at real scale: a sparsely-gated MoE model scaled to 600 billion parameters, trained across 2048 TPU v3 cores, using top-2 routing — each token sent to the two experts its router scores highest. The parallelism this requires is genuinely a new axis, not a variant of tensor or pipeline parallelism: each device hosts a distinct subset of the experts, and because a token's assigned expert may live on a different device than the one currently holding that token's activations, the routing step is an explicit all-to-all communication — every device sends out the tokens it needs routed elsewhere, and receives the tokens other devices have routed to the experts it holds. Switch Transformer (Fedus, Zoph, and Shazeer, arXiv:2101.03961) simplified the routing rule itself, from top-2 down to top-1 — a single expert per token — trading a little bit of the ensembling effect of consulting two experts for a simpler, cheaper routing step and less communication overhead.
Sparse, not just large
The GShard and Switch Transformer numbers above describe parameter count, not compute per token — that distinction is the entire point of routing. A 600-billion- parameter MoE model with top-2 routing does roughly the compute of a model with only two experts' worth of feed-forward parameters per token, not the compute of a dense 600-billion-parameter model. Expert parallelism exists to place that huge, mostly-idle parameter set across many devices without paying to move all of it for every token.
Four axes, composed
Lay all four splits — data, tensor, pipeline, and now sequence or expert — next to each other and a pattern emerges: none of them is "the" way to parallelize a model. Each answers a different bottleneck. Data parallelism answers "I have more data than time" by replicating the whole model. Tensor parallelism answers "one layer doesn't fit" by cutting inside it. Pipeline parallelism answers "the whole stack doesn't fit" by cutting across layers, at the cost of bubble time. Sequence parallelism answers "the input itself doesn't fit." Expert parallelism answers "the model is sparse — most of it, most of the time, is idle" by placing different experts on different devices and routing tokens to them.
| Axis | What doesn't fit | How it splits |
|---|---|---|
| Data | Not enough time to see all the data | Replicate the whole model, split the batch |
| Tensor | One layer's weights don't fit one device | Cut inside a layer's matrices |
| Pipeline | All the layers together don't fit | Cut across layers, pay in bubble time |
| Sequence (Ring Attention) | One sequence's activations/KV don't fit | Shard the sequence across a device ring |
| Expert (GShard, Switch) | Total parameters are huge but sparsely used | Place experts on devices, route tokens to them |
Real training runs at the frontier do not pick one axis and stop — they compose several at once. The Megatron-LM follow-up work already covered in Lectures 3 and 4 documents this composition as "3D parallelism" — data, tensor, and pipeline parallelism applied together in one run, each cutting along a different dimension of the same problem. Add expert parallelism to that mix for a sparse model and some writeups call the result "4D." Neither name marks a single paper's invention; both name an engineering composition of independently developed techniques, chosen and combined to match wherever a particular model and cluster actually run out of room.
Read the primary source
- Ring Attention with Blockwise Transformers for Near-Infinite Context, Liu, Zaharia, Abbeel.
- USP: A Unified Sequence Parallelism Approach for Long Context Generative AI.
- GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, Lepikhin et al.
- Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, Fedus, Zoph, Shazeer.
- Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM.