{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Module 11: Mixture of Experts" ], "id": "22f940b0-6bc8-4297-98ee-7f839bb52298" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [], "id": "b67beba6-269b-4df8-9219-cb7279b2a1e2" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "21940ea0-6781-4eba-94bf-1dc59eec5306" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Introduction\n", "\n", "Every model we have built so far is **dense**: each token is pushed\n", "through every single parameter. Double the parameters and you double the\n", "compute for *every* token. That is the wall the frontier ran into — and\n", "**Mixture of Experts (MoE)** is the most successful way around it.\n", "\n", "An MoE replaces the one feed-forward network in a transformer block with\n", "*many* parallel **expert** FFNs plus a small **router** that sends each\n", "token to only a few of them. The model can then hold a huge number of\n", "parameters — all the experts — while each token only pays for the\n", "handful it actually uses. Total parameters and per-token compute are\n", "**decoupled**.\n", "\n", "This is not a curiosity. **Mixtral 8×7B** (8 experts, top-2 routing),\n", "**DeepSeek-MoE**, and several frontier flagships are MoE models: they\n", "serve the quality of a very large model at the cost of a much smaller\n", "one.\n", "\n", "### What You’ll Learn\n", "\n", "- Why **total parameters ≫ active parameters** is the entire point of\n", " MoE\n", "- How a **top-k router** scores experts and dispatches each token\n", "- How to build a sparse **`MoELayer`** from scratch (gather → run →\n", " scatter)\n", "- Why routers **collapse** and how a **load-balancing loss** prevents it\n", "- The sparsity factor `E/k` and the memory-vs-compute trade-off it\n", " encodes\n", "\n", "### Prerequisites\n", "\n", "- [Module 06: Transformer](../m06_transformer/lesson.qmd) — the FFN that\n", " MoE replaces\n", "- [Module 07: Training](../m07_training/lesson.qmd) — the loss the aux\n", " term is added to\n", "- Comfort with `torch` tensor indexing (gather / scatter-add)\n", "\n", "## Intuition: A Dense FFN Runs Everything, Every Time\n", "\n", "Look back at the transformer block. Its feed-forward network is where\n", "most of the parameters live — and it fires in full for every token,\n", "whether the token is the word “the” or a rare technical term.\n", "Intuitively that is wasteful: different tokens want different\n", "transformations, but a dense FFN forces them all through the same one.\n", "\n", "MoE turns that single FFN into a *committee* of experts and hires a\n", "**router** to pick which experts see each token:\n", "\n", "- **Experts** — several independent FFNs (say 8). Each can specialize.\n", "- **Router** — a tiny linear layer that scores all experts for a token\n", " and keeps the **top-k** (often just 1 or 2).\n", "- **Combine** — the chosen experts’ outputs are summed, weighted by the\n", " router.\n", "\n", "Because only `k` of the `E` experts run per token, the compute is that\n", "of a `k`-expert FFN, no matter how many experts the model *owns*. Add\n", "experts to grow capacity; keep `k` fixed to keep the bill flat.\n", "\n", "> **Key Insight**\n", ">\n", "> MoE decouples **capacity** from **compute**. A dense model ties them\n", "> together — more parameters means more work per token. An MoE with `E`\n", "> experts and top-`k` routing has `E` experts’ worth of parameters but\n", "> does `k` experts’ worth of work per token. The ratio `E/k` is the\n", "> **sparsity factor**: free capacity, paid for in memory rather than\n", "> FLOPs.\n", "\n", "## The Math: Routing and Load Balancing\n", "\n", "**Routing.** For a token $x$, a linear gate produces one score per\n", "expert; a softmax turns the scores into probabilities, and we keep the\n", "top $k$:\n", "\n", "$$g(x) = \\text{softmax}(x W_g) \\in \\mathbb{R}^E, \\qquad \\mathcal{T} = \\text{top-}k(g(x))$$\n", "\n", "The layer’s output sums the chosen experts, weighted by their\n", "(renormalized) gate values:\n", "\n", "$$\\text{MoE}(x) = \\sum_{i \\in \\mathcal{T}} \\tilde{g}_i(x)\\, E_i(x), \\qquad \\tilde{g}_i = \\frac{g_i}{\\sum_{j \\in \\mathcal{T}} g_j}$$\n", "\n", "where $E_i$ is the $i$-th expert FFN. Only the $k$ experts in\n", "$\\mathcal{T}$ ever run.\n", "\n", "**Load balancing.** Left to itself the router **collapses** — it\n", "discovers a few strong experts early and sends everything to them,\n", "starving the rest (a starved expert gets no gradient and effectively\n", "dies). To counter this, we add an auxiliary loss (Shazeer 2017; Switch\n", "Transformer, Fedus 2021). Over a batch of $N$ tokens routed top-$k$:\n", "\n", "$$\\mathcal{L}_{\\text{aux}} = E \\cdot \\sum_{i=1}^{E} f_i \\, P_i$$\n", "\n", "where $f_i$ is the fraction of dispatch slots sent to expert $i$ and\n", "$P_i$ is the mean router probability for expert $i$. Both vectors sum to\n", "1, so the loss is **minimized at uniform routing** ($f_i = P_i = 1/E$),\n", "giving $\\mathcal{L}_{\\text{aux}} = 1$. Any imbalance pushes it above 1.\n", "Training minimizes\n", "$\\mathcal{L}_{\\text{task}} + \\alpha\\,\\mathcal{L}_{\\text{aux}}$ with a\n", "small $\\alpha$, so the router is nudged to keep every expert fed.\n", "\n", "## Code: An Expert and a Router from Scratch\n", "\n", "Everything lives in `moe.py`. First we load it and build the two pieces\n", "— an `Expert` (an ordinary GELU FFN) and the top-k `MoERouter`." ], "id": "bce2630e-40f8-443e-a413-76d66f72a1f2" }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "PyTorch 2.12.1+cu130 on cpu\n", "probs shape: (5, 8) (softmax over 8 experts)\n", "chosen experts: [[6, 5], [7, 0], [1, 0], [1, 6], [2, 4]]\n", "chosen weights: [[0.5210000276565552, 0.4790000021457672], [0.5210000276565552, 0.4790000021457672], [0.5059999823570251, 0.49399998784065247], [0.5180000066757202, 0.4819999933242798], [0.5059999823570251, 0.49399998784065247]]\n", "weights sum to 1? [1.0, 1.0, 1.0, 1.0, 1.0]" ] } ], "source": [ "import importlib.util\n", "import sys\n", "from pathlib import Path\n", "\n", "import torch\n", "\n", "# Load moe.py (directory name starts with a digit, so use importlib)\n", "spec = importlib.util.spec_from_file_location(\"moe\", Path(\"moe.py\").resolve())\n", "moe = importlib.util.module_from_spec(spec)\n", "sys.modules[\"moe\"] = moe\n", "spec.loader.exec_module(moe)\n", "\n", "device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n", "print(f\"PyTorch {torch.__version__} on {device}\")\n", "\n", "router = moe.MoERouter(embed_dim=64, num_experts=8, top_k=2)\n", "x = torch.randn(5, 64) # 5 tokens\n", "out = router(x)\n", "print(f\"probs shape: {tuple(out.probs.shape)} (softmax over 8 experts)\")\n", "print(f\"chosen experts: {out.topk_indices.tolist()}\")\n", "print(f\"chosen weights: {out.topk_weights.round(decimals=3).tolist()}\")\n", "print(f\"weights sum to 1? {out.topk_weights.sum(dim=-1).round(decimals=4).tolist()}\")" ], "id": "c0ef8db2" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Each token picked its 2 best experts, and the two kept weights were\n", "renormalized to sum to 1 — the discarded probability mass is simply\n", "dropped.\n", "\n", "## Code: The Sparse MoE Layer\n", "\n", "The layer ties experts and router together. The trick is the\n", "**dispatch**: for each expert, gather exactly the tokens routed to it,\n", "run the expert once on that batch, then **scatter-add** the weighted\n", "results back. That is the same gather → run → scatter pattern real MoE\n", "kernels use to avoid running every expert on every token." ], "id": "7946c1f7-4dc3-4815-b71b-af33d15826c5" }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "input shape: (2, 16, 64)\n", "output shape: (2, 16, 64) (same as input — a drop-in FFN)\n", "load-balancing aux loss: 1.0056 (1.0 = perfectly balanced)" ] } ], "source": [ "moe_layer = moe.MoELayer(embed_dim=64, ff_dim=256, num_experts=8, top_k=2)\n", "x = torch.randn(2, 16, 64) # (batch, seq, embed)\n", "output, aux = moe_layer(x)\n", "\n", "print(f\"input shape: {tuple(x.shape)}\")\n", "print(f\"output shape: {tuple(output.shape)} (same as input — a drop-in FFN)\")\n", "print(f\"load-balancing aux loss: {aux.item():.4f} (1.0 = perfectly balanced)\")" ], "id": "b07944dd" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The forward pass returns **two** things: the output *and* the\n", "load-balancing loss for this batch. During training you add\n", "`alpha * aux` to your task loss.\n", "\n", "**Verifying the dispatch math.** With `top_k=1`, each token’s output\n", "must be *exactly* its one chosen expert applied to it (its weight\n", "renormalizes to 1). We can check that directly:" ], "id": "2da4ff6a-89ba-48d1-8410-afe715b671b5" }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "top-1 output matches the chosen expert exactly: True" ] } ], "source": [ "layer1 = moe.MoELayer(embed_dim=32, ff_dim=64, num_experts=4, top_k=1)\n", "layer1.eval()\n", "x = torch.randn(1, 6, 32)\n", "out1, _ = layer1(x)\n", "\n", "flat = x.reshape(-1, 32)\n", "route = layer1.router(flat)\n", "manual = torch.stack([\n", " layer1.experts[route.topk_indices[t, 0]](flat[t:t+1])[0]\n", " for t in range(flat.shape[0])\n", "])\n", "print(\"top-1 output matches the chosen expert exactly:\",\n", " torch.allclose(out1.reshape(-1, 32), manual, atol=1e-5))" ], "id": "393bc694" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Total vs Active Parameters\n", "\n", "Here is the payoff, made concrete. `count_parameters()` reports what the\n", "layer *owns* versus what a single token *touches*:" ], "id": "b02ca82a-1a39-489b-831e-7dd3445371e7" }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "E= 2, k=2 → total 4.2M | active 4.2M | sparsity 1.0×\n", "E= 8, k=2 → total 16.8M | active 4.2M | sparsity 4.0×\n", "E=32, k=2 → total 67.2M | active 4.2M | sparsity 15.9×" ] } ], "source": [ "for E in (2, 8, 32):\n", " c = moe.MoELayer(embed_dim=512, ff_dim=2048, num_experts=E, top_k=2).count_parameters()\n", " print(f\"E={E:2d}, k=2 → total {c['total']/1e6:6.1f}M | \"\n", " f\"active {c['active']/1e6:5.1f}M | sparsity {c['sparsity_factor']:.1f}×\")" ], "id": "c3bb2338" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Going from 2 to 32 experts multiplies total parameters ~16× while the\n", "**active** parameters per token barely move. Drive the same trade-off\n", "yourself:" ], "id": "de251dd8-32a4-4ba1-a8ba-c66eb3df49ad" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "733e5ff7-4eea-4cda-8e5a-99cfeb7bb686" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "0165ae20-3a7c-47d6-b334-383b1457355c" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "5e951e08-b5d8-4555-ab09-98336af9d719" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "bf093c2b-297f-440a-8d4c-d50d4590debc" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "878948f1-694c-402a-856f-f29bb2b567a9" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "9ee6692a-8a10-4a62-8ca6-b0f639e833ef" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "4eca29b4-c34d-4715-8a70-79693de23e22" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "5b2c2af8-6d26-432e-93b2-b848258d8f58" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "a71e672c-f9cc-4be4-997a-7aeaa0d00777" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try This**\n", ">\n", "> 1. **Add experts, watch compute stay flat.** Slide `E` from 8 to 64\n", "> in the sparsity meter. The *total* bar grows; the *active* bar\n", "> barely moves. That gap is capacity you get almost for free (in\n", "> FLOPs — you still pay in memory).\n", "> 2. **Raise top-k.** Push `k` from 1 to 4. Active compute climbs and\n", "> the sparsity factor `E/k` shrinks — top-k is the compute dial.\n", "> 3. **Break the router.** In the routing simulator, drag the **skew**\n", "> up. Tokens pile onto a few experts, the others starve, and the\n", "> load-balancing loss climbs far above 1.0. This is the collapse the\n", "> auxiliary loss exists to prevent.\n", "> 4. **Balance it back.** Return skew to 0 and watch every bar settle\n", "> near the `1/E` line and the loss fall back to ≈ 1.0.\n", "\n", "## Common Pitfalls\n", "\n", "1. **Forgetting the auxiliary loss.** Without it the router collapses\n", " onto a few experts within a few hundred steps, and the rest become\n", " dead weight. Always add `alpha * aux_loss` (a common $\\alpha$ is\n", " 0.01).\n", "\n", "2. **Confusing total with active parameters.** An “8×7B” model does not\n", " run 56B parameters per token — with top-2 it runs about 13B. Report\n", " both numbers; they answer different questions (memory vs. FLOPs).\n", "\n", "3. **Renormalizing (or not) the top-k weights.** After picking the\n", " top-k, their gate values no longer sum to 1. Renormalize them (as\n", " here) so the combined output is a proper weighted average and its\n", " scale is stable.\n", "\n", "4. **Training instability from hard routing.** The top-k operation is\n", " non-differentiable in the *choice* itself; gradients flow only\n", " through the kept weights. This makes MoE routers more finicky to\n", " train than a dense FFN — the load-balancing loss and careful\n", " initialization matter.\n", "\n", "5. **Assuming MoE is always a win.** MoE trades FLOPs for memory and\n", " bandwidth. If you are memory-bound (a single GPU, long context),\n", " holding many experts can cost more than the compute it saves.\n", "\n", "## Exercises\n", "\n", "### Exercise 1: Sparsity of a real config\n", "\n", "Mixtral 8×7B uses `E=8`, `top_k=2`. What is its sparsity factor? Roughly\n", "how many parameters are active per token if the total is ~47B and the\n", "non-expert (attention + embeddings) share is ~2B?" ], "id": "740bd5f3-b2e2-47a1-b4a3-f92fd2f318bf" }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "sparsity factor E/k = 4×\n", "active params per token ≈ 13.2B (of 47B total)" ] } ], "source": [ "# E/k sparsity factor; active ≈ non_expert + (k/E) * expert_params\n", "total, non_expert = 47e9, 2e9\n", "expert_params = total - non_expert\n", "active = non_expert + (2 / 8) * expert_params\n", "print(f\"sparsity factor E/k = {8/2:.0f}×\")\n", "print(f\"active params per token ≈ {active/1e9:.1f}B (of {total/1e9:.0f}B total)\")" ], "id": "c5514456" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 2: Watch a router collapse\n", "\n", "Build a layer, feed it the *same* few tokens many times, and skew the\n", "gate by hand to one expert. Print `expert_utilization` before and after\n", "— most of the mass should move to a single expert.\n", "\n", "### Exercise 3: The load-balancing floor\n", "\n", "Show numerically that `load_balancing_loss` bottoms out at 1.0.\n", "Construct uniform `probs` and round-robin `topk_indices` for `E` experts\n", "and confirm the loss is ≈ 1.0; then concentrate the routing and confirm\n", "it rises.\n", "\n", "## Summary\n", "\n", "Key takeaways:\n", "\n", "1. **MoE decouples capacity from compute.** Many expert FFNs, a router\n", " that keeps only the top-`k` per token: total parameters ≫ active\n", " parameters.\n", "\n", "2. **The router is a tiny top-k gate.** A single linear layer scores\n", " experts; softmax + top-`k` selects them; the kept weights are\n", " renormalized to sum to 1.\n", "\n", "3. **Dispatch is gather → run → scatter.** Each expert runs once, on\n", " exactly the tokens routed to it — never all experts on all tokens.\n", "\n", "4. **Routers collapse without help.** A load-balancing auxiliary loss\n", " $E\\sum_i f_i P_i$ (minimum 1.0 at uniform routing) keeps every\n", " expert fed.\n", "\n", "5. **The sparsity factor is `E/k`.** It is free capacity paid for in\n", " memory, not FLOPs — the trade-off behind Mixtral, DeepSeek-MoE, and\n", " modern frontier MoEs.\n", "\n", "## What’s Next\n", "\n", "We now have both halves of the modern FFN story: the dense gated FFN\n", "(SwiGLU, in m06) and the sparse one (MoE, here). The frontier keeps\n", "building on these — MoE routing variants (expert choice, shared\n", "experts), long-context attention, and the alignment and reasoning stages\n", "that turn a pretrained model into an assistant.\n", "\n", "## Going Deeper\n", "\n", "- Shazeer et al., *Outrageously Large Neural Networks: The\n", " Sparsely-Gated Mixture-of-Experts Layer* (2017) — the modern MoE layer\n", " and its balancing loss. \n", "- Fedus, Zoph & Shazeer, *Switch Transformers* (2021) — top-1 routing at\n", " scale; the load-balancing loss used here.\n", " \n", "- Jiang et al., *Mixtral of Experts* (2024) — a strong open MoE (8\n", " experts, top-2), with the total-vs-active numbers this lesson quotes.\n", " " ], "id": "35278c6d-4b67-4da4-9caa-0bd128fe8ace" } ], "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "name": "python3", "display_name": "Python 3 (ipykernel)", "language": "python", "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3" }, "language_info": { "name": "python", "codemirror_mode": { "name": "ipython", "version": "3" }, "file_extension": ".py", "mimetype": "text/x-python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.14" } } }