How Mixture-of-Experts LLMs Work on Ship of Theseus

The Big Picture

Mon, 01 Jan 0001 00:00:00 +0000

Lesson 1 — The Big Picture: Sparse vs. Dense A 10-lesson technical course on the mechanics of Mixture-of-Experts language models. Audience: undergraduate level, already comfortable with dense decoder-only LLMs (you know next-token prediction, attention, the feed-forward network, and fine-tuning). Prerequisites: linear algebra (matrix products), multivariate calculus (gradients, the chain rule), basic probability. We use standard notation and define each new symbol as it appears. Course table of contents # Lesson Core technical content 1 The Big Picture (this file) Conditional computation; $P_{\text{tot}}$ vs. $P_{\text{act}}$; compute $\approx 6P_{\text{act}}D$ but memory $\propto P_{\text{tot}}$; the dense contrast 2 From the Dense FFN to a Layer of Experts The FFN holds $\tfrac23$ of params; $E$ replicated experts; the MoE-layer map; placement; shared / fine-grained experts 3 The Router / Gating Network $g(\mathbf x)=\mathrm{softmax}(W_g\mathbf x)$; top-$k$; noisy top-$k$; softmax-before vs. -after; the only new learned mechanism 4 Sparse Top-k Computation & the Combine The weighted sum $\sum_{i\in\mathcal T_k}g_i E_i(\mathbf x)$; sparse compute, dense memory; gradient through a discrete choice 5 Load Balancing & Auxiliary Losses Expert collapse; importance/load loss; Switch loss $\alpha E\sum_i f_iP_i$; z-loss; aux-loss-free bias 6 Capacity & Token Dropping Capacity factor $C_f$; capacity $=C_fkT/E$; overflow/dropping; dispatch & combine tensors 7 Expert Parallelism Sharding experts; the two all-to-all collectives; communication-bound training; GShard 8 Training Dynamics & Scaling Router instability & z-loss; router precision; emergent specialization; params-at-fixed-FLOPs scaling 9 Fine-Tuning MoE Overfitting of sparse models; expert dropout; selective FT; router freezing; vs. dense SFT/RLHF 10 Inference in Depth Memory-vs-compute asymmetry; batch routing & load imbalance; unchanged KV cache; offloading & serving A Appendix A — MoE Family Shazeer, GShard, Switch, GLaM, Mixtral, DeepSeek-MoE along the axes that differ B Appendix B — Dense vs. MoE Every contrast in one decision-oriented table C Appendix C — Acronyms Every acronym and symbol, expanded Two invariants run through every lesson and are worth stating before anything else:

Read the full post →

From the Dense FFN to a Layer of Experts

Mon, 01 Jan 0001 00:00:00 +0000

Lesson 2 — From the Dense FFN to a Layer of Experts Technical course on MoE LLM mechanics — Lesson 2 of 10. Recap of Lesson 1 A Mixture-of-Experts model is a dense LLM with one component changed: the per-block FFN becomes $E$ expert FFNs plus a router that selects $k\ll E$ of them per token, combining their outputs as $\mathrm{MoE}(\mathbf x)=\sum_{i\in\mathcal T_k(\mathbf x)}g_i(\mathbf x)E_i(\mathbf x)$. This splits the single parameter count into total $P_{\text{tot}}$ (capacity, memory) and activated $P_{\text{act}}$ (per-token compute), with compute $\approx 6P_{\text{act}}D$ but memory $\propto P_{\text{tot}}$. We now build the layer itself, starting from the exact component it replaces and counting precisely how the parameter budget changes.

Read the full post →

The Router and Gating Network

Mon, 01 Jan 0001 00:00:00 +0000

Lesson 3 — The Router / Gating Network Technical course on MoE LLM mechanics — Lesson 3 of 10. Recap of Lesson 2 An MoE layer replicates the dense FFN into $E$ experts $E_i(\mathbf x)$ and combines a selected $k$ of them as $\mathrm{MoE}(\mathbf x)=\sum_{i\in\mathcal T_k(\mathbf x)}g_i(\mathbf x)E_i(\mathbf x)$, dropped into the same residual block. The experts hold essentially all the parameters, so $P_{\text{tot}}$ scales with $E$ while $P_{\text{act}}$ stays pinned to $k$. Two pieces were assumed but not built: the gate weights $g_i(\mathbf x)$ and the selected set $\mathcal T_k(\mathbf x)$. Both come from one small network — the router — which this lesson constructs.

Read the full post →

Sparse Top-k Computation and the Combine

Mon, 01 Jan 0001 00:00:00 +0000

Lesson 4 — Sparse Top-k Computation & the Combine Technical course on MoE LLM mechanics — Lesson 4 of 10. Recap of Lesson 3 The router is a learned matrix $W_g\in\mathbb R^{E\times d}$ producing logits $\mathbf h=W_g\mathbf x$, a softmax distribution $\mathbf P=\mathrm{softmax}(\mathbf h)$ over experts, and a top-$k$ selection $\mathcal T_k(\mathbf x)$ whose renormalized probabilities are the gate weights $g_i$. The selection is discrete ($\arg\text{top-}k$ has zero gradient), so the router can be trained only through the continuous gate weights of the experts it chose. This lesson assembles the complete sparse forward pass, then derives exactly what gradient flows back — to the experts and to the router — making “trained through the weights, not the choice” precise.

Read the full post →

Load Balancing and Auxiliary Losses

Mon, 01 Jan 0001 00:00:00 +0000

Lesson 5 — Load Balancing & Auxiliary Losses Technical course on MoE LLM mechanics — Lesson 5 of 10. Recap of Lesson 4 The MoE forward pass computes only the $k$ selected experts and combines them as $\mathbf y=\sum_{i\in\mathcal T_k}g_iE_i(\mathbf x)$ — sparse in compute, dense in memory. Backward, only selected experts and only selected router logits receive gradient; the router gets a relative reweighting of its winners and never a signal to choose differently. That partial, self-reinforcing coverage means a popular expert gets more training, gets better, and gets chosen more — a positive feedback loop toward collapse. This lesson names that failure, measures it, and adds the explicit counter-pressure that is the defining peculiarity of MoE training.

Read the full post →

Expert Capacity, Token Dropping and the Training Forward Pass

Mon, 01 Jan 0001 00:00:00 +0000

Lesson 6 — Expert Capacity, Token Dropping & the Training Forward Pass Technical course on MoE LLM mechanics — Lesson 6 of 10. Recap of Lesson 5 A sparse router collapses without help, so MoE training adds an explicit balancing pressure: the Switch auxiliary loss $\alpha E\sum_i f_iP_i$ (load $\times$ mass), the router z-loss for logit stability, or aux-loss-free bias control. These shape the router’s intent toward uniform usage. But intent is not enough: the dispatched tokens must physically fit into expert buffers on real hardware, and those buffers have a fixed size chosen in advance. This lesson introduces expert capacity, the token dropping that happens on overflow, and the batched dispatch/combine tensors that turn the per-token math of Lesson 4 into an efficient parallel operation.

Read the full post →

Distributed Training: Expert Parallelism

Mon, 01 Jan 0001 00:00:00 +0000

Lesson 7 — Distributed Training: Expert Parallelism Technical course on MoE LLM mechanics — Lesson 7 of 10. Recap of Lesson 6 To run experts efficiently on parallel hardware, each is given a fixed capacity $C=\lceil C_f kT/E\rceil$; overflow tokens are dropped onto the residual path, and the batch is executed via a dispatch tensor (scatter tokens into fixed-size expert buffers) and a gate-weighted combine tensor (un-scatter the outputs). That formulation assumed every expert lived on one device. They do not. This lesson distributes the experts across many accelerators, turning dispatch and combine into the two all-to-all network collectives that make MoE training distinctively communication-bound.

Read the full post →

Training Dynamics, Stability and Scaling Laws

Mon, 01 Jan 0001 00:00:00 +0000

Lesson 8 — Training Dynamics, Stability & Scaling Laws Technical course on MoE LLM mechanics — Lesson 8 of 10. Recap of Lesson 7 Experts are sharded across devices (expert parallelism), so each MoE layer is bracketed by two all-to-all collectives that route tokens to their experts and back — putting MoE training on a communication-bound critical path and giving load balance a systems meaning (the layer runs at the speed of its busiest expert). With experts, router, balancing (Lesson 5), capacity (Lesson 6), and communication (Lesson 7) all in place, the full training apparatus exists. This lesson studies its behavior: the instabilities it introduces and their fixes, the specialization that does (and doesn’t) emerge, and the scaling laws that justify spending all this complexity.

Read the full post →

Fine-Tuning MoE Models

Mon, 01 Jan 0001 00:00:00 +0000

Lesson 9 — Fine-Tuning MoE Models Technical course on MoE LLM mechanics — Lesson 9 of 10. Recap of Lesson 8 MoE training runs the same engine as dense training but with shakier dynamics — the router softmax can flip discrete selections, fixed by the z-loss and fp32 router precision — and its payoff is the scaling trade-off: raising the expert count lowers loss at fixed FLOPs, converting memory budget into quality. Specialization emerges weakly. With a pretrained sparse base model in hand, this lesson adapts it: supervised fine-tuning and alignment, where sparsity changes how adaptation behaves and introduces failure modes the dense recipe never sees.

Read the full post →

Inference in Depth

Mon, 01 Jan 0001 00:00:00 +0000

Lesson 10 — Inference in Depth Technical course on MoE LLM mechanics — Lesson 10 of 10. Recap of Lesson 9 Adapting an MoE uses the same SFT/RLHF/DPO machinery as a dense model, but sparsity makes it overfit (capacity is the large $P_{\text{tot}}$ on small data, with uneven expert gradient coverage), countered by expert dropout, selective/LoRA fine-tuning, and the uniquely-MoE option of freezing the router. The output is one frozen, aligned $\theta$. We now serve it. As in the companion course, nothing here updates a parameter; every knob is [CONFIG] and every buffer is [STATE] — but the Lesson 1 memory-vs-compute asymmetry now becomes the central, practical fact of inference.

Read the full post →

Appendix A: The MoE Model Family

Mon, 01 Jan 0001 00:00:00 +0000

Appendix A — A Field Guide to the MoE Family Technical course on MoE LLM mechanics — Appendix to Lessons 2, 5, and 8. Where this fits The lessons built one generic MoE and traced its consequences. Real models differ along a handful of axes that the lessons introduced piecemeal; this appendix collects the major models and lines them up on those axes, with a short note on what each one contributed. Read it after Lesson 8, or dip in whenever a model name appears. The axes that actually differ:

Read the full post →

Appendix B: Dense vs. MoE Side-by-Side

Mon, 01 Jan 0001 00:00:00 +0000

Appendix B — The Dense vs. MoE Cheat Sheet Technical course on MoE LLM mechanics — Appendix to all lessons. Where this fits This appendix collects every contrast drawn across the ten lessons into one place, organized by where in the stack the difference lives. Use it as a lookup table when you need the short answer, with a pointer back to the lesson that derives it. The single root cause to keep in mind: a dense model computes all of $\theta$ for every token; an MoE computes only a router-selected $k/E$ slice of its experts per token. Every row below is a consequence of that one difference (Lesson 1).

Read the full post →

Appendix C: Acronyms

Mon, 01 Jan 0001 00:00:00 +0000

Appendix C — Acronyms & Symbols Technical course on MoE LLM mechanics — reference appendix for all lessons. Where this fits A lookup table to keep open alongside the lessons. It covers the symbols and acronyms specific to MoE; for the shared transformer vocabulary (attention, $Q/K/V$, LayerNorm, AdamW, cross-entropy, perplexity, etc.) see the companion dense course’s acronym appendix. Terms below are grouped by what they describe. Core MoE symbols Symbol Meaning Lesson $E$ number of experts per MoE layer [CONFIG] 1, 2 $k$ number of experts activated per token (top-$k$) [CONFIG] 1, 3 $E_i(\mathbf x)$ expert $i$ — an FFN with its own weights [LEARNED] 2 $W_g\in\mathbb R^{E\times d}$ router (gating) matrix [LEARNED] 3 $\mathbf h=W_g\mathbf x$ router logits (one score per expert) [STATE] 3 $\mathbf P=\mathrm{softmax}(\mathbf h)$ router distribution over experts [STATE] 3 $\mathcal T_k(\mathbf x)$ set of top-$k$ selected expert indices for a token [STATE] 3, 4 $g_i(\mathbf x)$ gate weight: renormalized prob. of selected expert $i$ [STATE] 3, 4 $\mathbf x$ token representation entering the MoE sub-layer 2–4 Parameter accounting Symbol Meaning Lesson $P_{\text{tot}}$ total parameters — capacity, drives memory 1 $P_{\text{act}}$ activated parameters per token — drives compute 1 $k/E$ sparsity ratio (fraction of experts active) 1, 4 $d$ model (residual) width 2 $d_{ff}$ expert hidden width (typically $4d$) [CONFIG] 2 $m$ granularity factor (fine-grained experts) [CONFIG] 2 $\mathcal S$ set of shared (always-on) experts [CONFIG] 2 Balancing Symbol Meaning Lesson $f_i$ fraction of batch tokens dispatched to expert $i$ [STATE] 5 $P_i$ average router mass on expert $i$ over the batch [STATE] 5 $\mathcal L_{\text{aux}}=\alpha E\sum_i f_iP_i$ Switch auxiliary load-balancing loss 5 $\alpha$ auxiliary-loss weight [CONFIG] 5 $\mathcal L_z=\beta\,\overline{(\log\sum_i e^{h_i})^2}$ router z-loss (logit stability) 5, 8 $\beta$ z-loss weight [CONFIG] 5 $b_i$ per-expert selection bias (aux-loss-free balancing) 5 $\gamma$ bias update rate [CONFIG] 5 Capacity & communication Symbol Meaning Lesson $T$ number of tokens in a batch 6 $C_f$ capacity factor ($\ge 1$) [CONFIG] 6 $C=\lceil C_f kT/E\rceil$ expert capacity (fixed buffer size per expert) 6 $D$ (tensor) dispatch tensor — scatters tokens into expert buffers [STATE] 6 $\tilde D$ combine tensor — gate-weighted un-scatter [STATE] 6 $G$ expert-parallel degree (devices over which experts are sharded) [CONFIG] 7 $D$ (scalar) number of training tokens (in $C\approx 6P_{\text{act}}D$) 1, 8 Acronyms Acronym Expansion Lesson MoE Mixture of Experts all FFN Feed-Forward Network (the expert’s internal structure) 2 top-$k$ selecting the $k$ highest-scoring experts 3 all-to-all collective where every device exchanges data with every other 7 EP Expert Parallelism 7 DP / TP / PP Data / Tensor / Pipeline Parallelism (inherited from dense) 7 z-loss router log-partition penalty (numerical stability) 5, 8 GShard the framework that standardized expert parallelism 7, A ST-MoE “Stable and Transferable MoE” (z-loss, fine-tuning) 8, A GLaM “Generalist Language Model” (MoE decoder LM) 8, A LoRA Low-Rank Adaptation (parameter-efficient fine-tuning) 9 SFT / RLHF / DPO the alignment pipeline (inherited from dense) 9 KV cache cached keys/values for fast decode (inherited, unchanged) 10 GQA / MQA Grouped- / Multi-Query Attention (KV-cache savings, inherited) 10 $P_{\text{tot}}$ / $P_{\text{act}}$ total / activated parameters 1 The two invariants, restated Sparse activation: all of $\theta$ is stored and trained, but only a $k/E$ slice of the experts is computed per token — so memory follows $P_{\text{tot}}$, compute follows $P_{\text{act}}$. A free router collapses: sparse routing requires an explicit balancing pressure (auxiliary loss, z-loss, or bias control) that dense training never needs. ← Course index · Reference for all lessons

Read the full post →