Phase 02 — Mixture of Experts From Scratch

The architecture that lets you add knowledge without adding cost per token — and the one that turns a modelling decision into a memory problem, a communication problem, a stability problem, and a data problem all at once. Feinberg's slides put both halves in one breath: "MoE scaling laws are better, but have implications for token hunger. We're running out of internet!" And the serving half, from the interview: an MoE "uses a lot more parameters," so when you shard it, "that token might live on the first TPU, but it needs to go to the last TPU."

Why this phase exists

MoE is the single most consequential architecture decision in modern pre-training, and it is the one most often described in a way that makes it sound free. It is not free. It is a cost-shifting device, and a pre-training lead has to know exactly where the cost lands:

You gain	You pay in
More parameters at the same FLOPs per token	HBM — every expert must be resident
Better loss per unit of compute	Communication — all-to-all every layer if you shard experts
Specialization across the corpus	Stability — routers collapse by default
	Dropped tokens — capacity is a fixed buffer, and overflow is silent
	Data hunger — MoE's compute-optimal `D` is larger, and unique data is the scarce resource

Every one of those five is something you implement and measure in this phase's lab. And the last two are the ones that surprise people: a production MoE can be silently dropping 5–10% of its tokens, and no exception is ever raised.

This phase also sets up the pivot in Feinberg's Flash 2.0 story. The communication cost you compute here is precisely the wall his team hit — and the fix was not a better kernel but changing which axis you shard along, covered in the transcript dissection, Claim 12.

Concept map

                          token (d_model vector)
                                   │
                          ┌────────▼────────┐
                          │     ROUTER      │  d_model x n_experts — tiny, and the
                          │  logits→softmax │  most fragile part of the whole model
                          └────────┬────────┘
                                   │ top-k, renormalized gates
              ┌────────────────────┼────────────────────┐
              ▼                    ▼                    ▼
        ┌──────────┐         ┌──────────┐         ┌──────────┐
        │ EXPERT 0 │  ...    │ EXPERT i │  ...    │ EXPERT E │   only k of these run
        └──────────┘         └──────────┘         └──────────┘
              │                    │                    │
              └────────── weighted sum by gate ─────────┘
                                   │
                    + shared expert (runs for EVERY token)
                                   │
                                   ▼
                                output

  CONSTRAINTS BOLTED ON TOP
  ├─ capacity   = cf × tokens × k / E     overflow is DROPPED (silently)
  ├─ L_balance  = E · Σ f_i·P_i           1.0 = perfect; larger = collapsing
  ├─ L_z        = mean(logsumexp²)        keeps logits from saturating the softmax
  └─ comms      = 2 all-to-alls per layer if experts are sharded across chips

What you will be able to do

Implement a top-k router with renormalized gates, and explain what not renormalizing silently does to the layer's output scale.
Write the Switch load-balance loss and explain why it multiplies a non-differentiable token fraction by a differentiable probability mass — and why balanced routing gives exactly 1.0.
Write the router z-loss and explain the two distinct failures it prevents (bf16 overflow, and a saturated softmax whose gradient vanishes).
Compute expert capacity, and quantify the drop-rate / wasted-compute trade-off that tuning capacity_factor actually buys.
Explain why a fully-dropped token is a silent quality loss, and how a shared expert removes that failure mode entirely.
Separate total parameters (memory) from active parameters (FLOPs) and never confuse them again.
Compute the all-to-all traffic of naive expert parallelism and state, in seconds, why it is fatal for an interactive product.
Simulate router collapse and demonstrate the auxiliary loss preventing it.

The lab

Lab	What you build
Lab 01 — Router, Load Balancing, Capacity & the MoE Layer	A complete MoE layer in pure stdlib: stable softmax, top-k router with gate renormalization, Switch load-balance loss, router z-loss, capacity/drop/pad accounting, a full forward pass with shared expert and residual fallback, total-vs-active parameter counting, expert-parallel communication cost, and a router-collapse simulator that shows the aux loss working

Success criteria. LAB_MODULE=solution pytest test_lab.py -v → 53 passed; python solution.py prints the eight-part worked example, including the collapse demonstration.

Deliverables checklist

Balanced routing gives a load-balance loss of exactly 1.0 in my implementation.
My softmax survives logits of 1000 without producing NaN.
I can state the capacity formula and predict the drop rate for a given imbalance.
I have watched a router collapse with aux_weight = 0 and recover with 0.01.
I can compute total vs active parameters and say which goes into 6ND.
I can compute the all-to-all bytes for a realistic MoE and convert it to seconds.
I can explain what a shared expert buys, in one sentence.

Key takeaways

Parameters scale with E; FLOPs scale with k. That single sentence is what MoE is.
top_k == n_experts is a dense model wearing an MoE costume. Sparsity ratio 1.0.
Routers collapse by default. Rich-get-richer is the natural dynamic; the auxiliary loss is what stops it, and its coefficient (~0.01) is one of the most finicky hyperparameters in pre-training.
L_balance = 1.0 means perfectly balanced. Memorize the calibration point — it makes the dashboard readable at a glance.
Dropped tokens are silent. No exception, no log line; the token rides the residual and your loss is quietly worse. Watch drop_rate like a hawk.
Capacity factor trades dropping against wasted compute. There is no setting that avoids both.
A shared expert guarantees every token gets some FFN, which removes the fully-dropped failure mode at the cost of always-on compute.
Naive expert parallelism costs seconds of pure network time. That is the wall, and the fix is to change the sharding axis, not the kernel.
MoE trades a compute problem for a data problem. Better scaling law, larger optimal D, and unique tokens are what is actually running out.

References

Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated MoE Layer, 2017 — https://arxiv.org/abs/1701.06538
Fedus, Zoph & Shazeer, Switch Transformers, 2021 — https://arxiv.org/abs/2101.03961 (the load-balance loss, capacity factor, k=1 routing)
Zoph et al., ST-MoE: Designing Stable and Transferable Sparse Expert Models, 2022 — https://arxiv.org/abs/2202.08906 (the router z-loss)
Lepikhin et al., GShard, 2020 — https://arxiv.org/abs/2006.16668 (expert parallelism and the all-to-all)
Clark et al., Unified Scaling Laws for Routed Language Models, 2022 — https://arxiv.org/abs/2202.01169 (the MoE scaling law Feinberg cites)
DeepSeek-AI et al., DeepSeek-V3 Technical Report, 2024 — https://arxiv.org/abs/2412.19437 (shared experts, fine-grained experts, auxiliary-loss-free balancing)
Jiang et al., Mixtral of Experts, 2024 — https://arxiv.org/abs/2401.04088
Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides

Frontier Pre-Training Lead — Scaling Laws, MoE, Distillation & Inference Co-Design