Part II — Systems

Mixture of Experts — parameters without the FLOPs

Lesson 09 gave you the last tool you need to train an arbitrarily large dense model: tensor, pipeline, and sequence parallelism, composed into a 3D mesh that splits any matmul, any layer stack, any sequence across hundreds of GPUs. So you can now build a 70B, a 400B, a model of any size the cluster can hold. But there is a tax you cannot escape by splitting harder: in a dense model every token flows through every parameter, so the FLOPs per token grow in lockstep with the parameter count, and the budget caps both at once. This lesson breaks that lockstep. The idea is sparsity — only fire a fraction of the network per token — and it lets total parameters grow while the FLOPs you pay stay flat.

1 · The decoupling: capability tracks total params, cost tracks active params

The whole of MoE rests on prying apart two quantities that a dense model welds together. In a dense Transformer they are the same number N: the parameters that store knowledge are exactly the parameters each token multiplies through, so capability and per-token cost are one knob. MoE introduces a second number.

The whole approach rests on an empirical bet, not a theorem: that model quality is driven far more by N_total than by N_act. The intuition is that knowledge is sparse and localizable — any given token only needs a small, relevant slice of what the model knows (a coding token doesn't consult the French-poetry weights), so storing more in inactive experts adds capability the token never pays FLOPs to consult. A model can know a great deal as long as the knowledge is stored somewhere and the router can find the right sliver. So if you could grow N_total while holding N_act fixed, you would buy capability at constant FLOPs/token — a "free" parameter from the budget's point of view. That gap, N_total ≫ N_act, is the entire prize. But it is an empirical regularity that holds in practice, not a guarantee — it can break: if the router collapses onto a few experts, or experts learn redundant rather than specialized knowledge (capacity waste), the extra parameters stop buying capability and you are paying memory and comm for ballast. The rest of the lesson is how to build a network where the bet pays off, and what it costs you elsewhere (memory and communication — there is no actual free lunch, only a moved one).

2 · The router: replace the MLP with E experts and a top-k gate

Where does the gap come from mechanically? Recall from lesson 2 that the per-block MLP holds about two-thirds of every block's parameters and is applied independently at each token position — it is the network's per-token feature store. MoE attacks exactly this sublayer. Replace the single shared MLP with E independent MLPs — the experts — each the same shape as the original. Then, per token, a small router (a single linear layer W_r of shape d × E) decides which experts that token visits.

# x: (tokens, d)   E experts, each an MLP(d -> d_ff -> d)
logits  = x @ W_r                 # (tokens, E)   router scores, W_r is d×E
probs   = softmax(logits, dim=-1) # (tokens, E)   gate weights over experts
top     = topk(probs, k)          # pick the k highest-scoring experts / token
top     = [(e, g/sum_of_topk_g) for (e,g) in top]  # renormalize gates over the chosen k (Mixtral)
y = 0
for (e, g) in top:                # k of them, k ≈ 1 or 2
    y += g * expert[e](x)         # run ONLY the chosen experts, weight by gate
# tokens not routed to expert e never touch its weights → its FLOPs aren't paid

Read what changed. The router runs for every token (it is tiny: d·E FLOPs, negligible against an expert MLP's ~16d²). But only the top-k experts actually execute. With E = 8 experts and k = 2, each token touches 2 of the 8 MLPs: the block stores 8× the MLP parameters but pays for only 2 of them per token. The gate weight g (the softmax probability, renormalized over the chosen k) scales each expert's contribution so the router is differentiable and learns which expert is good for which token — by backprop, with no explicit supervision of the assignment.

Why top-k and not all? Routing to all E would be a dense model with E parallel MLPs — capability up, but FLOPs up E× too, defeating the point. Routing to k = 1 (Switch Transformer) minimizes FLOPs but gives the router no way to blend specialties and makes training twitchy. k = 2 is the workhorse (Mixtral, GShard): enough to mix two experts and stabilize gradients, still only 2/E of the compute. The active MLP parameters are k/E of the total expert parameters — that ratio is the decoupling, made concrete.

3 · The failure that needs fixing: expert collapse, load balancing, capacity

Train the router naively and it breaks in a specific, predictable way: expert collapse. Early on, a few experts get marginally better, so the router sends them slightly more tokens, so they get more gradient and improve faster, so the router sends them even more — a rich-get-richer loop. Within a few thousand steps the router funnels nearly all tokens to a handful of experts while the rest receive almost nothing, never train, and become dead weight. You have paid for E× parameters and are effectively running a dense model with the others as ballast. The router's argmax is a positive-feedback amplifier, and left alone it always collapses.

The fix is to make balance an explicit objective. Two approaches dominate:

The capacity factor exists for a systems reason that section 4 makes clear: experts run as fixed-shape batched matmuls on hardware, so each must be given a statically known number of tokens. Perfect balance is impossible per batch, so you over-provision by cap and drop the rest. A drop is a small quality hit (that token's MLP update is skipped this layer); at cap = 1.25 with a healthy balancing term, drop rates sit at a few percent — the price of turning a dynamic routing decision into a static tensor shape.

4 · Systems: expert parallelism and the all-to-all bottleneck

Now the bill MoE actually moves the cost to. Two facts collide. First, memory holds all experts: N_total bytes of weights are resident even though each token uses k/E of them — the 16-bytes/param ledger of lesson 5 now applies to the full expanded parameter count, so an 8× expert expansion is an ~8× MLP-memory bill. Second, the experts are too many to replicate, so you shard them across GPUs: expert parallelism (EP) places different experts on different devices.

That sharding creates a new communication pattern the dense stack never had. A token arrives on the GPU that ran attention for it, but the expert it was routed to lives on some other GPU. So before the expert MLP you must ship every token to wherever its chosen expert sits, and after, ship the result back. Every GPU sends a different slice to every other GPU — an all-to-all collective, twice per MoE layer (dispatch and combine).

All-to-all is the MoE analogue of the all-reduce/all-gather that defined Part II's earlier comm walls — and it is the new bottleneck. Its volume scales with token count and is acutely sensitive to network topology: cheap over NVLink inside a node, brutal across the slower inter-node fabric, which is why expert parallelism, like tensor parallelism (lesson 9), wants to stay within high-bandwidth islands. The arithmetic-intensity lens from lesson 6 returns here: all-to-all moves a lot of bytes for relatively little compute, so a poorly placed MoE layer is communication-bound and the saved FLOPs are eaten by stalled GPUs waiting on the network. The kernel-level view of these collectives — how the all-to-all and grouped-GEMM expert kernels are written, fused, and debugged — lives in gpu_kernels · 16 (debugging a kernel / system design), and the topology of placing experts across a cluster is treated in the reinforcement-learning track's distributed-topology lessons.

5 · Real models and the dense-vs-MoE trade

The decoupling is not theoretical — it ships in frontier models. The number that matters is the active/total ratio: how much of the model each token pays for.

The verdict is honest, not cheerleading: MoE buys capability-per-FLOP — invaluable when compute is the binding constraint and you have the memory and fast interconnect to absorb the premium. It is the wrong trade when memory is the binding constraint (it multiplies the very thing lesson 5 showed already overflows the device) or when your interconnect can't hide the all-to-all. You moved the cost off the FLOP axis; you did not delete it.

6 · Drive the economics

The widget makes the decoupling tangible. Move E (experts), k (top-k), the expert hidden size, and the number of layers. Watch the two bars: total params (blue, what you store, what drives capability) and active params/token (green, what each token touches, what sets FLOPs). The point to internalize: crank E with k fixed and the blue bar explodes while green — and FLOPs/token with it — stays dead flat. That flat green line at a rising blue ceiling is the free parameter, and the gap between them is exactly the capability you bought without paying compute. Then look at total memory climbing with the blue bar: that is where the bill landed.

Notice the asymmetry the widget exposes: raising k moves both bars (you pay for more experts/token), but raising E moves only the blue one. That is why frontier MoEs push E high and keep k small — they are buying parameters off the cheap axis. And it is why this lesson hands you a problem rather than a solution: you now have four knobs that all trade differently against the budget.

Failure modes & checklist

Checkpoint

Where this points next

MoE just handed you the most powerful lever in the course — buy capability off the cheap (FLOP) axis by spending the expensive (memory + comm) one. With it, Part II is complete: you can now specify what to train at any scale — dense or sparse, any size the cluster can hold, with the FLOP/parameter trade tuned to your budget. But specifying the model and executing the run are different problems. That training run is not an afternoon job: it is weeks of wall-clock on thousands of GPUs, and at that scale the rare becomes routine — a GPU dies mid-step, the loss spikes and threatens to diverge, a node drops off the fabric, and every hour the job is stalled or silently corrupted is thousands of dollars and days of schedule burned. Keeping a run alive and healthy — checkpointing, monitoring, detecting and recovering from failures, reading the loss curve — is its own engineering discipline, and it is what stands between a good configuration and a finished model. Next: 11 · Running the job — executing a weeks-long distributed run.

Interview prompts

• What exactly does MoE decouple, and why does that beat the dense budget? (§1 — total parameters (capacity ≈ capability, what memory holds) from active parameters (the k/E that each token touches, = 2·N_act FLOPs/token). A dense model welds them (N_act=N_total); MoE grows total while pinning active, buying capability at constant FLOPs.)
• Walk the router forward pass and say why k=2 is typical. (§2 — router linear W_r (d×E) → softmax over experts → top-k; run only those k experts, weight by renormalized gate g, sum. k=1 (Switch) is cheapest but unstable and can't blend; k=2 mixes two specialties and stabilizes gradients at 2/E compute.)
• What is expert collapse and how do you prevent it? (§3 — a rich-get-richer feedback loop where the router funnels tokens to a few experts that then train faster, leaving the rest dead. Fix with a load-balancing aux loss (α·E·Σ f_ep_e, minimized at uniform usage) or an aux-loss-free per-expert router bias nudged by recent load.)
• Why is there a capacity factor, and what happens at overflow? (§3 — experts run as fixed-shape batched matmuls, so each needs a statically known token count; cap≈1.25 over-provisions the per-expert buffer and overflow tokens are dropped (skip the MLP via the residual). Higher cap = fewer drops, more compute; it's the cost of making dynamic routing a static tensor shape.)
• What new communication does MoE introduce and when does it dominate? (§4 — expert parallelism shards experts across GPUs, so each MoE layer needs two all-to-all collectives (dispatch tokens to their expert's GPU, combine results back). It dominates and makes the layer comm-bound when EP spans slow inter-node links; keep EP inside NVLink islands.)
• Mixtral is "8×7B" — why isn't it 56B, and what are its active/total figures? (§5 — attention and embeddings are shared, only the MLPs are replicated, so total ≈47B not 56B; with k=2 of 8 experts, active ≈13B per token — a ~3.6× total/active ratio.)
• When is MoE the wrong choice? (§5 — when memory is the binding constraint (it multiplies the parameter count the 16 B/param ledger already overflows) or when the interconnect can't hide the all-to-all; MoE trades FLOPs for memory and comm, so it only wins when compute is what's scarce.)