2.9 Mixture-of-Experts (MoE) Architectures¶

There is a brutal fact at the center of the scaling story: in a dense Transformer, every parameter does work on every token. If you want a model that “knows more” — more facts, more languages, more code idioms — you grow the feed-forward networks (FFNs), and the cost of each forward pass grows in lockstep. A 70-billion-parameter dense model spends roughly 70 billion parameters’ worth of floating-point operations (FLOPs) on the word “the” and on a subtle line of CUDA, equally. That is wasteful, and it is the wall that Mixture-of-Experts (MoE) is built to break.

The MoE idea is almost embarrassingly simple to state: decouple the number of parameters from the FLOPs per token. Replace the single FFN in a Transformer block with a bank of FFNs called experts, and add a small learned router that, for each token, picks just a few of them. A token is processed by, say, 2 experts out of 256. The model now contains the knowledge of 256 FFNs but pays for only 2 of them per token. You get the memory footprint and representational capacity of a giant model with the compute bill of a small one. This is conditional computation — the network decides, per input, which sub-networks to activate — and it is the architectural choice behind GShard, Switch Transformer, Mixtral, DeepSeek-MoE, Qwen-MoE, and most of the frontier “trillion-parameter” models that nonetheless run at the speed of a much smaller dense model.

The catch — and this chapter is largely about the catch — is that picking is hard. The router is a discrete, non-differentiable decision sitting in the middle of a network we train with smooth gradients. Tokens must be routed in balanced batches so that GPUs stay busy and no expert is starved or overwhelmed. Experts must be spread across devices, which turns every MoE layer into a communication problem (the expert parallelism we preview at the end). Get the load balancing wrong and the model collapses to using one expert, throwing away the whole point. We will build the mechanism from the FFN up, write a correct toy MoE layer you can run, work the numbers on capacity and parameter counts, and tour the landmark designs that got each piece right.

This chapter assumes you know the Transformer block — especially that the FFN is two linear layers with a nonlinearity — and the softmax from Probability, Statistics & Information Theory. It pairs closely with Scaling Laws (why sparsity changes the compute-optimal frontier) and with Distributed Training II (how experts are sharded in practice).

From One FFN to a Bank of Experts¶

Recall the Transformer block: attention mixes information across tokens, then a position-wise FFN transforms each token independently. For a token representation \(x \in \mathbb{R}^{d_\text{model}}\) the standard FFN is

with \(\phi\) a nonlinearity (GELU, or a gated SwiGLU as in modern models — see The Transformer Block). Typically \(d_\text{ff} = 4\,d_\text{model}\), so the FFN holds about \(8\,d_\text{model}^2\) parameters and is usually two-thirds of the parameters in the whole model. The FFN is also where most of the “knowledge” is thought to live: keys-and-values memory, factual associations, n-gram-like circuits. So if you want to add capacity cheaply, the FFN is exactly where you want to do it.

A MoE layer replaces that one FFN with \(E\) experts, each an independent FFN with its own weights \(\{\operatorname{FFN}_e\}_{e=1}^{E}\), plus a router (also called the gate) \(g: \mathbb{R}^{d_\text{model}} \to \mathbb{R}^{E}\) that scores how relevant each expert is to the token. The layer output is a weighted combination of the experts the router selects:

where \(\mathcal{T}(x) \subseteq \{1,\dots,E\}\) is the (small) set of selected experts for token \(x\) and \(g_e(x)\) is its gate weight. In a dense MoE (the original 1991 Jacobs–Jordan formulation, and a “soft” baseline today) \(\mathcal{T}(x) = \{1,\dots,E\}\): you run all experts and blend them. That is differentiable and clean but defeats the purpose — you pay for all \(E\) FFNs. The breakthrough that made MoE useful at LLM scale is sparsity: make \(|\mathcal{T}(x)| = k\) with \(k \ll E\) (typically \(k=1\) or \(k=2\)), so each token touches only a handful of experts.

Two vocabulary items you will see constantly. Total parameters counts every expert’s weights — this is the model’s “size on disk” and its capacity. Active parameters counts only what a single token actually uses — this drives the FLOPs and therefore latency. Mixtral 8×7B, for instance, has on the order of 47B total parameters but activates roughly 13B per token (it is not 8×7=56B, because attention and embeddings are shared and only the FFNs are replicated). That ratio — total-to-active — is the sparsity dial, and it is the headline of every MoE model card. The whole bet of MoE is that a model’s quality tracks total parameters while its cost tracks active parameters, so you want that ratio large.

The Router: Gating, Top-k, and the Softmax-Then-Select Question¶

The router is a single linear layer — astonishingly small relative to what it controls. It maps each token to one logit per expert:

For \(E=256\) experts and \(d_\text{model}=4096\), that is about a million parameters governing the routing of a model with hundreds of billions. The router’s job is to turn the logit vector \(h(x)\) into (a) a choice of which \(k\) experts to run and (b) a set of weights with which to combine their outputs. There are two design knobs, and the order in which you pull them matters.

Top-k selection. Take the indices of the \(k\) largest logits: \(\mathcal{T}(x) = \operatorname{top\text{-}k}\big(h(x)\big)\). This is the discrete, non-differentiable heart of the layer. Crucially, gradients flow only through the experts that were selected — an expert that loses the top-k comparison gets no gradient signal this step. (This is also why the gate weights, not the indices, must carry gradient: the selected experts’ weights \(g_e\) scale their outputs, so \(\partial \mathcal{L}/\partial W_g\) flows through them. The top-k operator itself is treated as a stop-gradient on the index decision; the network learns routing only through the magnitudes of the kept weights and the load-balancing loss we add later.)

Softmax: before or after top-k? Two conventions exist and they are not equivalent:

• Softmax-then-top-k (GShard, Switch): compute \(p = \operatorname{softmax}(h(x))\) over all \(E\) experts, then keep the \(k\) largest probabilities as the gate weights. The kept weights are a slice of a full distribution; they do not sum to 1.
• Top-k-then-softmax (Mixtral): select the top-\(k\) logits first, then softmax only over those \(k\). The gate weights now sum to 1, which keeps the output magnitude stable regardless of how confident the router was.

Mixtral’s choice — renormalizing over the chosen experts — is the more common modern default because it decouples “how much total signal flows” from “how spread the router’s confidence was.” With softmax-then-top-k, a router that is uncertain (flat distribution) produces small gate weights and a weak FFN contribution; renormalizing fixes that. Here is the gate, both ways, so you can see the difference is a single line:

import torch
import torch.nn.functional as F

def route_softmax_then_topk(logits, k):
    # GShard/Switch style: softmax over ALL experts, then keep the top-k slice.
    # The kept weights do NOT sum to 1 (they are a slice of a full distribution).
    probs = F.softmax(logits, dim=-1)                  # (tokens, E)
    weights, idx = torch.topk(probs, k, dim=-1)        # (tokens, k)
    return weights, idx

def route_topk_then_softmax(logits, k):
    # Mixtral style: pick top-k logits, then softmax over JUST those k.
    # The kept weights DO sum to 1 -> stable output magnitude.
    topk_logits, idx = torch.topk(logits, k, dim=-1)   # (tokens, k)
    weights = F.softmax(topk_logits, dim=-1)           # (tokens, k), rows sum to 1
    return weights, idx

logits = torch.tensor([[2.0, 1.0, 0.1, -1.0, 0.5, 3.0, -2.0, 0.2]])  # 1 token, 8 experts
w1, i1 = route_softmax_then_topk(logits, k=2)
w2, i2 = route_topk_then_softmax(logits, k=2)
print("softmax-then-topk:", i1.tolist(), w1.round(decimals=3).tolist())  # weights < 1, no sum-to-1
print("topk-then-softmax:", i2.tolist(), w2.round(decimals=3).tolist())  # weights sum to 1

Noisy / top-k gating. The original sparse-MoE paper (Shazeer et al., 2017) added tunable Gaussian noise to the logits before top-k:

The noise does two jobs: it breaks ties so different tokens explore different experts early in training, and it acts as a regularizer that spreads load. Most modern models drop the learned noise (they rely on the auxiliary balancing loss instead) but the lesson — that routing needs exploration to avoid premature collapse onto a few experts — recurs throughout the field. DeepSeek-V3 reintroduced a different exploration mechanism: a per-expert bias added to the routing scores used for selection only (not for the gate weights), nudged up or down each step to equalize load without polluting the gradient — an “auxiliary-loss-free” balancing trick we return to below.

A baseline that should worry you. Hash layers (Roller et al., 2021) throw the learned router away entirely and dispatch each token by a fixed random hash of its token id — no gating, no aux loss, perfectly balanced by construction — and still capture a large share of MoE’s gain over a dense baseline. The lesson is not that routing is pointless but that much of the benefit comes from the extra capacity plus balanced dispatch, so a learned router has to earn its keep. Whenever you build one, run the hash baseline next to it; if your router does not beat a hash, it is not learning anything.

A Toy MoE Layer From Scratch¶

Let us build a correct, runnable sparse MoE FFN layer. The pedagogical subtlety in every MoE implementation is the same: top-k routing is ragged — different tokens go to different experts, so you cannot just run a clean batched matmul over the whole sequence. There are two ways to handle this, and you should know both.

1. The dense/masked loop (simple, what we write first): loop over experts; for each expert, find the tokens routed to it, run that expert on them, scatter the results back. Clear and correct; not the fastest because of the gather/scatter, but it is exactly what grouped-GEMM kernels (Megablocks, the Triton grouped matmul) optimize under the hood.
2. The all-experts-then-mask trick (only for tiny \(E\)): run every expert on every token and zero out the ones you did not select. Trivial to write, but it throws away the FLOPs savings — useful only for unit tests and teaching.

We implement the loop version, which is what scales:

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class Expert(nn.Module):
    """A single expert: a standard 2-layer FFN (SwiGLU-free for clarity)."""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):                      # x: (n_tokens_for_this_expert, d_model)
        return self.w2(F.gelu(self.w1(x)))

class SparseMoE(nn.Module):
    """Top-k sparse MoE layer with capacity-based token dropping and an
    auxiliary load-balancing loss (Switch Transformer style)."""
    def __init__(self, d_model, d_ff, n_experts, k=2, capacity_factor=1.25):
        super().__init__()
        self.d_model = d_model
        self.n_experts = n_experts
        self.k = k
        self.capacity_factor = capacity_factor
        self.router = nn.Linear(d_model, n_experts, bias=False)   # the gate W_g
        self.experts = nn.ModuleList(
            [Expert(d_model, d_ff) for _ in range(n_experts)]
        )

    def forward(self, x):
        # x: (batch, seq, d_model). Flatten to a flat token list -- routing is per-token.
        B, T, D = x.shape
        x = x.reshape(-1, D)                       # (N, D), N = B*T
        N = x.shape[0]

        # 1) ROUTE: logits -> top-k experts -> renormalized gate weights (Mixtral style).
        logits = self.router(x)                    # (N, E)
        topk_logits, topk_idx = torch.topk(logits, self.k, dim=-1)   # (N, k)
        gate = F.softmax(topk_logits, dim=-1)      # (N, k), each row sums to 1

        # 2) CAPACITY: max tokens any single expert will accept this batch.
        #    capacity = capacity_factor * (tokens * k / E), rounded up.
        capacity = math.ceil(self.capacity_factor * (N * self.k) / self.n_experts)
        capacity = max(capacity, 1)

        out = torch.zeros_like(x)                  # accumulate expert outputs here
        # Counters for the auxiliary loss and for diagnostics.
        # f_e = fraction of tokens dispatched to expert e (hard count, top-1 view).
        # P_e = mean router probability mass on expert e (soft).
        router_prob = F.softmax(logits, dim=-1)    # (N, E) full distribution for aux loss
        load = torch.zeros(self.n_experts, device=x.device)

        # 3) DISPATCH: loop over experts; process the tokens routed to each.
        for e in range(self.n_experts):
            # Which (token, slot) pairs picked expert e? slot in {0..k-1}.
            sel = (topk_idx == e)                  # (N, k) bool
            token_ids, slot_ids = sel.nonzero(as_tuple=True)   # variable length
            load[e] = token_ids.numel()            # tokens wanting expert e (pre-drop)

            if token_ids.numel() == 0:
                continue
            # CAPACITY DROP: keep only the first `capacity` tokens for this expert.
            if token_ids.numel() > capacity:
                token_ids = token_ids[:capacity]
                slot_ids = slot_ids[:capacity]

            ex_in = x[token_ids]                    # (n_e, D) gather
            ex_out = self.experts[e](ex_in)         # (n_e, D) run THIS expert only
            w = gate[token_ids, slot_ids].unsqueeze(-1)        # (n_e, 1) gate weight
            out.index_add_(0, token_ids, w * ex_out)           # scatter-add back

        # 4) AUXILIARY LOAD-BALANCING LOSS (Switch): encourage uniform routing.
        f = load / load.sum().clamp(min=1)         # dispatched fraction per expert
        P = router_prob.mean(dim=0)                # mean router prob per expert
        aux_loss = self.n_experts * torch.sum(f * P)   # minimized when both are uniform

        return out.reshape(B, T, D), aux_loss

# --- smoke test: it runs and shapes are right ---
torch.manual_seed(0)
moe = SparseMoE(d_model=32, d_ff=64, n_experts=8, k=2, capacity_factor=1.25)
x = torch.randn(4, 16, 32)              # batch 4, seq 16
y, aux = moe(x)
print(y.shape, float(aux))             # torch.Size([4, 16, 32]) <some positive scalar>

A few details in that code are load-bearing and worth dwelling on. The index_add_ is what makes top-\(k>1\) correct: a token selected by two experts contributes both weighted outputs to the same output row, and because the gate weights were renormalized to sum to 1, the result is a true convex combination. The capacity clamp token_ids[:capacity] is the token dropping mechanism — tokens beyond an expert’s quota silently get no FFN contribution this layer (their output stays whatever the residual stream carried in, since we add to a zero buffer that then joins the residual). In a real implementation the dropped token still flows through the residual connection, so it is degraded, not destroyed. And the auxiliary loss combines a hard count f with a soft probability P on purpose — we explain exactly why next.

Load Balancing: The Make-or-Break Problem¶

Left to its own devices, a sparse router collapses. Early in training, one expert is — by random initialization — slightly better. The router sends it more tokens, so it gets more gradient, so it improves faster, so it attracts even more tokens. This is a textbook rich-get-richer feedback loop, and its fixed point is a model that routes (almost) everything to one expert while the rest atrophy, never receiving gradient. You have then paid for \(E\) experts and trained one. Worse, in a distributed setting, the favored expert’s GPU is saturated while the others idle — the slowest device sets the pace, so imbalance directly wastes wall-clock time. Load balancing is therefore not a nicety; it is the difference between MoE working and MoE failing.

The auxiliary load-balancing loss¶

The dominant fix, from GShard and refined by Switch, is to add a differentiable penalty that pushes routing toward uniform. Define, over a batch of \(N\) tokens and \(E\) experts:

• \(f_e\) = the fraction of tokens dispatched to expert \(e\) (a hard count, \(\sum_e f_e = 1\) for top-1; for top-\(k\) it sums to \(k\) before normalization).
• \(P_e = \frac{1}{N}\sum_{i} p_e(x_i)\) = the mean router probability assigned to expert \(e\) (soft, from the full softmax).

The Switch auxiliary loss is

with a small coefficient \(\alpha\) (on the order of \(10^{-2}\)). Two questions: why the product \(f_e P_e\), and why mix a hard count with a soft probability?

The product is minimized at uniform load, and it is worth being precise about in what sense. At a routing fixed point the hard counts track the soft mass, \(f_e \approx P_e\), so \(\sum_e f_e P_e \approx \sum_e P_e^2 = \lVert P \rVert_2^2\) — and minimizing a squared norm subject to \(\sum_e P_e = 1\) gives \(P_e = 1/E\) for every \(e\), with value \(1/E\). The factor \(E\) out front rescales that to exactly \(1\) at perfect balance, so \(\alpha\) means the same thing across different \(E\); a perfectly collapsed router gives \(E \cdot 1 = E\) instead. Read literally at a single step, with \(f\) detached, the gradient on each expert’s probability mass is just \(\alpha E f_e\): constant pressure proportional to how overloaded that expert currently is. It is the dynamics of that pressure, not a static optimum, that walks the router toward uniform.

The hard/soft mix is the clever part. We want to penalize the hard dispatch counts \(f_e\) — those are what actually cause imbalance — but \(f_e\) comes from a top-k/argmax and is non-differentiable. So we keep \(f_e\) as a constant (it provides the direction: which experts are overloaded) and multiply by the differentiable \(P_e\) (which provides the gradient: nudging the router’s probabilities down on overloaded experts and up on starved ones). The loss “reads” imbalance from the hard counts and “writes” the correction through the soft probabilities. It is a beautiful little gradient-routing trick and it is why MoE trains at all.

def switch_aux_loss(router_logits, expert_idx, n_experts, alpha=1e-2):
    """router_logits: (N, E) pre-softmax. expert_idx: (N,) top-1 choice per token."""
    P = F.softmax(router_logits, dim=-1).mean(dim=0)        # (E,) soft mass, differentiable
    # f: hard dispatch fraction (a count; we detach so only P carries gradient).
    f = torch.bincount(expert_idx, minlength=n_experts).float()
    f = (f / f.sum()).detach()                              # (E,), constant direction
    return alpha * n_experts * torch.sum(f * P)

A second, smaller term often accompanies it: the router z-loss (from ST-MoE), \(\mathcal{L}_z = \beta \cdot \frac{1}{N}\sum_i \big(\log\sum_e e^{h_e(x_i)}\big)^2\), which penalizes large router logits. This keeps the router from saturating into a near-one-hot regime where small numerical perturbations flip the top-k choice, dramatically improving training stability in bf16/fp8 (see Mixed Precision).

Expert-choice routing: invert the problem¶

There is an elegant alternative that makes imbalance impossible by construction. In standard token-choice routing, each token picks its top-\(k\) experts — and nothing stops every token from picking the same expert. Expert-choice routing (Zhou et al., 2022) flips it: each expert picks its top-\(C\) tokens, where \(C\) is exactly its capacity. Now every expert processes precisely \(C\) tokens — perfect balance, no auxiliary loss, no dropping in the usual sense. The trade-off: a token might be chosen by many experts or by none, so some tokens get more compute than others (the model decides which tokens are “hard”), and the scheme is awkward for autoregressive inference because it needs to see the whole batch of tokens to do the selection, which breaks causal, one-token-at-a-time decoding. Expert-choice is therefore most natural in encoder-style or training-time settings; token-choice with an aux loss remains the default for decoder LLMs.

Capacity Factor, Token Dropping, and the Worked Numbers¶

The capacity of an expert is the maximum number of tokens it will accept in a batch. With \(N\) tokens, top-\(k\) routing, and \(E\) experts, the average tokens per expert is \(Nk/E\). We multiply by a capacity factor \(C_f \ge 1\) to leave slack for imbalance:

Why does capacity exist at all? Because GPUs want static, rectangular tensors. To run expert \(e\) as one efficient matmul, you need a fixed-size buffer of tokens for it. Capacity is that buffer size. If more than capacity tokens route to an expert, the overflow is dropped (skips the FFN, passes through on the residual). If fewer arrive, the buffer is padded with zeros (wasted FLOPs). So \(C_f\) trades two evils: too low and you drop many tokens (hurting quality); too high and you waste compute and memory on padding. Typical training values are \(C_f \in [1.0, 2.0]\); Switch used around 1.0–1.25, GShard often 2.0. Modern systems increasingly use dropless MoE (Megablocks) with grouped/block-sparse GEMMs that handle ragged sizes directly, eliminating both dropping and padding — but the capacity concept is essential for understanding the classics and most serving paths.

The takeaway in one line: capacity factor governs the dropping-vs-padding trade-off; the total-to-active ratio governs the capacity-vs-cost trade-off. Both are knobs you tune, and both show up directly in training loss and serving latency.

The Landmark Designs: GShard, Switch, Mixtral, DeepSeek-MoE¶

The ideas above were not discovered all at once; each landmark model fixed a specific pain point. Knowing the lineage lets you reason about why a given model made its choices.

GShard (2020): MoE goes to scale¶

GShard (Lepikhin et al., Google) put MoE into a 600B-parameter multilingual translation Transformer and established the recipe that everyone still uses: top-2 routing, an auxiliary load-balancing loss, a capacity factor with token dropping, and — critically — expert parallelism via an all-to-all communication primitive that shuffles tokens to wherever their experts live. GShard is where MoE stopped being a curiosity and became a scaling tool, and where the systems problem (the all-to-all, previewed below) was first confronted head-on.

Switch Transformer (2021): top-1 is enough¶

Switch (Fedus, Zoph, Shazeer, Google) made a deliberately aggressive simplification: route each token to exactly one expert (\(k = 1\)). The conventional wisdom had been that you needed top-2 so the router could compare two experts and get a gradient signal for the choice; Switch showed top-1 trains fine with the right tricks. The payoff is roughly halved communication and compute per token. Switch also contributed the cleaner auxiliary loss formulation above, careful bf16 stability work (selectively casting the router to fp32), and the router z-loss. It scaled to 1.6T parameters. If you implement MoE once from scratch, implement Switch; it is the minimal complete design.

Mixtral 8×7B (2023): MoE for open decoder LLMs¶

Mixtral (Mistral AI) was the model that made MoE mainstream for the open-source decoder-only LLM. Eight experts per layer, top-2 routing, top-k-then-softmax gating (renormalized weights), with attention and embeddings shared across experts so only the FFNs are replicated. The naming “8×7B” is a friendly fiction: it is not 56B parameters but \(\approx 47\)B total (because the non-FFN parts are not multiplied), activating \(\approx 13\)B per token. Mixtral demonstrated that an MoE could match or beat a much larger dense model (Llama-2 70B) at a fraction of the active compute, on standard chat/reasoning benchmarks — the existence proof that sold the architecture to practitioners.

DeepSeek-MoE and the fine-grained + shared-expert idea (2024)¶

DeepSeek-MoE introduced two refinements that the frontier has largely adopted, and they are worth understanding precisely because they attack a real limitation of the GShard/Switch design: coarse experts force redundancy. If you only have 8 big experts and you want to combine, say, “code knowledge” with “math knowledge” on a given token, your only options are coarse. Two changes fix this:

1. Fine-grained experts. Instead of \(E\) big experts, use \(mE\) experts each \(1/m\) the size, and increase \(k\) proportionally. The active parameter count is unchanged, but the number of combinations of experts a token can express explodes combinatorially (\(\binom{mE}{mk}\) grows fast). More, smaller experts give the router a richer, more composable palette and empirically improve specialization. DeepSeek-V3 takes this far: 256 routed experts, top-8 per token.

2. Shared experts. Designate a small number of experts (often 1) that every token always uses, in addition to its routed experts. The intuition: common, universal knowledge (“how to form a plural,” “basic syntax”) should not have to be redundantly relearned inside many specialized experts. The shared expert absorbs the common case; the routed experts then specialize on the residual. This both improves quality and reduces the redundancy that pure routed-MoE suffers.

DeepSeek-V3 further replaced the auxiliary loss with the auxiliary-loss-free bias-adjustment scheme mentioned earlier: a per-expert bias \(b_e\) is added to the routing logits for selection only, and after each step \(b_e\) is decremented for overloaded experts and incremented for underloaded ones. This achieves balance without the aux loss’s gradient interference, which can slightly degrade quality by fighting the language-modeling objective. The architecture also pairs MoE with Multi-head Latent Attention (MLA) for KV-cache compression — see Multi-Head Attention, MQA, GQA & MLA.

A small comparison table to anchor the lineage:

MoE in Practice: The Libraries That Implement This¶

You will almost never ship the dispatch loop we wrote. Each layer of the MoE stack has a mature open-source owner, and knowing which library owns which concern is most of the practical knowledge.

Modelling and fine-tuning — HuggingFace transformers. Mixtral, Qwen3-MoE, DeepSeek-V3, OLMoE and Llama 4 all ship as transformers model classes whose MoE block (MixtralSparseMoeBlock and its siblings) is structurally the loop above: a nn.Linear gate, an expert ModuleList, top-\(k\), and an index_add_ scatter. The non-obvious API detail is that the auxiliary loss is opt-in — you must ask the model to return router logits before the balancing term is folded into the loss:

from transformers import AutoConfig, AutoModelForCausalLM

cfg = AutoConfig.from_pretrained("mistralai/Mixtral-8x7B-v0.1")
print(cfg.num_local_experts, cfg.num_experts_per_tok)  # -> 8, 2  (E and k)
print(cfg.router_aux_loss_coef)                        # the alpha in alpha * E * sum(f*P)

cfg.output_router_logits = True        # <-- REQUIRED: without it, no aux loss is added
model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-v0.1", config=cfg)

out = model(input_ids=input_ids, labels=labels)
# out.loss        = LM cross-entropy + router_aux_loss_coef * load-balancing loss
# out.router_logits = tuple of (N, E) tensors, one per MoE layer -- feed these
#                     straight into the ExpertLoadMonitor at the end of this chapter.

The reference implementation of the balancing term lives in transformers.models.mixtral.modeling_mixtral.load_balancing_loss_func; read it next to the switch_aux_loss we derived — it is the same \(\alpha E \sum_e f_e P_e\), generalized to top-\(k\) and to attention masks. Forget output_router_logits=True when fine-tuning an MoE and you silently train with no balancing at all, which is one of the most common ways to wreck a routed model on a small dataset.

Pretraining at scale — Megatron-Core, DeepSpeed-MoE, MegaBlocks, Tutel. These own dispatch kernels and expert parallelism. Megatron-Core exposes the entire design space of this chapter as flags, which makes its CLI a good checklist of the knobs that matter:

#!/bin/bash
# Megatron-Core: a DeepSeek-style fine-grained MoE with a shared expert.
# Every flag below is a knob this chapter derived from first principles.

MOE_ARGS=(
  --num-experts 64                        # E: routed experts per MoE layer
  --moe-router-topk 6                     # k: experts per token
  --moe-router-load-balancing-type aux_loss   # 'none' => the bias-based scheme
  --moe-aux-loss-coeff 1e-2               # alpha in alpha * E * sum(f * P)
  --moe-z-loss-coeff 1e-3                 # router z-loss (ST-MoE), logit hygiene
  --moe-expert-capacity-factor 1.25       # omit this flag entirely => dropless
  --moe-shared-expert-intermediate-size 2048  # always-on shared expert (DeepSeekMoE)
  --moe-grouped-gemm                      # one batched GEMM over the ragged groups
  --moe-token-dispatcher-type alltoall    # the expert-parallel shuffle
)

torchrun --nproc-per-node 8 pretrain_gpt.py \
  "${MOE_ARGS[@]}" \
  --expert-model-parallel-size 8 \
  --tensor-model-parallel-size 1
  # ... plus the usual data / optimizer / schedule flags

DeepSpeed-MoE offers the same capability as a drop-in module (deepspeed.moe.layer.MoE) plus expert-slicing for inference; Microsoft Tutel contributes adaptive, dynamically-switched all-to-all kernels; and MegaBlocks (megablocks.layers.dmoe.dMoE) is the dropless block-sparse implementation — it reformulates the ragged dispatch as one block-sparse matmul so there is no capacity factor, no padding, and no dropped tokens. If you want to write the kernel rather than call it, the grouped-GEMM pattern is the canonical intermediate exercise in Writing GPU Kernels with Triton.

Serving — vLLM and SGLang. Both implement fused MoE kernels (a single grouped GEMM over the sorted, expert-major token list) and both let you choose whether the MoE layer is sharded by tensor parallelism — every rank holds a slice of every expert — or by expert parallelism, where each rank owns whole experts and the all-to-all does the shuffling. vLLM turns the latter on with --enable-expert-parallel; SGLang exposes an equivalent expert-parallel size flag. Which is faster depends on batch size and interconnect, and is the subject of Serving Mixture-of-Experts.

Sparse upcycling: an MoE without pretraining one¶

You do not have to train an MoE from random initialization. Sparse upcycling (Komatsuzaki et al., 2023) takes a trained dense checkpoint, replaces each FFN with \(E\) copies of that same FFN, adds a freshly-initialized router, and continues pretraining. At step zero the MoE layer computes almost exactly what the dense layer did — no capability is thrown away — and the experts then diverge under continued training. Several open MoE releases (the Qwen-MoE line among them) were produced this way, because it converts an existing dense run into extra capacity for a fraction of a from-scratch MoE budget. Two initialization details decide whether it works:

def upcycle_ffn_to_moe(dense_ffn, d_model, d_ff, n_experts, k=2, jitter=1e-2):
    """Sparse upcycling: one trained dense FFN -> an E-expert MoE layer.
    Every expert starts as a COPY of the dense FFN, so the upcycled layer starts
    from the dense model's function rather than from noise."""
    moe = SparseMoE(d_model, d_ff, n_experts, k=k)
    sd = dense_ffn.state_dict()
    with torch.no_grad():
        for e in range(n_experts):
            moe.experts[e].load_state_dict(sd)          # exact clone of the dense FFN
            # (1) Break the symmetry: identical experts get identical gradients and
            #     would never differentiate. A little noise lets them specialize.
            for p in moe.experts[e].parameters():
                p.add_(jitter * p.std() * torch.randn_like(p))
        # (2) Small -- NOT zero -- router weights. Near-uniform routing preserves the
        #     dense function; nonzero weights break per-token ties so every expert is
        #     selected sometimes and therefore receives gradient. A zero router would
        #     always pick experts 0..k-1 and starve the rest permanently.
        moe.router.weight.normal_(mean=0.0, std=1e-3)
    return moe

torch.manual_seed(0)
dense = Expert(d_model=32, d_ff=64)
upcycled = upcycle_ffn_to_moe(dense, d_model=32, d_ff=64, n_experts=8, k=2)
x = torch.randn(2, 16, 32)
y_moe, _ = upcycled(x)
y_dense = dense(x.reshape(-1, 32)).reshape(2, 16, 32)
print(float((y_moe - y_dense).abs().mean()))   # ~3e-3: starts where dense left off

The gate weights sum to 1 and every expert is (nearly) the same function, so the top-\(k\) convex combination reproduces the dense FFN — which is exactly why upcycling does not lose the parent model’s quality on day one.

Expert Parallelism: A Systems Preview¶

So far we have treated all experts as living on one device. At frontier scale they cannot — 256 experts of 100M+ parameters each will not fit on one GPU, and even if they did, you would want their compute spread out. Expert parallelism (EP) shards the experts across devices: GPU 0 holds experts 0–31, GPU 1 holds 32–63, and so on. This creates the defining systems challenge of MoE, and it is worth previewing here even though Distributed Training II covers it in depth.

The problem: a token on GPU 0 may route to an expert on GPU 5. Before the MoE layer can run, every token must be sent to the device holding its chosen expert, and after the experts run, every result must be sent back to the token’s home device for the residual add and the next layer. These two shuffles are all-to-all collectives (see Parallel Computing & Collective Communication) — the most communication-intensive pattern there is, because in the worst case every device talks to every other device.

This is exactly why load balancing maps directly to throughput. If expert 3 (on GPU 5) is overloaded, GPU 5 does more work than the others; because the all-to-all is a barrier, every GPU waits for the slowest. Imbalance does not just hurt quality — it idles your cluster. The all-to-all is also bandwidth-hungry, so production MoE training overlaps it with computation, places experts to minimize cross-node traffic, and tunes the capacity factor to bound the message sizes. DeepSeek-V3’s auxiliary-loss-free balancing and its node-limited routing (restricting each token to experts on a bounded number of nodes) are, at bottom, communication optimizations dressed as routing rules. The architecture and the systems are inseparable in MoE; you cannot reason about one without the other.

For inference, EP interacts with the serving stack: MoE models have low active-parameter FLOPs (great for decode) but high total-parameter memory (you must keep all experts resident in HBM, since any token might need any of them) and high all-to-all traffic. This makes MoE serving a memory-bandwidth and communication problem more than a compute one — the opposite balance from dense models — and motivates the specialized kernels and expert-placement strategies of Serving Mixture-of-Experts: Expert Parallelism & All-to-All Inference.

Training, Stability, and When to Reach for MoE¶

A few practitioner realities round out the picture. MoE models are less FLOP-efficient to train per step but reach a given loss in fewer FLOPs — the scaling-laws story is that, at a fixed compute budget, sparse models sit on a better loss-vs-FLOPs frontier than dense ones, which is the whole economic argument for them. But they come with sharp edges:

• Training instability. Routers are prone to spiky losses; the router z-loss, fp32 router computation, and careful initialization (small router weights so early routing is near-uniform) are standard mitigations. See Training Stability & Loss Spikes.
• Fine-tuning is finicky. Sparse models can overfit downstream tasks faster and route inconsistently on small datasets; freezing the router or using a higher capacity factor during fine-tuning helps. Distilling a sparse model into a dense one (Distillation & Compression) is a common deployment path when you want MoE’s training efficiency but dense serving simplicity.
• Memory dominates serving. As above, all experts must be resident, so MoE shifts the bottleneck from compute to HBM capacity and bandwidth. Expert offloading (CPU↔GPU paging of cold experts) and quantizing experts more aggressively than attention are active areas.
• When to use it. Reach for MoE when you are compute-bound in training or latency-bound in serving but have memory to spare — exactly the frontier-pretraining and high-throughput-serving regimes. If memory is your scarce resource (edge, single-GPU), a dense model of the same total size is simpler and serves with less machinery. This is precisely why the book’s capstone model, Stack-100M, is deliberately dense: at 100M parameters on one GPU, memory — not FLOPs — is the binding constraint, so MoE’s machinery buys nothing. The Retrospective & Scale-Up Path works out where MoE becomes the right fork on the way past 1B and gives a fine-grained + shared-expert FFN you can drop into that codebase.

Instrumenting the tip: an ExpertLoadMonitor¶

The tip above is only useful if you actually instrument it. Below is a drop-in ExpertLoadMonitor that mirrors the TrainingMonitor pattern from Training Stability & Loss Spikes — same log_every throttling, same log_step returning a metrics dict — but computes the three signals that predict MoE collapse before the loss does: normalized router entropy, the per-expert dispatch-fraction histogram, and a starvation alert.

import math
import torch

class ExpertLoadMonitor:
    """
    Tracks per-expert dispatch load and router entropy across MoE layers.
    Wire it by having each SparseMoE layer stash its pre-softmax router
    logits (shape (N, E), N = tokens in the batch, E = n_experts) into a
    dict keyed by layer name, then call log_step(logits_by_layer) once
    per training step.
    """
    def __init__(self, n_experts: int, k: int, log_every: int = 10,
                 starve_frac: float = 0.2, patience: int = 3):
        self.n_experts = n_experts
        self.k = k
        self.log_every = log_every
        self.threshold = starve_frac / n_experts  # e.g. 0.2/8 = 0.025
        self.patience = patience
        self.step = 0
        self.starve_streak = {}  # (layer_name, expert_idx) -> consecutive windows below threshold

    @torch.no_grad()
    def log_step(self, logits_by_layer: dict) -> dict:
        self.step += 1
        if self.step % self.log_every != 0:
            return {}

        metrics = {}
        alerts = []
        for name, logits in logits_by_layer.items():
            N, E = logits.shape
            probs = torch.softmax(logits, dim=-1)                       # (N, E)
            tok_H = -(probs * probs.clamp_min(1e-9).log()).sum(dim=-1)  # (N,) nats
            norm_entropy = (tok_H.mean() / math.log(E)).item()          # in [0, 1]

            idx = torch.topk(logits, self.k, dim=-1).indices            # (N, k)
            counts = torch.bincount(idx.reshape(-1), minlength=E).float()
            total = counts.sum()
            if total == 0:
                continue  # empty batch edge case
            f = counts / total                                          # (E,), sums to 1
            fair = 1.0 / E
            starved = (f < self.threshold).nonzero(as_tuple=True)[0].tolist()

            for e in range(E):
                key = (name, e)
                if e in starved:
                    self.starve_streak[key] = self.starve_streak.get(key, 0) + 1
                    if self.starve_streak[key] >= self.patience:
                        alerts.append(
                            f"{name}: expert {e} starved for {self.starve_streak[key]} "
                            f"windows (f={f[e]:.4f} < {self.threshold:.4f})"
                        )
                else:
                    self.starve_streak[key] = 0

            metrics[f"moe/{name}/router_entropy"] = norm_entropy
            metrics[f"moe/{name}/max_load_ratio"] = (f.max() / fair).item()  # >~2-3x => imbalance
            metrics[f"moe/{name}/min_load_ratio"] = (f.min() / fair).item()  # ->0 => a dying expert
            metrics[f"moe/{name}/n_starved"] = len(starved)

        metrics["moe/alerts"] = alerts
        for a in alerts:
            print("[MoE ALERT]", a)
        return metrics


# --- Demo: a healthy layer vs. a collapsed layer ---
E, k, N = 8, 2, 4096
monitor = ExpertLoadMonitor(n_experts=E, k=k, log_every=1, starve_frac=0.2, patience=3)

torch.manual_seed(0)
healthy_logits = torch.randn(N, E)          # uniform-ish routing

torch.manual_seed(1)
collapsed_logits = torch.randn(N, E)
collapsed_logits[:, 0:4] += 4.0             # experts 0-3 dominate
collapsed_logits[:, 4:8] -= 4.0             # experts 4-7 starve

for _ in range(3):  # 3 windows == patience -> alert fires on the last call
    out = monitor.log_step({
        "layer0.healthy": healthy_logits,
        "layer1.collapsed": collapsed_logits,
    })

print(out)

Expected output: layer0.healthy reports router_entropy ~= 0.83, max_load_ratio ~= 1.0, n_starved = 0; layer1.collapsed reports router_entropy ~= 0.54, experts 4-7 at min_load_ratio = 0.0, n_starved = 4, and on the third call four [MoE ALERT] lines fire for layer1.collapsed (none for layer0.healthy). The starvation rule in words: an expert holding fewer than starve_frac/E = 0.2/8 = 0.025 of dispatched tokens — 20% of its fair share of 1/E = 0.125 — counts as starving; patience consecutive logging windows below that threshold trigger the alert.

Further reading¶

• Shazeer, Mirhoseini, Maziarz, Davis, Le, Hinton, Dean — Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (2017). The paper that introduced noisy top-k gating and modern sparse MoE.
• Lepikhin, Lee, Xu, Chen, Firat, Huang, Krikun, Shazeer, Chen — GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (2020). Top-2 routing, capacity, and expert-parallel all-to-all at scale.
• Fedus, Zoph, Shazeer — Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (2021). Top-1 routing, the auxiliary loss formulation, and bf16 stability tricks.
• Zoph, Bello, Kumar, et al. — ST-MoE: Designing Stable and Transferable Sparse Expert Models (2022). The router z-loss and a careful study of stability and fine-tuning.
• Zhou, Lei, Liu, et al. — Mixture-of-Experts with Expert Choice Routing (2022). The expert-choice scheme that balances load by construction.
• Jiang, Sablayrolles, Roux, et al. (Mistral AI) — Mixtral of Experts (2024). The open decoder-only MoE with top-2 renormalized gating.
• Dai, Deng, Zhao, et al. (DeepSeek) — DeepSeekMoE: Towards Ultimate Expert Specialization (2024), and the DeepSeek-V3 Technical Report (2024). Fine-grained and shared experts; auxiliary-loss-free balancing.
• Gale, Narayanan, Zaharia, et al. — MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (2022). Block-sparse / grouped-GEMM kernels that make MoE dropless.
• Roller, Sukhbaatar, Szlam, Weston — Hash Layers For Large Sparse Models (2021). Fixed hash routing as the baseline every learned router should have to beat.
• Komatsuzaki, Puigcerver, Lee-Thorp, et al. — Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints (2023). Reuse a dense run to bootstrap an MoE.
• Muennighoff, Soldaini, Groeneveld, et al. (AI2) — OLMoE: Open Mixture-of-Experts Language Models (2024). A fully open MoE — data, code, logs and checkpoints — and the best end-to-end reproduction target.
• Jacobs, Jordan, Nowlan, Hinton — Adaptive Mixtures of Local Experts (1991). The original mixture-of-experts formulation that started it all.

Exercises¶

1. In one or two sentences, explain the rich-get-richer feedback loop that makes an un-regularized sparse router collapse onto a single expert. Then explain why the auxiliary loss \(\mathcal{L}_\text{aux} = \alpha E \sum_e f_e P_e\) multiplies a hard dispatch count \(f_e\) by a soft probability \(P_e\) instead of using either quantity alone.

2. Consider an MoE layer with \(E = 16\) experts, top-\(k = 2\) routing, \(d_\text{model} = 2048\), \(d_\text{ff} = 8192\), processing a batch of \(N = 4096\) tokens with capacity factor \(C_f = 1.5\). Treat each expert as a 2-matrix FFN (\(W_1, W_2\), ignore biases). Compute: (a) the average tokens per expert; (b) the per-expert capacity; © the number of tokens dropped at this layer if the router sends 900 tokens to one hot expert; (d) the total parameters in the layer’s experts; (e) the active expert-parameters used per token; (f) the total-to-active ratio.

3. Take the router logits from the chapter for a single token over 8 experts: \(h = [2.0,\ 1.0,\ 0.1,\ -1.0,\ 0.5,\ 3.0,\ -2.0,\ 0.2]\) and \(k = 2\). Compute the selected expert indices and their gate weights under both conventions — softmax-then-top-k (GShard/Switch) and top-k-then-softmax (Mixtral). State numerically why the two differ and what property renormalization guarantees.

4. DeepSeek-style fine-grained experts split each expert into \(m\) smaller ones and raise \(k\) by the same factor \(m\), keeping active parameters fixed. Starting from \(E = 8\), \(k = 2\), apply \(m = 4\). (a) What are the new \(E\) and \(k\)? (b) Show the active parameter count is unchanged. © Compute the number of distinct expert combinations a token can select before and after (\(\binom{E}{k}\)), and comment on why this is the point of the design.

5. Modify the chapter’s SparseMoE to add a single shared expert (DeepSeek-MoE style): an FFN that every token always passes through, added to the routed contribution. The routed path keeps its top-\(k\) over the remaining routed experts. Write the changed __init__ and forward, and explain in one line what problem the shared expert solves.

class SparseMoEWithShared(SparseMoE):
    def __init__(self, d_model, d_ff, n_experts, k=2,
                 capacity_factor=1.25, n_shared=1):
        super().__init__(d_model, d_ff, n_experts, k, capacity_factor)
        # Always-on experts; run densely on every token.
        self.shared = nn.ModuleList(
            [Expert(d_model, d_ff) for _ in range(n_shared)]
        )

    def forward(self, x):
        B, T, D = x.shape
        # Routed (sparse) path: reuse the parent's full implementation.
        routed_out, aux_loss = super().forward(x)   # (B, T, D), scalar

        # Shared (dense) path: every token, no routing, no capacity, no gate.
        flat = x.reshape(-1, D)                      # (N, D)
        shared_out = torch.zeros_like(flat)
        for se in self.shared:
            shared_out = shared_out + se(flat)       # sum over shared experts
        shared_out = shared_out.reshape(B, T, D)

        return routed_out + shared_out, aux_loss

# --- smoke test ---
torch.manual_seed(0)
moe = SparseMoEWithShared(d_model=32, d_ff=64, n_experts=8, k=2, n_shared=1)
x = torch.randn(4, 16, 32)
y, aux = moe(x)
print(y.shape, float(aux))     # torch.Size([4, 16, 32]) <positive scalar>

6. Implement the DeepSeek-V3 auxiliary-loss-free bias-adjustment balancer as a small helper. Given per-expert dispatch counts from the last step and a running per-expert bias vector \(b\), the bias is added to the routing logits for selection only (top-\(k\) uses \(h + b\)) while the gate weights still come from the unbiased logits \(h\). After each step, decrement \(b_e\) for overloaded experts and increment it for underloaded ones by a fixed step \(\gamma\). Write (a) a select_with_bias(logits, bias, k) function and (b) an update_bias(bias, counts, gamma) function, and state in one line why the bias is kept out of the gate weights.

import torch
import torch.nn.functional as F

def select_with_bias(logits, bias, k):
    """
    logits: (N, E) unbiased router logits h(x).
    bias:   (E,)   per-expert selection bias b (updated across steps).
    Returns (gate, idx): gate weights from the UNBIASED logits over the
    selected set (Mixtral-style renorm), and the selected indices.
    """
    biased = logits + bias                          # (N, E), selection only
    idx = torch.topk(biased, k, dim=-1).indices     # (N, k) selection uses h + b
    # Gate weights come from the UNBIASED logits at the chosen indices.
    chosen = torch.gather(logits, -1, idx)          # (N, k)
    gate = F.softmax(chosen, dim=-1)                # (N, k), rows sum to 1
    return gate, idx

@torch.no_grad()
def update_bias(bias, counts, gamma):
    """
    bias:   (E,) current per-expert bias, updated in place.
    counts: (E,) tokens dispatched to each expert this step.
    gamma:  fixed step size (e.g. 1e-3).
    Overloaded experts (above mean load) get bias decremented;
    underloaded experts get it incremented.
    """
    mean_load = counts.float().mean()
    # sign(mean - count): +1 if underloaded, -1 if overloaded, 0 if exactly mean.
    bias += gamma * torch.sign(mean_load - counts.float())
    return bias

# --- demo: an initially imbalanced router self-corrects ---
torch.manual_seed(0)
N, E, k = 4096, 8, 2
logits = torch.randn(N, E)
logits[:, 0] += 3.0                    # expert 0 is a hot favorite
bias = torch.zeros(E)
for step in range(200):
    gate, idx = select_with_bias(logits, bias, k)
    counts = torch.bincount(idx.reshape(-1), minlength=E)
    update_bias(bias, counts, gamma=1e-2)
print("final dispatch counts:", counts.tolist())   # far more uniform than at step 0
print("final bias:", bias.round(decimals=2).tolist())  # negative on expert 0