EP Math — Expert Parallelism (MoE)

What expert parallelism actually shards

Start with the object being split. An MoE feed-forward layer contains E experts, each a full FFN with its own weights. A top-k router picks k experts per token (commonly k = 1 or k = 2). Expert parallelism lays those E experts across D devices, so each device owns E / D experts. With E = 64 experts and D = 8 devices, every device stores 8 experts and nothing else of the expert pool.

This is the opposite of tensor parallelism, which splits a single weight matrix across devices so all of them cooperate on one computation. EP keeps each expert whole and instead partitions which experts live where. The consequence: a token routed to expert 37 must be moved to whichever device holds it before processing, then moved back. The router’s per-token decision sets up an irregular many-to-many pattern that a single collective — the all-to-all — is built to serve.

Notation: global tokens N and per-device tokens T

Two token counts run through every formula, and blurring them is the easiest way to get the math wrong, so pin them down once. Let N be the tokens in the batch handled by the whole EP group and T the tokens on one device, so N = D · T; let d be the hidden dimension — the length of each token’s activation vector. Before the MoE layer each device holds its T tokens as a [T, d] tensor. The router assigns each token to k experts, producing T · k token-to-expert assignments per device (N · k across the group). Those assignments are what the all-to-all delivers: each is a d-vector that must travel from the token’s device to the target expert’s device. Volume is counted as assignments × d × bytes — use T for one device, N for the group.

The dispatch all-to-all: bytes on the wire

The first collective is dispatch: send each token’s activation to the device hosting its chosen expert. A device begins with T · k assignments, each a d-vector. In bfloat16 (2 bytes/element) the data a single device sends is:

dispatch_send(per device)
  = T · k · d · 2 bytes    (produced)
  × (D - 1) / D             (fraction sent)

The (D - 1) / D factor appears because, with balanced routing, a device’s assignments spread uniformly over all D devices; the 1 / D share bound for its own experts never touches the network, and ingress mirrors egress on average. The key scaling fact: dispatch volume grows with k and d but is independent of E — doubling the experts does not, by itself, move more bytes.

The combine all-to-all: the return trip

After each expert applies its FFN, every processed token must go home. That is a second all-to-all, combine, the exact reverse of dispatch: the expert outputs (still d-vectors) travel back to the token’s original device, where its k expert results are weighted by the router’s gate values and summed.

Because the FFN preserves the hidden dimension, combine moves the same volume as dispatch, so one MoE layer costs two all-to-alls of equal size:

layer_comm(per device)
  ≈ 2 × T · k · d · 2 bytes × (D - 1) / D

This is the number that dominates MoE scaling. The expert matmuls are the same FLOPs a dense FFN would do for the routed tokens, but the two all-to-alls are pure overhead with no compute analogue in a dense model. An all-to-all also stresses bisection bandwidth harder than a tensor-parallel all-reduce, and only helps if it overlaps with expert compute — which is why EP is kept modest and confined to a fast interconnect domain.

Capacity factor: fixing buffer sizes

Routing is data-dependent, so the number of tokens landing on any one expert varies from batch to batch. But collectives need fixed, pre-agreed tensor shapes: you cannot launch an all-to-all whose message sizes are only known at runtime without expensive synchronization. The fix is to give every expert a fixed capacity — a buffer sized for a set number of tokens:

expected load per expert  μ = N · k / E
capacity per expert       C = f · N · k / E   (round up)

Here f is the capacity factor, a small multiplier above 1 (typically 1.0 to 1.5). Every expert’s buffer holds exactly C tokens; the all-to-all then moves a dense tensor of shape roughly [E, C, d] across the group. This trades a little wasted bandwidth — padding under-full experts up to C — for predictable, statically-shaped collectives, almost always the right trade.

Token drop: what overflow costs

Fixed capacity has a sharp edge. If more than C tokens route to one expert in a batch, the surplus does not fit and is dropped — that expert does not process them. A dropped token is not an error; its activation passes through the layer via the residual connection unchanged, as if the expert output were zero. But it missed the computation the router wanted, which costs a little quality.

The capacity factor is the dial that trades bandwidth and memory against drops. With perfectly uniform routing, f = 1.0 would waste nothing and drop nothing. Real routing is lumpy, so at f = 1.0 the popular experts overflow while the quiet ones sit padded, and drop rates of several percent are common. Nudging f to 1.25 gives 25% headroom and usually pushes drops near zero — at the price of 25% more all-to-all volume and buffer memory. Training tolerates modest drops; inference runs at higher f because a dropped token degrades a user-visible answer.

Load balance: why experts drift uneven

Everything above assumes routing is roughly balanced. Left alone it is not: a freely-trained router tends to collapse onto a few favorite experts. That is doubly bad under EP — the favored experts’ devices become compute hotspots while others idle, and they overflow capacity and drop tokens. The all-to-all runs at the speed of its busiest link, so one hot device stalls the whole group.

The standard remedy is an auxiliary load-balancing loss added during training, commonly α · E · Σ_i f_i · P_i, where f_i is the fraction of tokens dispatched to expert i and P_i the mean gate mass it received. The product is minimized when both are uniform at 1 / E, so the term pressures the router toward even utilization — filling the fixed capacity buffers evenly and keeping both drops and device idle time low.

Placing experts across devices

Expert placement is mostly an arithmetic and memory problem. With E experts on D devices you want E divisible by D so each holds an equal E / D experts; an uneven split leaves some devices with more expert weight and compute, a hotspot in hardware rather than in the router. Per-device expert memory is just (E / D) × params_per_expert — and this is the whole reason EP exists: it is how a model with hundreds of experts fits at all.

A worked memory example: 64 experts, each an FFN with hidden width 4d over d = 4096, is about 2 × 4 × d^2 ≈ 134M parameters per expert, so 8.6B in experts alone. Spread over D = 8 devices that is ~1.1B per device, versus the full 8.6B on each if replicated. EP converts a memory wall into a communication cost — the trade the all-to-all math is really pricing.

A worked communication example

Put numbers on one layer. Take d = 4096, top-k = 2, T = 2048 tokens per device, D = 8 devices, bfloat16. The data one device produces for dispatch is:

T · k · d · 2 = 2048 · 2 · 4096 · 2 = 33.6 MB
× (D-1)/D = 7/8  → ~29 MB sent

Combine sends the same, so this device moves about 58 MB per MoE layer. A model with 16 MoE layers shuffles on the order of ~0.9 GB per device per forward pass; against a ~600 GB/s intra-node link that is roughly 1.5 ms of pure communication, small only if it overlaps with compute. Note again what is absent: E. Growing from 8 to 64 experts changes memory and load balance, not the bytes on the wire, as long as k, d, and the token count hold.

How EP composes with DP, TP, and PP

EP never runs alone; it is one axis of a multi-dimensional mesh. Typically EP applies only to the MoE FFN sublayer, while attention and dense layers use tensor parallelism (TP), the model is replicated with data parallelism (DP), and depth is split with pipeline parallelism (PP). The device count multiplies out as total = DP × PP × TP × EP, and the EP group is the set of devices that jointly hold one full copy of the expert pool and run the all-to-all among themselves.

The practical rule is to place the EP all-to-all on the fastest links: EP usually stays within a node (or NVLink domain) while DP and PP, whose collectives are rarer or smaller, span the slower cross-node fabric. Get this wrong — an EP group straddling a slow link — and the all-to-all math above, benign on paper, becomes the dominant cost of the entire training step.

CPU and small-model implications

On a CPU or a single small device the all-to-all disappears — there is nothing to shuffle between, so ‘EP’ degenerates into looping over experts locally. That flips the bottleneck: the cost is no longer bandwidth but the scattered, small, irregularly-sized matmuls that top-k routing produces, which map badly onto wide CPU vector units and thrash the cache. The math’s lesson holds off the cluster too — MoE buys parameters cheaply in FLOPs but pays in data movement and irregularity, whether that is an inter-device all-to-all or a cache-unfriendly gather on one chip. If your deployment cannot absorb that movement, a dense model of equal active size is usually the more honest choice.