Lab 02 — Mixture of Experts (MoE)
Phase: 02 — Model Architecture Mastery | Difficulty: ⭐⭐⭐⭐⭐ | Time: 3–4 hours
This lab implements the MoE layer used in GPT-4, Mixtral, and Switch Transformer.
What you build
- TopKRouter — linear routing with top-K expert selection and Switch Transformer auxiliary load-balancing loss
- ExpertLayer — FFN expert (w1 → GELU → w2)
- MoEBlock — routes each token to K experts, combines weighted outputs
- compute_expert_utilization — detects load imbalance
Key concepts
| Concept | What to understand |
|---|---|
| Conditional compute | Each token uses only K of N experts (sparse activation) |
| Aux loss | n_experts * Σ(f_i * p_i) — penalizes load imbalance |
| Expert collapse | When all tokens route to same expert — training failure mode |
| Top-1 vs Top-2 | Switch Transformer uses top-1; Mixtral/GPT-4 uses top-2 |
Files
| File | Purpose |
|---|---|
lab.py | Stubs with TODO markers |
solution.py | Complete working implementation |
test_lab.py | pytest suite |
requirements.txt | Dependencies |
Run
pip install -r requirements.txt
python lab.py
pytest test_lab.py -v