L9.2 · LESSON 2 OF 6

Routing and load balance

Explain expert routing, long-context costs, state-space models and multimodal token streams under explicit assumptions.

THE IDEA

When experts are not used evenly

Nothing forces tokens to spread evenly over the experts. If many tokens choose the same one, that expert (and the device holding it) becomes the bottleneck while others wait. Systems cap how many tokens an expert accepts per batch (its capacity) and either drop or reroute the overflow. Training often adds an auxiliary balance loss that rises when routing is skewed.

capacity factor1

Eight tokens choose one expert each: 0,0,0,0,1,1,2,3. Experts work in parallel, so the step takes as long as the busiest one.

expert	tokens	over capacity (2)
0	4	2
1	2	0
2	1	0
3	1	0

average load = 2; busiest = 4 → imbalance 2× (time is set by the busiest expert) capacity = ceil(1 × 8 / 4) = 2; dropped tokens = 2

A small number of active parameters is not a complete latency model: one popular expert can be a bottleneck, and tokens beyond capacity may skip the expert. A balance term in the loss encourages even use.

Eight tokens, four experts. Raise the capacity factor to stop dropping tokens.

Predict first. Can one popular expert become a bottleneck?

Chapter context

New architectures alter where computation and memory are spent. Compare them by tracing one request through their actual operations rather than assuming every improvement reduces every cost.

TRY IT YOURSELF

Part 1. Do it hands-on

Count expert loads, imbalance and dropped tokens, then code the load statistics and the balance loss.

capacity factor1

Eight tokens choose one expert each: 0,0,0,0,1,1,2,3. Experts work in parallel, so the step takes as long as the busiest one.

expert	tokens	over capacity (2)
0	4	2
1	2	0
2	1	0
3	1	0

average load = 2; busiest = 4 → imbalance 2× (time is set by the busiest expert) capacity = ceil(1 × 8 / 4) = 2; dropped tokens = 2

tokens for expert 0 imbalance dropped at capacity 2

To do: Eight tokens are sent to experts 0,0,0,0,1,1,2,3 (one expert each). How many tokens does expert 0 get?
To do: Imbalance = busiest expert's load / average load. Average = 8 / 4 = 2. Type the imbalance
To do: Each expert accepts at most capacity = capacity_factor × tokens / experts = 1.0 × 8 / 4 = 2 tokens; extras are dropped (skipped). How many tokens are dropped?

Part 2. Do it in code

expert_loads(assignments, n_experts): how many tokens each expert receives (a list of counts).
imbalance(loads): max(loads) divided by the mean load. ValueError if the total is 0.
expert_capacity(tokens, n_experts, capacity_factor, k=1): ceil(capacity_factor · tokens · k / n_experts).
dropped_tokens(loads, capacity): tokens above the capacity, summed over experts.
balance_loss(fractions, probs): n_experts · Σ fractions_i · probs_i, where fractions_i is the share of tokens sent to expert i and probs_i the router's average probability for it. (Uniform routing gives 1.)

import math
import numpy as np

def expert_loads(assignments, n_experts):
    pass

def imbalance(loads):
    pass

def expert_capacity(tokens, n_experts, capacity_factor, k=1):
    pass

def dropped_tokens(loads, capacity):
    pass

def balance_loss(fractions, probs):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Can one popular expert become a bottleneck?

AYesBNo

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

For a toy MoE with shared weights S, E equal experts of P weights and k experts active per token:

N_{total}=S+EP,\quad N_{active}\approx S+kP.

This is a parameter-accounting model, not a full FLOP or communication model. Naively materialized attention scores use $T^2$ entries per head.

Worked example

With 2 billion shared parameters and eight experts of 1 billion each, total parameters are 10 billion. If two experts are selected per token, approximately 4 billion parameters participate along that token's path. The other experts still require storage or loading. A 4096×4096 naive score matrix contains 16,777,216 entries per head.

Reference code

The following example uses the Python standard library and does not call a remote service.

shared, experts, per_expert, active=2,8,1,2
assert shared+experts*per_expert == 10
assert shared+active*per_expert == 4
assert 4096**2 == 16777216

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Active parameters and resident parameters differ. Linear sequence complexity does not make every model property linear. Context-window capacity and demonstrated context use are separate claims.