Expert Parallel Math Deep

Routing as an expectation over a batch

Drop from the per-token picture to the per-batch one, the level at which a device feels load. For token x the router produces logits ℓ = x · W_r (W_r: [h, E]) and gate probabilities p_i(x) = softmax(ℓ)_i; top-k dispatch sends x to the k highest scorers. What a device physically receives, though, is a count. Over a batch of T tokens, per expert:

P_i = (1/T) Σ_x p_i(x)     # importance: mean gate mass (differentiable)
f_i = n_i / (T·k)           # load: dispatch fraction (a hard count)

Here n_i is the token count for expert i, and since each token makes k dispatches, Σ_i n_i = T·k, so Σ_i f_i = 1; likewise Σ_i P_i = 1. The ideal is uniform, P_i = f_i = 1/E. The gap between P_i (what the router wants, differentiably) and f_i (what the hardware gets, non-differentiably) is the seam the balancing machinery is built around.

Why imbalance is a parallelism problem, not a quality one

On one device a lopsided f is a modeling nuisance; across an expert-parallel group it is a wall-clock disaster, and the reason is synchronization. The dispatch all-to-all, the expert matmul, and the combine all-to-all form a barrier-bounded pipeline: no device leaves the layer until the last finishes. If the busiest expert holds n_max tokens while the mean is μ = T·k/E, the layer runs a factor n_max / μ slower than a balanced one — every idle device burning money.

Left alone the router makes this worse: a slightly preferred expert gets more tokens, sharpens on more gradient signal, and is preferred more strongly next step — a rich-get-richer collapse onto a few live experts. Nothing in the task loss opposes it, and you cannot backpropagate a fix directly, because f_i comes from an argmax/top-k selection whose gradient is zero almost everywhere. The counts that cause the straggler are exactly the quantity you cannot differentiate — the impasse the auxiliary loss is engineered around.

Deriving the load-balancing auxiliary loss

Start from what we want to minimize: unevenness in the hard load f. A natural, dimensionless measure is Σ_i f_i^2. With Σ f_i = 1 fixed, Jensen puts its minimum at the uniform f_i = 1/E (value 1/E); a single hot expert drives it to 1. The right objective — and useless alone, because f_i has no gradient.

The trick is to replace one of the two factors with its differentiable twin. The importance P_i tracks the same preference as f_i (both are large exactly when expert i is favored) but flows gradients. That yields the Switch-style / GShard auxiliary loss:

L_aux = α · E · Σ_i f_i · P_i

  f_i  treated as a constant weight (no grad)   ->   ∂L_aux/∂θ = α·E·Σ_i f_i · ∂P_i/∂θ

The gradient pushes probability mass down on experts that are already overloaded (large f_i), redistributing future tokens. The leading E normalizes the range to [1, E]: the uniform optimum gives E · Σ(1/E)(1/E) = 1, total collapse gives E, independent of expert count. The coefficient α (~0.01) keeps it a gentle nudge that never overwhelms the task loss.

Capacity and the exact drop-rate formula

The auxiliary loss reduces skew but cannot guarantee it, and an all-to-all needs fixed-size buffers anyway. Both yield to a hard per-expert capacity:

C  = CF · (T·k / E) = CF · μ      (tokens per expert; CF = capacity factor, ~1.0-2.0)
dropped_i = max(0, n_i - C)          drop_rate = ( Σ_i max(0, n_i - C) ) / (T·k)

Tokens arriving at a full expert are dropped — they skip the layer through the residual connection, contributing nothing but also costing nothing. Under-filled experts are the opposite problem: their buffers are padded up to C so every all-to-all message is the same shape. The group therefore always moves E · C = CF · T · k slots no matter how many carry real tokens, making CF literally the overhead factor. So capacity does three jobs with one dial: it bounds the straggler at n_max ≤ C = CF·μ, fixes every buffer size statically, and sets how many real tokens you sacrifice. The best CF is the smallest one whose drop rate the loss can absorb.

Worked example: aux loss, drops, and utilization

Take E = 8, top-k = 2, T = 65536, so total dispatches T·k = 131072 and mean load μ = 16384. A skewed step produces per-expert loads n = [30000, 24000, 20000, 16000, 14000, 12000, 8000, 7072] (sum = 131072). With CF = 1.25, C = 20480.

drops   = (30000-20480) + (24000-20480) = 9520 + 3520 = 13040
drop_rate = 13040 / 131072 ≈ 9.95%

slots moved   = E·C        = 163840          # padded, fixed
real carried  = 131072 - 13040 = 118032
utilization   = 118032 / 163840 ≈ 72%          # 28% of the wire moves padding

aux (if P_i mirrors f_i, an illustration only): E·Σ f_i^2 = 8 · 0.1508 ≈ 1.21

Two caveats. The aux figure assumes P_i ≈ f_i purely to get a number — the loss exists because they differ, so read 1.21 as “~21% above the balanced floor of 1, shrinking under the gradient,” not an equality. And the two costs are distinct: ~10% of tokens went uncomputed (a quality tax), while 28% of bandwidth shipped empty padding (a throughput tax).

The throughput model: all-to-all meets expert compute

Now put communication and compute on the same clock. Because buffers are padded to C, both terms are set by capacity, not by the real token count. An FFN of width 4h costs about 16h^2 forward FLOPs per token, and the busy device runs n_max of them; the two all-to-alls each move one h-vector per token both ways at per-device bandwidth W (bytes b each):

T_compute = n_max · 16h^2 / P                    # P = realized FLOP/s; n_max ≤ C
T_comm    = 2 · ((D-1)/D) · C · h · b / W        # dispatch + combine, padded to C

The all-to-all volume itself is the systems companion’s subject; what matters here is that T_comm is activation-sized and on the critical path (compute waits on dispatch), so without micro-batch overlap the layer pays them additively: T_layer ≈ T_comm + T_compute. The decisive knob is W: on intra-node NVLink the collectives nearly vanish; across nodes at ~10 GB/s they can match or exceed the expert math they bracket — which is why the group is placed on the fastest links available.