Lab 01 — Router, Load Balancing, Capacity & the MoE Layer

Build a complete Mixture-of-Experts layer in pure Python, then break it on purpose: watch the router collapse, watch tokens get dropped silently, and watch the auxiliary loss save it.

The problem

MoE is usually explained as "more parameters, same FLOPs." That framing is true and useless, because it omits every part that makes MoE hard: routers collapse by default, capacity buffers drop tokens without raising anything, and sharding experts across chips costs seconds of network time per forward pass.

This lab implements the mechanism and all three failure modes, so that "MoE" stops being a word and becomes something you have debugged.

What you build

Key concepts

Files

Run

pytest test_lab.py -v                       # your lab.py — red until you implement
LAB_MODULE=solution pytest test_lab.py -v   # the reference — must be green (53 passed)
python solution.py                          # the worked example

Where to start

1. softmax → get test_softmax_survives_huge_logits green. Logits of 1000 must not produce NaN. This is the most common real-world MoE crash.
2. route_token → test_renormalized_gates_sum_to_one and its raw counterpart. Both behaviours are tested because both exist in the wild.
3. load_balance_loss → aim for test_balanced_routing_gives_aux_loss_of_exactly_one. If you do not get exactly 1.0, your formula is wrong. Remember f divides by total slots (tokens × k), not token count.
4. apply_capacity → first-come-first-served in token order; the tests depend on that determinism.
5. moe_forward → the residual fallback for fully-dropped tokens is the line that matters.
6. simulate_collapse → the money test.

The traps:

• Naive softmax → NaN on large logits. Tested.
• Dividing f by token count instead of slot count → your L_aux will not be 1.0.
• capacity_factor < 1.0 must raise — it drops tokens even under perfect balance.
• top_k == n_experts must give sparsity exactly 1.0. If not, your accounting is off.
• Unstable tie-breaking in the router destroys reproducibility. Break by expert index.

Success criteria

• LAB_MODULE=solution pytest test_lab.py -v → 53 passed.
• Your lab.py reaches 53 passed.
• python solution.py runs and you can explain all eight sections.
• You have seen aux_weight=0.0 collapse to max/mean = 4.10 and 0.01 hold at 1.00.
• You can state the capacity formula and predict a drop rate from an imbalance.
• You can convert an MoE's all-to-all traffic into seconds and say why it is fatal.

How this maps to the real stack

What is not a lie: the load-balance formula and its 1.0 calibration, the capacity formula, the total-vs-active distinction, and the all-to-all byte count. Those are exact, and those are what get asked about.

Extensions

1. Implement expert-choice routing. Invert the assignment — each expert picks its top tokens. Balance becomes guaranteed by construction, no aux loss, no dropping. Then discover why it cannot be used for autoregressive decode.
2. Implement auxiliary-loss-free balancing (DeepSeek-V3): maintain a per-expert bias added to router logits, nudged up or down by observed load. Compare final quality against the aux loss at matched balance — the claim is that the aux loss costs quality by forcing wrong routing, and you can measure that.
3. Add a real gradient. Wire this into Phase 03's autograd (or PyTorch) and actually train a tiny MoE on a toy task. Watch the aux-loss coefficient matter.
4. Model the pipelined-prefill fix. Extend expert_parallel_comm_bytes with a pipelined alternative that shards layers instead of experts, and reproduce the order-of-magnitude improvement in the transcript dissection's Claim 12.
5. Sweep the capacity factor against real routing entropy. Generate routings at varying levels of imbalance and plot drop rate vs capacity_factor — that curve is the one you actually use to pick the hyperparameter.

Interview / resume bullets

• "Implemented a Mixture-of-Experts layer from first principles in pure Python — numerically stable top-k routing with gate renormalization, Switch-style load-balancing loss, ST-MoE router z-loss, capacity-factor dispatch with drop/pad accounting, shared experts, and expert-parallel communication cost modelling — verified by 53 tests covering boundary and determinism cases."
• "Reproduced and quantified the three canonical MoE failure modes: router collapse via rich-get-richer feedback (and its prevention by an auxiliary loss), silent token dropping through fixed capacity buffers, and the all-to-all communication wall of naive expert parallelism (32 GB and ~2.6 s per 8k-token prefill on a 60-layer model)."
• Interview-ready: "Active parameters for FLOPs, total parameters for memory. A 700B/52B MoE trains like a 52B model and stores like a 700B one — and confusing the two makes your capacity plan wrong by an order of magnitude."