Capacity Factors, Token Dropping, and Stability

"They gave me a buffer for exactly five hundred tokens. On a good batch I sat half empty; on a bad one I turned good tokens away at the door. Nobody asked me how I felt about it."

An Expert at Capacity

The previous sections built a layer that can route each token to a chosen expert and exchange tokens across machines so that every expert receives its assigned work. That description is complete at the level of mathematics, where an expert simply processes however many tokens happen to be routed to it. It is incomplete at the level of hardware. A GPU kernel that multiplies a batch of tokens by an expert's weight matrix must know the batch dimension before it launches, and the all-to-all collective that ships tokens between nodes must allocate send and receive buffers before it runs. Neither can resize itself in the middle of a step to absorb whatever the router decided this time. The number of tokens per expert, however, is data-dependent and varies batch to batch, which is exactly the quantity these static structures cannot accommodate. Expert capacity is how the system squares that circle.

1. Why Experts Need a Fixed Capacity Beginner

Imagine you are sizing the receive buffer on each device before a batch arrives. You do not yet know how many tokens the router will send to the expert living there, because routing depends on the tokens themselves. You have two honest options. You could allocate a worst-case buffer large enough for the unlikely event that every token in the batch picks this one expert, which wastes almost all of that memory on every normal batch. Or you could pick a fixed, modest buffer size and decide in advance what happens when more tokens arrive than fit. Real systems take the second option, and the fixed size they choose is called the expert capacity.

Capacity is set relative to the average, because the average is the load each expert would carry under perfect balance. With $T$ tokens in a batch (counting each token once per expert it is routed to, so a top-$k$ router contributes $k$ copies) and $E$ experts, the balanced share is $T/E$ tokens per expert. The capacity factor $c$ scales that share into the actual buffer size,

so $c = 1$ provisions exactly the balanced load and nothing more, while $c = 1.5$ leaves each expert room for fifty percent more tokens than it would receive under perfect balance. The ceiling makes the buffer an integer number of token slots. This one number, multiplied across every expert and every layer, fixes the memory footprint and the communication volume of the whole MoE block before a single token is routed, which is precisely what the static all-to-all and the static matmul kernels require.

2. Token Dropping and Padding Waste Intermediate

Because the buffer is fixed, every batch resolves into one of two failure modes per expert, and usually both happen somewhere in the layer at once. When an expert is assigned more tokens than its capacity, the surplus is dropped: those tokens skip the expert entirely and pass through unchanged on the residual connection, as if the MoE layer were the identity for them. The forward computation still produces an output, so training does not crash, but those tokens received no expert transformation at this layer, which is a real loss of model quality. When an expert is assigned fewer tokens than its capacity, the unused slots are padded with zeros so the buffer is full-shaped for the kernel; the expert dutifully multiplies those zero rows by its weights and discards the result, which is compute and communication spent on nothing.

Both quantities have clean definitions. Let $n_e$ be the number of tokens routed to expert $e$ and let $\text{cap} = \lceil c\,T/E \rceil$ be the capacity. The fraction of tokens dropped across the layer is

the total overflow summed over experts and normalized by the batch. The padding waste is the mirror image, the fraction of provisioned slots that go unused,

These two numbers move in opposite directions as you turn the capacity knob, and the figure below shows why. The diagram references the same picture the demo in Section 4 will quantify.

The asymmetry of the two costs matters. A dropped token is a quality cost that the optimizer pays silently: the loss is slightly worse than it could have been, and no error is raised. Padding is a throughput cost that shows up directly in the wall-clock of every step and in the bytes moved by the all-to-all of Section 17.5, because padded slots are real rows that travel across the network. A team tuning an MoE layer is therefore balancing an invisible quality leak against a visible speed tax, which is why the capacity factor is one of the most-tuned hyperparameters in sparse training.

3. The Capacity-Factor Trade-Off, and Why Load Balance Sets the Floor Intermediate

Turning the capacity factor up has a monotone effect on each cost: higher $c$ means a taller buffer, so fewer tokens overflow (drop rate falls) but more slots sit empty (padding waste rises), and the all-to-all moves more bytes per step. Turning it down does the reverse. There is no setting that zeroes both costs unless the load is perfectly balanced, and it never is. The crucial dependency, the one that ties this section back to Section 17.6, is that the load distribution decides how expensive the trade-off is. Under near-balanced routing, a modest $c$ close to $1$ drops almost nothing and wastes little, because every expert sits near the average the capacity was sized for. Under skewed routing, where a few popular experts attract a large share of tokens, even a generous $c$ leaves the hot experts overflowing while the cold ones pad heavily, so you pay on both sides at once.

This is the practical reason the load-balancing loss of Section 17.6 exists. Its job is not balance for its own sake; it is to flatten the load distribution enough that a small capacity factor suffices, which keeps both the drop rate and the padding waste low simultaneously. Capacity and balance are two halves of one mechanism: the balancing loss reshapes the demand, and the capacity factor provisions the supply, and they are tuned together. Typical production values for the capacity factor sit between $1.0$ and $2.0$ during training, with inference often using a different (sometimes larger) value because a served request cannot tolerate a dropped token the way a training batch can, a point developed in Section 17.8.

# DeepSpeed-MoE: capacity and dropping handled inside the layer.
from deepspeed.moe.layer import MoE

moe_layer = MoE(
    hidden_size=4096,
    expert=expert_mlp,            # the per-expert feed-forward module
    num_experts=64,
    k=1,                          # top-1 routing
    capacity_factor=1.25,         # train-time buffer = ceil(1.25 * tokens / experts)
    eval_capacity_factor=2.0,     # larger buffer at inference to avoid drops
    use_residual=False,
)
# The layer dispatches tokens, runs the expert all-to-all, drops overflow onto
# the residual, pads underflow, and exposes the drop rate through its metrics.

4. From Scratch: Drop Rate and Padding Waste Versus Capacity Intermediate

The cleanest way to feel the trade-off is to route a batch under both a balanced and a skewed load and read off the two costs as the capacity factor varies. The code below uses only NumPy. It draws a top-1 routing assignment from a probability vector that is either uniform (balanced) or tilted so half the experts are six times more attractive (skewed), counts the tokens per expert, then applies the drop-rate and pad-waste formulas from Section 2 across a sweep of capacity factors.

import numpy as np

def route_and_measure(num_tokens, num_experts, capacity_factor, skew, rng):
    avg = num_tokens / num_experts
    capacity = int(np.ceil(capacity_factor * avg))         # ceil(c * T / E)
    base = np.ones(num_experts)
    base[: num_experts // 2] *= skew                       # half the experts attract `skew`x more
    probs = base / base.sum()
    assign = rng.choice(num_experts, size=num_tokens, p=probs)   # top-1 routing
    counts = np.bincount(assign, minlength=num_experts)          # tokens per expert
    dropped = np.maximum(counts - capacity, 0).sum()             # overflow above the buffer
    padding = np.maximum(capacity - counts, 0).sum()            # empty slots below the buffer
    total_slots = capacity * num_experts
    return capacity, dropped / num_tokens, padding / total_slots

rng = np.random.default_rng(0)
N, E = 8192, 16
print(f"tokens={N}  experts={E}  (top-1 routing)")
print(f"{'cap_factor':>10} {'load':>9} {'capacity':>9} {'drop_rate':>10} {'pad_waste':>10}")
for skew, label in [(1.0, 'balanced'), (6.0, 'skewed')]:
    for cf in (1.0, 1.25, 1.5, 2.0):
        cap, drop, pad = route_and_measure(N, E, cf, skew, rng)
        print(f"{cf:>10.2f} {label:>9} {cap:>9d} {drop*100:>9.2f}% {pad*100:>9.2f}%")

tokens=8192  experts=16  (top-1 routing)
cap_factor      load  capacity  drop_rate  pad_waste
      1.00  balanced       512      1.55%      1.55%
      1.25  balanced       640      0.00%     20.00%
      1.50  balanced       768      0.00%     33.33%
      2.00  balanced      1024      0.00%     50.00%
      1.00    skewed       512     36.17%     36.17%
      1.25    skewed       640     23.24%     38.59%
      1.50    skewed       768     10.44%     40.29%
      2.00    skewed      1024      0.00%     50.00%

Read the two halves of the table against each other and the lesson of Section 3 is stark. In the balanced regime, drops vanish by $c = 1.25$ and everything above that is wasted slots, so a tight capacity is both cheap and safe. In the skewed regime, the drop rate stays double-digit until $c = 1.5$ and only reaches zero at $c = 2.0$, where padding waste hits fifty percent because the cold experts are half empty by construction. The same capacity factor buys a clean run under balance and a leaky, wasteful one under skew. This is the quantitative form of the claim that the balancing loss of Section 17.6 and the capacity factor must be tuned as a pair.

5. Training Stability: Router Collapse and Loss Spikes Advanced

Capacity governs efficiency and the slow quality leak of dropped tokens. A second, sharper hazard is unique to sparse models: the router can destabilize training outright. The router is a small learned layer that produces a logit per expert and selects the top one, and that discrete top-$k$ selection sits in a feedback loop with the experts. If an expert gets chosen slightly more often early in training, it receives more gradient, improves faster, and so looks even more attractive to the router, which chooses it still more. Left unchecked this positive feedback drives router collapse, where the model funnels nearly all tokens to a tiny clutch of experts and the rest go cold and untrained, defeating the entire point of having many experts. The load-balancing loss of Section 17.6 is the first line of defense against collapse, but it is not the whole story.

The second symptom is loss spikes. Because the router's logits feed a softmax and a hard selection, large logit magnitudes make routing decisions brittle: a small weight update can flip many tokens from one expert to another, causing a sudden jump in the loss that can in the worst case diverge the run. The standard remedy is the router z-loss, an auxiliary penalty on the magnitude of the router logits. With logits $\ell_{i,e}$ for token $i$ and expert $e$ across a batch of $T$ tokens, the z-loss penalizes the log-partition-function so the logits stay small and the softmax stays numerically calm,

It is added to the task loss with a small coefficient and was a key stabilizer in the ST-MoE work that made large sparse models trainable without the loss spikes that plagued earlier attempts. Two further practices round out the toolkit. Careful router initialization keeps the initial logits small and the early assignments near-uniform, so no expert wins the feedback race before the experts have learned anything. And lower-precision-safe routing computes the router's softmax and selection in float32 even when the rest of the model runs in bfloat16 or float16, because the router's discrete decision is exactly the place where a few bits of rounding can flip an assignment and inject noise into the feedback loop. These are cheap insurance: the router is a tiny fraction of the model's compute, so keeping it in higher precision costs almost nothing.

6. How Capacity, Balance, and the All-to-All Interlock Intermediate

It is worth stating plainly how the three mechanisms of this chapter compose, because they are easy to treat as separate knobs when they are really one coupled system. The all-to-all of Section 17.5 moves a fixed-shape tensor whose size is set by the capacity, so the capacity factor directly determines the communication volume per step: doubling $c$ doubles the bytes the all-to-all ships, padding included. The load balance of Section 17.6 determines how much of that fixed-shape traffic carries real tokens versus padding, and how many real tokens get dropped before they ever enter the collective. And the stability machinery of Section 5 keeps the router producing the kind of spread-out assignments that let a small capacity factor work, by preventing the collapse that would make any fixed buffer overflow on the hot experts and starve on the cold ones.

So the design loop is circular in a productive way. A well-balanced, stable router lets you run a small capacity factor, which shrinks the all-to-all and speeds every step; a small capacity factor in turn pressures the router, via dropped-token quality loss, to stay balanced. When any one of the three fails, the symptom appears in the others: an unstable router shows up as a spiking drop rate, a weak balancing loss shows up as padding waste you cannot tune away, and an over-large capacity factor shows up as an all-to-all that dominates the step time. Reading the three together, rather than tuning each in isolation, is what separates an MoE layer that delivers its promised throughput from one that merely runs. The serving picture in Section 17.8 revisits these same trade-offs under the harder constraint that an inference request cannot quietly drop its tokens, and the broader comparison against dense distributed models in Chapter 16 puts the whole sparse approach in context.