Context & Expert Parallelism

We've split three of the four things you can split: the batch (data parallelism), the tensors (tensor parallelism), and the layers (pipeline parallelism). This chapter covers the fourth — the sequence — and then a fifth axis that isn't a way of splitting an existing computation at all, but of changing the computation so there's less of it per token: mixture of experts.

They're paired here because both are about the parts of the transformer the first three strategies handle least well: very long sequences, and very large feed-forward blocks.

Context parallelism: splitting the sequence#

From chapter 2, activation memory has a term that grows with sequence length, and at long context it dominates everything. Tensor parallelism splits the hidden dimension but not the sequence; at 128k or 1M tokens, even with everything else sharded, one device still holds a full-length sequence's worth of activations and runs out of memory.

Context parallelism splits along the sequence: each of $p$ devices owns a contiguous chunk of tokens, holding only that chunk's activations and KV pairs. The MLP and norms are per-token, so they shard trivially. Attention is the hard part, because token $i$ must attend to every token $j \le i$ — including tokens living on other devices.

Ring attention solves this without ever materializing the full sequence on one device.Liu et al., Ring Attention with Blockwise Transformers for Near-Infinite Context, 2023. It composes directly with flash attention — each device runs a flash-attention kernel over the KV blocks as they arrive. Arrange the devices in a ring. Each holds its query chunk fixed and passes its KV chunk around the ring, one hop per step. At each step a device computes attention between its local queries and whichever KV chunk it currently holds, accumulating results using exactly the online-softmax running-max-and-sum machinery from chapter 4 — because that's precisely the algorithm for combining attention over blocks you see one at a time. After $p$ steps every query has attended to every key, and the KV communication overlaps with the attention compute, so on fast links it nearly disappears.

The wrinkle is load imbalance from causal masking. With a naive contiguous split, the device holding the last chunk attends to the whole sequence while the device holding the first chunk attends to almost nothing — the causal mask makes later tokens far more expensive. Left alone, the busiest device sets the pace and the rest idle. The fix is a zigzag assignment: give each device both an early chunk and a late chunk, so the light and heavy halves balance out and every device does roughly equal work. It's the same "flatten the per-stage curve" lesson as Llama 3's pipeline rebalancing, in a different costume.

Mixture of experts: conditional computation#

Everything so far keeps the computation fixed and splits it. MoE changes the computation. The observation: a dense model runs every parameter on every token, which is why chapter 1's decode has to read all 16 GB of weights per token. But most tokens plausibly don't need most of the network. What if each token only used a fraction of the parameters?

Replace the feed-forward block — which is two-thirds of a transformer's parameters — with $E$ parallel FFN "experts" and a small router that sends each token to only $k$ of them.Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, 2017, is the modern origin; Switch Transformer, GLaM, Mixtral, and DeepSeek-V3 are the line of descent. The router is a learned linear layer producing a softmax over experts; you keep the top $k$ and combine their outputs weighted by the gate. Two common settings:

• Top-1 (Switch Transformer): each token to a single expert. Maximum sparsity, simplest routing.
• Top-2 (GShard, Mixtral, GLaM): each token to two experts, their outputs blended. A little more compute, meaningfully better quality.

The payoff is the decoupling of total parameters from active parameters. Mixtral 8×7B has ~47B total parameters but activates only ~13B per token; it costs like a 13B model to run and knows like something much larger. That ratio is the entire pitch: more capacity at fixed FLOPs per token.

Expert parallelism, and where it hurts#

Experts are large and there are many, so you shard them: expert parallelism puts different experts on different devices. Each device hosts a subset of the $E$ experts, and here's the cost — after the router decides where each token goes, the tokens must physically travel to the device holding their assigned expert, and the results must travel back. That's an all-to-all communication, twice per MoE layer (dispatch and combine).

All-to-all is a demanding collective — every device potentially sends data to every other — and MoE's traffic is data-dependent, which is what makes it hard. Which brings up the defining problem of MoE systems:

Load balancing. Nothing forces the router to spread tokens evenly. Left to itself it collapses — a few popular experts receive most tokens (and their devices become the bottleneck) while others sit nearly idle, wasting the very capacity you added them for. Worse, it's self-reinforcing: an expert that gets more tokens trains faster, gets better, and attracts still more. Every MoE system needs an explicit counterweight:

• Auxiliary load-balancing loss. Add a term to the training objective that penalizes uneven token distribution, nudging the router toward using all experts. Standard since Switch Transformer. The downside is that it's a training objective fighting the modeling objective — you're deliberately degrading routing quality to buy balance.
• Capacity factor and token dropping. Cap each expert at a fixed number of tokens per batch (capacity = factor × average). Tokens over the cap for their chosen expert are dropped — passed through via the residual with no expert applied. Bounds the memory and compute per device at the cost of occasionally not processing a token's FFN at all.
• Auxiliary-loss-free balancing (DeepSeek-V3). Instead of a loss term, add a learned per-expert bias to the routing scores and adjust it dynamically to equalize load — balancing the routing without polluting the gradient.DeepSeek-AI, DeepSeek-V3, 2024. Their ablations show it beats the auxiliary-loss approach on both balance and quality, because it stops the balance objective from corrupting the model objective. This is the current best answer, and the reason is exactly that it decouples the two objectives the aux-loss approach forces into conflict.

Expert parallelism vs tensor parallelism#

Both shard the FFN, so when do you use which? Tensor parallelism splits every expert's matrices across devices with heavy all-reduce traffic; expert parallelism keeps each expert whole and routes tokens to it with all-to-all traffic. In practice large MoE models use both, plus data and pipeline parallelism — expert parallelism across nodes (all-to-all tolerates it better than TP's all-reduce), tensor parallelism within a node, and the router's balance tuned so the all-to-all doesn't hotspot. Which is the natural cue for the next chapter: how all five of these compose at once.