← Inference From First Principles

Expert Routing and Load Balancing

What a trained router actually does to a serving system: how three balancing methods differ, what a capacity limit costs in quality on learned routing, how placement from measured load cuts the straggler, and what communication would add.

Three small mixture-of-experts language models were trained identically except for how they were encouraged to balance their experts. All three reached the same quality on held-out text, to within the difference between seeds. Then each was given the same fixed capacity per expert at inference, as a serving system with static shapes would, for batches of 64 tokens. The model with no balancing term lost 0.28 bits per character at a capacity factor of 2.0; the two balanced models lost 0.05. At a factor of 1.25 the loss was 0.68 bits against 0.33. The weights were not different in kind, nor was the quality without a limit. What differed was how evenly the router spread its tokens, and that is invisible until an execution engine with a finite budget meets it.

Chapter 19 had to work with an untrained router and synthetic skew, because no trained mixture-of-experts model was available. This chapter removes that limit in the only way open to me without downloading a large model: it trains a small one on the book’s own text, three ways, three seeds each, and measures balance, the cost of capacity, the value of expert placement, and what communication would add. The model is a laboratory device. It has 5.4 million parameters, it was trained for about 94 seconds, and it overfits; nothing below says how a large model routes. What it provides is learned routers on which the mechanisms can be measured and the Chapter 19 questions answered at the scale where they can be.

Where this sits

Chapter 19 gave the layer, four ways to run it, and a capacity limit that drops overflowing assignments; its open questions were what dropping does to quality and what a real router’s balance looks like. Chapter 11’s three ceilings and Chapter 15’s batching make this chapter’s measurements matter: a batch is the unit that meets a capacity, and a step waits for its slowest device. Chapter 18’s fixed-shape discipline is the reason capacity exists in the first place.

The Inference Lab piece added here is the routing laboratory: a small trainable mixture model with three balancing variants, evaluation under capacity limits, trace collection, placement algorithms that decide on one split of data and are scored on another, and the communication arithmetic. Chapter 21 turns from the feed-forward layers to attention.

For the Pi series, the idea to keep is that balance is a property a model is given by its training and a serving system inherits. The runtime cannot change which experts the router prefers; it can choose where experts live, how much slack to allow, and how to treat the overflow. All three are decisions about a distribution that has to be measured, on traffic that was not used to choose them.

Three layers of the problem

The starter material rightly separated three things that are easily confused, and the measurements below follow its separation.

  1. The router produces scores and top-\(k\) choices. How even those choices are is learned.
  2. A capacity policy decides what happens when an expert receives more assignments than a batch can hold (Chapter 19: drop them, pass the token on through the residual connection).
  3. Placement maps experts to devices; a distributed layer waits for the busiest device, and tokens move between devices to reach their experts.

A balanced-looking histogram of the first says little about the other two: balance on average over a corpus is not balance within a batch, and balance across experts is not balance across devices. The experiments are arranged to show each gap.

The laboratory

The model is a character-level transformer: 4 layers, \(d = 192\), 4 heads, context 128, every feed-forward block replaced by 16 experts of width 192 with top-2 routing and a Mixtral-style softmax over the selected logits (5,428,224 parameters). It trains for 2,000 steps of 32 windows with AdamW, seeds 0, 1 and 2. Three variants differ in a single place:

  • Plain: top-2 routing, no balancing term.
  • Auxiliary loss: Switch Transformer’s loss, \(\alpha N \sum_i f_i P_i\) with \(\alpha = 0.01\), where \(f_i\) is the share of assignments to expert \(i\) and \(P_i\) its mean router probability. It is minimised by uniform routing; \(f\) is not differentiable but \(P\) is, so the gradient pushes probability away from busy experts. The loss is added to training.
  • Bias (auxiliary-loss-free balancing): each expert has a bias added to its score for the top-\(k\) selection only; the gate that weights the expert’s output is computed from the unbiased logits. After every batch the bias moves by \(u \cdot \text{sign}(\bar c - c_i)\), where \(c_i\) is expert \(i\)’s count and \(\bar c\) the mean, with \(u = 0.005\) here (the paper’s setup used 0.001 over far more steps; I fixed 0.005 in advance and did not tune it). No gradient is involved.

I checked all three descriptions against the papers: the Switch loss and its \(\alpha\), the bias-for-selection-only rule and sign update of Wang et al. (a preprint at the time of reading), and DeepSeekMoE’s expert-level loss of the same form plus a device-level variant I did not implement.

Data discipline. The text is this book. Chapters 1 to 14, with the example code and tests that go with them, train the models (739,076 characters, so prose, equations, tables and Python). Chapters 15 and 16 are a calibration split, used for placement and nothing else. Chapters 17 to 19 are the test split (227,957 characters), used only for the numbers reported. The test text contains no character absent from the training vocabulary of 135.

The tests tie this model back to Chapter 19: with no capacity limit its mixture layer equals Chapter 19’s reference implementation, and with a limit it equals Chapter 19’s capacity dispatch, including the count of dropped assignments.

What the three variants do

Quality first, because a balancing method that costs quality would settle the question at once. Bits per character on the test split, mean and range over three seeds:

Variant Test bits per character Calibration
Plain 2.104 (2.090 to 2.113) 1.917
Auxiliary loss 2.103 (2.081 to 2.129) 1.911
Bias 2.119 (2.103 to 2.133) 1.926

The three are indistinguishable at this scale. (All overfit: the final training cross-entropy was about 0.8 nats, 1.2 bits, against 2.1 on the test text. That limits what the quality numbers can show, though not the routing statistics.)

Balance on the test split, the busiest expert’s share of assignments over the mean share, layer by layer:

Variant Layer 1 Layer 2 Layer 3 Layer 4
Untrained router (reference) 1.4 to 1.6 1.5 to 1.9 1.8 to 2.8 1.6 to 2.2
Plain 3.73 2.49 3.57 6.76
Auxiliary loss 1.34 1.17 1.17 1.19
Bias 1.07 1.09 1.11 1.10

Three observations. First, training without a balancing term made the router less balanced than random initialisation: the plain model’s layer 4 sent 42.3% of all assignments to one expert (uniform is 6.25%) and its routing entropy was 0.63 of the maximum, while the untrained router was within a factor of three of even. Learned routing concentrates. Second, both balancing methods fixed it, the bias method more tightly (within 11% of even in every layer, against 17 to 34% for the auxiliary loss) and with no term in the loss. Third, no expert died in any model: all 16 experts in all layers received tokens on the test split, even in the plain model, so “the busiest expert is overloaded” and “experts are unused” are different failures, and only the first appeared here.

The routers did learn something

Balance alone would be satisfied by a random router, so I also measured whether the routing means anything. Taking the first-choice expert and the class of the input character (letter, digit, space, newline, other), the normalised mutual information between them is a measure of how much the expert choice tells you about what kind of character is being processed:

Layer 1 Layer 2 Layer 3 Layer 4
Untrained 0.18 to 0.22 0.16 to 0.23 0.19 to 0.21 0.18 to 0.24
Plain 0.68 0.66 0.57 0.15
Auxiliary loss 0.65 0.62 0.63 0.50
Bias 0.59 0.62 0.55 0.46

Training roughly tripled the information in the first three layers. The interesting row is the plain model’s fourth layer: its collapse onto one dominant expert took its specialisation with it (0.15, no more than a random router), while the balanced models kept 0.46 to 0.50. In this toy, balance and specialisation were not in tension; the balancing terms preserved the specialisation that collapse destroyed. I would not extend that to large models, where the literature reports tension between the auxiliary loss and quality (the aux-loss-free paper’s motivation); here the measurement is of one small model on one small corpus.

What a capacity limit costs

This is the Chapter 19 question. For each model I ran the test text through the same network with every expert limited to \(\lceil c \cdot Gk/E \rceil\) assignments per group of \(G\) tokens, \(c\) being the capacity factor. Assignments beyond the limit are dropped and the token passes through that expert’s layer unchanged. Test bits per character, with the dropped share of assignments in brackets, for groups of 512 tokens (4 windows) and, in the second table, 64:

Groups of 512 Plain Auxiliary loss Bias
Factor 2.0 2.337 (20.1%) 2.114 (0.7%) 2.129 (0.8%)
Factor 1.5 2.529 (28.0%) 2.159 (3.1%) 2.164 (3.2%)
Factor 1.25 2.697 (33.5%) 2.259 (6.7%) 2.246 (6.7%)
Factor 1.0 2.935 (40.5%) 2.538 (14.7%) 2.537 (14.9%)
Groups of 64 Plain Auxiliary loss Bias
Factor 2.0 2.388 (21.8%) 2.151 (2.6%) 2.162 (3.0%)
Factor 1.5 2.602 (30.2%) 2.268 (7.3%) 2.277 (7.8%)
Factor 1.25 2.781 (35.9%) 2.440 (12.7%) 2.451 (13.2%)
Factor 1.0 3.049 (43.1%) 2.770 (21.7%) 2.787 (22.3%)

With no limit the three models scored 2.10, 2.10 and 2.12. So the quality cost of dropping, which Chapter 19 listed as unmeasured, is large and strongly dependent on the router: for the plain model a factor of 2.0 (twice the average load) drops a fifth of all assignments and loses about 0.2 to 0.3 bits per character; for the balanced models the same factor drops under 1% at groups of 512 and loses about 0.01 bits. A capacity factor is therefore not a property of the layer; it is a property of the layer and the router it serves, and a value chosen on one model can be disastrous on another.

Two further points. Average balance is not batch balance. If the balanced models’ routing were uniform within every group, the exact binomial expression of Chapter 19 says a factor of 1.25 would drop 0.1% of assignments at groups of 512 and 4.5% at 64. They dropped 6.7% and 12.7%. Corpus-level balance coexists with bursts: stretches of code route differently from stretches of prose, and consecutive characters are not independent draws. And the cost does not buy back with larger groups alone: moving from 64 to 512 tokens helped the balanced models substantially (12.7% to 6.7% dropped at a factor of 1.25) and the plain one hardly at all (35.9% to 33.5%), because the plain model’s problem is a persistent hot expert and not the sampling noise a bigger group averages away.

The remedies all cost something. Raise the factor and the fixed-shape dispatch pays more arithmetic for empty slots (with \(C = T\), nothing drops at \(E\) times the arithmetic, as Chapter 19 measured). Process overflow later instead of dropping it and the request waits. Balance the router at training time and the cost moves out of the server, where it is a model property and not an engineering choice.

Placement

If experts live on several devices, a layer finishes when its busiest device does. I placed the 16 experts on 4 devices four ways and scored each by the mean over layers and windows of busiest-device load over mean device load, in consecutive windows of 64 and 512 test tokens. Placement “by load” used the calibration split to give each device four experts by the greedy rule of taking experts from the busiest down and assigning each to the least loaded device with room; the scores are on the test split, which was not used to choose anything. For reference, perfectly uniform routing, which cannot be reached because a window is finite, would give 1.18 at 64 tokens and 1.06 at 512.

Placement (windows of 64 / 512) Plain Auxiliary loss Bias
Contiguous blocks 1.75 / 1.71 1.34 / 1.25 1.31 / 1.20
Round-robin 1.90 / 1.87 1.32 / 1.22 1.33 / 1.22
By calibration load 1.39 / 1.32 1.29 / 1.19 1.28 / 1.17
By load, plus 2 replicas 1.41 / 1.37 1.30 / 1.21 1.30 / 1.21

The pattern: for the unbalanced router, placement decided from measured load cut the straggler by nearly a quarter (1.71 to 1.32 at 512 tokens), much more than either fixed rule. Round-robin, which looked good on Chapter 19’s synthetic skew, was the worst choice for this router (it was marginally better than blocks for the balanced ones), because there is no reason for popular experts to be neighbours in index order and every reason for them to be arbitrary. The decision transferred: scoring the same placement on the calibration split itself gave 1.33, 1.19 and 1.18, almost exactly the test scores, so on this corpus the load pattern was stable across chapters.

For the balanced routers placement had little to give. By-load placement reached 1.17 to 1.19 at 512 tokens against a uniform floor of 1.06; blocks were already at 1.20 to 1.25. Balancing the router at training time did most of what runtime placement could do.

Replication did not help. Duplicating the two busiest experts onto the least loaded devices and splitting their load made the straggler slightly worse in all three variants (for example 1.37 against 1.32 for the plain model at 512 tokens). My replication rule is crude (a fixed two experts, even split, chosen once from calibration), and the result is a statement about that rule on this model. It does not contradict the DeepInfer note’s report of redundant experts: that note, a first-party description, says each GPU hosts redundant experts for hot ones, that load is gathered over a window and experts are reassigned periodically (every 1,000 steps in prefill, 10,000 in decode), and gives slowest-rank-over-mean ratios of 1.045 to 1.090 on its authors’ systems. Those figures are not comparable with mine, which describe a 16-expert toy, and I have not reproduced them. The note, as read, does not say whether weight movement is asynchronous or describe the all-to-all pattern; the starter’s statement that it did was wrong and is withdrawn.

Communication

When experts are spread over devices, tokens have to travel to them. If tokens start spread evenly over \(D\) devices and each assignment goes to a uniformly chosen device, a fraction \((D-1)/D\) of assignments leaves the device, and a device sends, per mixture layer (dispatch plus combine),

\[ 2 \cdot t \cdot k \cdot d \cdot b \cdot \frac{D-1}{D} \]
bytes for \(t\) tokens per device, \(k\) experts per token, width \(d\) and \(b\) bytes per value. For Mixtral's \(d = 4096\) and \(k = 2\) in fp16:
Tokens per device 2 devices 4 devices 8 devices
64 1.0 MB 1.6 MB 1.8 MB
512 8.4 MB 12.6 MB 14.7 MB
4,096 67 MB 101 MB 117 MB

The only link on this machine I can measure is the host-to-device copy, 12.4 GB/s pinned (12.7 the other way). Standing that in for an inter-device link, which it is not, these volumes take 0.08 to 0.15 ms, 0.7 to 1.2 ms and 5.4 to 9.5 ms. Chapter 19 measured 7.3 ms for a Mixtral-shaped layer at 64 tokens on one GPU, so a link of that speed would add 1 to 2 per cent at 64 tokens per device; at 512 tokens (0.7 to 1.2 ms) I have no matching compute measurement (the nearest, at 256 tokens, was 9.85 ms), and at 4,096 tokens per device the transfer alone is several milliseconds. A faster interconnect moves all of this. This is arithmetic with a measured constant, not a measurement of an all-to-all, and it ignores latency, overlap with compute, and the fact that the traffic is only as even as the routing: the same traffic skew that makes a straggler concentrates the bytes on one link.

Expert choice

A different way to get balance by construction is to invert the choice: each expert picks its highest-scoring tokens (Zhou et al., 2022), so every expert’s load is exactly its capacity and a token may get none or several experts. On 512 test tokens from the plain model’s layer-2 scores, with a capacity of 64 per expert, every expert received exactly 64; 2.1% to 3.7% of tokens received no expert at all, one token received 12, and only 54.7% to 57.8% of the token-choice assignments were also chosen by the experts. So the method trades the straggler for a different irregularity. It also needs the scores of every token in the group before it can decide, which in an autoregressive decoder includes tokens not yet generated, so using it for generation needs care that I did not investigate. I applied it to a model trained with token choice; a model trained for expert choice is a different thing, and this is a measurement of the selection rule only.

What the numbers do not show

  • A toy. 5.4 million parameters, 16 experts, characters, 2,000 steps, overfit. Larger models may balance differently without help, may need the auxiliary loss for different reasons, and have more layers in which imbalance can accumulate.
  • One corpus, one split. The book’s text, which mixes prose, equations and code; the calibration-to-test transfer of placement was stable here and need not be elsewhere, particularly under traffic that drifts.
  • Three seeds. The ranges are real (for the plain model’s last layer, 6.3 to 7.6 times the mean), and none of the differences between the two balanced methods is established.
  • Neither balancing method was tuned. \(\alpha\) was Switch’s value, and \(u\) was set in advance; the paper’s ablations suggest the update rate matters.
  • No communication was run. The table is arithmetic; the placement experiment counts assignments, not time.
  • Capacity cost depends on the order tokens are dropped. My rule keeps the earliest tokens in a group; other rules (lowest gate weight first) would change the numbers.
  • No device-level balance loss, no sigmoid gating, no shared experts in the trained model.

Inference Lab increment

The Lab now has:

  • build_model: the trainable mixture language model with the plain, aux and bias routers, capacity-limited evaluation in groups, and per-layer traces.
  • train_model, evaluate, collect_traces, windows: training and measurement with a split discipline.
  • load_balance_stats, nmi, bias_update: balance and specialisation measures and the bias rule.
  • place_block, place_round_robin, place_by_load, replicate_hottest, device_loads_from_counts, straggler: placement.
  • alltoall_bytes_per_device, expert_choice_mask.
  • experiment, reference and link commands.

The rule this adds: balance is learned, capacity is a budget against it, and placement is a decision about measured load. The serving system chooses the second and third; it should choose them from traffic it measured on one set of requests and verify on another, and it should treat the capacity factor as belonging to the model it was tuned on.

What you should now be able to do

  • Describe the Switch auxiliary loss and the auxiliary-loss-free bias method, and say which parts affect selection and which affect the gate.
  • Measure a router’s balance and specialisation from traces, and say why a balanced histogram does not imply balanced batches.
  • Estimate the quality cost of a capacity limit on a given router, and explain why the same factor can be harmless for one model and ruinous for another.
  • Place experts on devices from measured load, validate the placement on traffic that did not choose it, and compare it with fixed rules.
  • Compute the all-to-all volume of a mixture layer and judge whether a link makes it matter.
  • Say what a result from a small trained model licenses and what it does not.

Key terms

  • Auxiliary loss: a training term that penalises uneven routing.
  • Auxiliary-loss-free balancing: a per-expert selection bias updated from observed load, with no gradient.
  • Routing collapse: the router sending most tokens to few experts.
  • Capacity factor: slots per expert as a multiple of the average load.
  • Specialisation: expert choice that depends on the input; measured here by mutual information with character class.
  • Placement: the map from experts to devices.
  • Straggler factor: busiest device’s load over the mean.
  • Redundant expert: a copy of an expert on another device, to share its load.
  • Expert choice: routing in which experts select tokens.

Reader exercises

  1. Tune the bias update rate \(u\) on the calibration split only, then report the test numbers once. How much of the bias method’s advantage depends on \(u\)?
  2. Implement the device-level balance loss of DeepSeekMoE and train with experts grouped into four devices. Does it make placement by load unnecessary?
  3. Change the order in which overflow is dropped to lowest gate weight first. How do the dropped share and the bits per character change at a factor of 1.25?
  4. Make the traffic drift: calibrate on chapters 15 to 16 and test on a stretch that is all code. How does the by-load placement degrade, and how often would you need to re-place?
  5. Improve replication: choose the number of replicas from the calibration load, split by capacity, and place copies to minimise the maximum. Can it beat by-load placement for the plain model?
  6. Train the plain model with 64 experts of width 48 and top-8. Does the last-layer collapse persist?
  7. Measure the autoregressive generation of the model with and without capacity limits at batch sizes 1, 4 and 16: how often is a generated token different?
  8. Estimate, for Mixtral’s shape and the 12.4 GB/s link, the number of tokens per device at which the all-to-all matches the compute time of Chapter 19’s layer C.

References

Evidence in this chapter

Claim Category
The Switch loss, the bias-for-selection-only rule and sign update, device-level balance loss, expert-choice routing Established background (the papers’ text)
The three balancing variants reach the same test quality (2.10, 2.10, 2.12 bits per character) Book result, 3 seeds, a toy model
Without balancing the router became less balanced than at initialisation (6.8x the mean in layer 4); both balancing methods held it within 1.07 to 1.34x Book result, 3 seeds, a toy model
Trained routers carry information about the input (NMI 0.55 to 0.68 in layers 1 to 3 against 0.2 untrained); collapse removed it in layer 4 Book result, a toy model
Capacity limits cost 0.28 bits per character (plain) against 0.05 (balanced) at factor 2.0, groups of 64 Book result, answers Chapter 19’s open question for this model
Corpus-level balance does not give batch-level balance (6.7% dropped against 0.1% for uniform routing, factor 1.25, groups of 512) Book result
Placement by measured load cut the straggler for the plain router (1.71 to 1.32) and transferred from calibration to test; round-robin was worst Book result, a toy model
Replication of two hot experts did not help under the rule used Book result for that rule only
The DeepInfer note’s redundant-expert design and ratios of 1.045 to 1.090 Reported by the authors, not reproduced; asynchronous movement not stated by the note
All-to-all volumes and times at 12.4 GB/s Arithmetic with one measured constant; no all-to-all run
Large trained models, real multi-device runs, drifting traffic Not measured

Implementation status

  • Manuscript: content/books/inference/20-chapter.md
  • Runnable examples: examples/ch20/routing.py (original hand-made-policy simulator, superseded), examples/ch20/routing_lab.py (trainable mixture model, balance variants, capacity, placement and communication arithmetic)
  • Tests: tests/test_ch20_routing.py, tests/test_ch20_routing_lab.py (balance statistics, the bias rule balancing fixed scores, auxiliary-loss value at uniform and collapsed routing, mutual information, placements and replication, expert-choice capacity, equality with Chapter 19’s reference and capacity dispatch, groups independent, bias moving only in training, a small model learning; CPU, no GPU needed)
  • Evidence: research/ch20/sources.md, research/ch20/results.json, research/evidence-register.md
  • Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It makes no claim about large trained models, real multi-device communication or any production scheduler. Chapter 21 moves from the feed-forward layers to attention: multi-head, multi-query and multi-latent variants, and what each does to the KV cache that Chapters 10 to 16 managed.