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.

Layer efficiency in one equation

Fold it into one number: how much of a perfectly balanced, zero-communication machine do you get? Compare the ideal per-device compute time T_ideal = μ · 16h^2 / P against the real layer time. The efficiency factors into a straggler term and a communication term:

η = T_ideal / T_layer = ( μ / n_max ) · ( T_compute / (T_comm + T_compute) )
        \_____ straggler _____/   \________ communication ________/

Plug the running example: h = 4096, P = 200 TFLOP/s, D = 8, b = 2, inter-node W = 10 GB/s. Per token 16h^2 ≈ 2.68e8 FLOPs. The busy expert is capped at n_max = C = 20480, so T_compute ≈ 27.5 ms; the padded all-to-alls cost T_comm ≈ 29 ms, giving T_layer ≈ 56.5 ms. Then η = (16384/20480) × (27.5/56.5) = 0.80 × 0.49 ≈ 39%. Even with a bounded straggler, the layer delivers barely over a third of its peak — the capacity padding stretched compute by 25% and cross-node communication ate half the clock. This is the number expert parallelism actually optimizes, and the factored form is the to-do list: shrink the straggler by flattening load (stronger aux loss) or lowering CF; shrink communication with faster links, a larger h (T_compute scales as h^2 but T_comm only as h), or by overlapping the collectives with compute. Growing the group D trades the two — easier memory, bigger all-to-all.

Router stability and the z-loss

One numerical footnote earns its place because it interacts with balancing. As the router sharpens, the raw logits ℓ can grow without bound — the softmax only cares about differences — and large logits make the gate brittle in low precision, amplifying the very collapse the aux loss fights. The router z-loss penalizes the log-partition directly:

L_z = (1/T) Σ_x ( log Σ_j exp(ℓ_j(x)) )^2

It pulls the logsumexp toward zero, keeping logits small and the gate well-conditioned without dictating which expert wins; a tiny coefficient (~1e-3) suffices. The two losses are complementary: the aux loss shapes where mass goes so no device straggles, and the z-loss keeps the magnitudes sane so bf16 routing stays stable over a long run.

Pitfalls and CPU-SLM notes

A few traps recur. Token dropping is silent: a too-low CF starves the model and shows up only as a stubborn loss gap, so monitor drop rate directly, not the aux-loss value alone. The aux loss is a nudge, not a guarantee: α too small lets experts collapse, too large distorts routing — both cost quality. Efficiency is a product, so one weak factor caps the whole thing: perfect balance cannot rescue an all-to-all on a slow link, nor infinite bandwidth a 5:1 straggler. On CPU or modest hardware the communication term dominates — bandwidth is scarce and small hidden sizes make the 1/h amortization work against you — so a sparse MoE that looks cheap on paper (few active FLOPs) can run slower than the dense model it replaced once communication and the straggler are paid. The active-parameter FLOP count is not the runtime; the efficiency equation is.