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.