Expert Parallelism and All-to-All

Dense data parallelism has one collective per step and it is an all-reduce. A sparse mixture-of-experts layer has a different one, twice per layer, and it is an all-to-all. That single change is why MoE training is a systems problem rather than a modelling problem, and why the routing function ends up constrained by network topology.

The arrangement is simple to state. Experts are partitioned across devices; every device holds the full attention stack and a subset of experts. When a token reaches an MoE layer, its router picks \(k\) experts that are mostly on other devices, so the activations must be shipped there, processed, and shipped back.

The two exchanges

Dispatch. Each device sorts its local tokens by destination expert and sends each group to the device owning that expert. Every device sends to every other device, which is an all-to-all.

Combine. After the expert FFNs run, the outputs travel back along the reverse permutation, weighted by the router's gate values. A second all-to-all.

Both are on the critical path of every MoE layer, and unlike all-reduce, an all-to-all does not compress: it moves \(O(k \cdot T \cdot d)\) bytes for \(T\) tokens of hidden size \(d\), and the volume grows with the number of experts each token selects. GShard established the pattern while scaling a translation model past 600B parameters on 2,048 TPU v3 chips in four days (Lepikhin et al., 2020, arXiv:2006.16668), and every large MoE since has been an exercise in making these two collectives cheaper.

Three ways the system fights back

Overlap the communication with computation. DeepSeek-V3 splits each pipeline chunk into attention, all-to-all dispatch, MLP, and all-to-all combine, and schedules them so that communication for one micro-batch hides behind computation for another, reporting near-zero all-to-all overhead at 64-way expert parallelism spanning 8 nodes (DeepSeek-AI, 2024, arXiv:2412.19437).

Constrain routing to the topology. Inside a node, NVLink gives roughly 160 GB/s; between nodes, InfiniBand gives about 50 GB/s, a ratio of about 3.2 to 1 on the H800 systems DeepSeek used. Their response was node-limited routing: a token may be dispatched to at most 4 nodes, chosen before the per-expert selection. The router is no longer free; it is a function whose codomain is shaped by the interconnect.

Remove the token-dropping compromise. Early MoE implementations fixed a per-expert capacity so the all-to-all had a static shape, and dropped any token that overflowed. MegaBlocks reformulated the expert computation as block-sparse matrix multiplication so no token needs dropping, reporting up to 40% speedup over Tutel and 2.4x over Megatron-LM (Gale et al., 2022, arXiv:2211.15841). Tutel took the complementary route of adapting parallelism and pipelining at runtime, reporting 4.96x and 5.75x speedup for a single MoE layer on 16 and 2,048 A100 GPUs respectively (Hwang et al., 2022, arXiv:2206.03382).

When it breaks

Load imbalance is a latency problem, not just a quality problem. The all-to-all is a barrier: the layer finishes when the slowest device finishes. One expert receiving three times its share of tokens makes every other device wait, so the auxiliary load-balancing loss is protecting throughput as much as it is protecting expert specialisation. This is why the imbalance shows up as a wall-clock regression before it shows up in the loss curve.

Expert parallelism composes awkwardly with everything else. It must be combined with data, tensor, pipeline, and sequence parallelism into a single device mesh, and the collectives interleave. An all-to-all inside a pipeline stage interacts with the pipeline schedule; expert parallelism across a tensor-parallel group changes which ranks participate in which collective. Getting the mesh wrong produces a correct model that trains at a fraction of the achievable rate.

Small batches make it worse. The all-to-all cost per token is roughly fixed, while compute per device falls with batch size. At small batch, the exchange dominates, which is why MoE inference at low concurrency has such poor economics compared to MoE training. DeepSpeed-MoE addressed the inference side specifically, reporting 7.3x better latency and cost than prior MoE inference approaches and 4.5x faster, 9x cheaper inference than quality-equivalent dense models (Rajbhandari et al., ICML 2022, arXiv:2201.05596).

All experts must be resident. Expert parallelism distributes parameters but does not reduce their total. The aggregate memory across the mesh is the full parameter count, and a device failure takes specific experts offline rather than a replaceable replica.