Mixture of Experts — scaling model capacity without activating every parameter

Yesterday we explored context parallelism: splitting a long sequence across GPUs while preserving attention between its tokens. Before that, tensor and pipeline parallelism divided model computation, and FSDP divided training state.

Today we'll study a different reason to distribute work: not every token needs to use every part of the model.

That is the intuition behind Mixture of Experts (MoE).

Instead of sending every token through one enormous feed-forward network, let a small router choose a few specialized networks for that token.

MoE can increase a model's total parameter count much faster than its computation per token. But sparse computation creates a new systems problem: the selected experts may live on different GPUs.

Dense versus sparse feed-forward layers¶

A conventional Transformer block contains attention and a feed-forward network (FFN). In a dense model, every token goes through the same FFN:

\[ h_{\mathrm{out}}=h+\operatorname{FFN}(h). \]

An MoE block replaces that FFN with a collection of \(E\) expert networks, often structurally similar FFNs. A router computes a score for each expert:

\[ z= W_rh,\qquad p=\operatorname{softmax}(z). \]

For top-\(k\) routing, only the \(k\) selected experts run on this token. A simplified output is

\[ \operatorname{MoE}(h) =\sum_{e\in\operatorname{TopK}(p)} g_e(h)\operatorname{FFN}_e(h), \]

where \(g_e\) is a routing weight, often normalized over the selected experts.

With eight experts and top-2 routing, all eight experts exist, but only two execute for a given token at that layer. Attention and other shared components still execute normally.

“Expert” does not necessarily mean a human-interpretable specialty such as mathematics or French. Experts are learned computational pathways, and their specialization may be subtle or distributed.

Why parameter count and compute separate¶

Imagine 32 MoE layers, each with eight experts of 120 million parameters. Suppose the model also has 6 billion shared parameters, including attention and other non-expert components.

Its total parameter count is approximately

\[ N_{\mathrm{total}} =6\text{B}+32(8)(0.12\text{B}) =36.72\text{B}. \]

With top-2 routing, the parameters used for one token's expert computations across those layers are approximately

\[ N_{\mathrm{active}} =6\text{B}+32(2)(0.12\text{B}) =13.68\text{B}. \]

So this illustrative model stores 36.72B parameters but activates roughly 13.68B per token.

That is not an exact FLOP count. Attention cost depends on context length, routing has overhead, and implementations can share or reuse parameters differently. But it explains why MoE models are often described using both total parameters and active parameters. Comparing only the first number can be misleading.

Crucially, inactive experts do not disappear from memory. Their weights still have to live somewhere, typically across the GPUs serving the model.

Expert parallelism: move tokens to the right GPU¶

Suppose eight experts are placed on eight GPUs. After attention, each GPU holds hidden states for some tokens. The router may choose experts that live elsewhere.

A simplified execution looks like this:

tokens on several GPUs
          |
     compute router
          |
     choose top-2 experts
          |
  ALL-TO-ALL token dispatch
          |
GPU 0: expert 0    GPU 1: expert 1    ...    GPU 7: expert 7
          |
     expert computation
          |
  ALL-TO-ALL return outputs
          |
   weighted combination

This is expert parallelism (EP): experts are distributed across workers, and token activations travel to the workers holding the selected experts.

Compare it with tensor parallelism. TP splits the matrix multiplication inside one layer across GPUs. EP lets different GPUs run different expert networks on different token assignments.

An expert can itself be tensor-parallel, and EP can coexist with context, pipeline, and data parallelism. The challenge is choosing a layout whose communication costs do not erase sparse-compute savings.

Worked example: 4,096 tokens, eight experts, top-2 routing¶

Suppose one MoE layer processes 4,096 tokens, with hidden dimension \(d=4096\) and BF16 activations. Each token selects two of eight experts.

There are

\[ 4096\times2=8192 \]

token-to-expert assignments. With perfectly balanced routing, each expert receives

\[ 8192/8=1024 \]

assignments.

Each dispatched hidden state occupies

\[ 4096\times2=8192\text{ bytes}=8\text{ KiB}. \]

Sending both selected experts a copy of every token's hidden state produces a total dispatch payload of

\[ 8192\times8192\text{ bytes}=64\text{ MiB}. \]

Returning one expert output per assignment involves another roughly 64 MiB of payload. This is aggregate logical payload across the expert group, not 128 MiB per GPU; local assignments need not cross a network, and real implementations add metadata, packing, and synchronization.

The surprising point: MoE avoids running six of eight experts per token, but it may introduce two all-to-all communication phases in every MoE layer. On a slow interconnect, moving activations can become the bottleneck.

The load-balancing problem¶

Routers are learned. They do not automatically distribute tokens evenly.

Imagine the same 8,192 assignments, but expert 0 receives 1,800, while the other seven average about 913 each. Even if all experts have equal compute capacity, expert 0 becomes a straggler: the layer cannot finish until its work finishes.

Some MoE systems impose an expert capacity, often expressed relative to the average load. With a capacity factor of 1.25:

\[ C=1.25\times1024=1280 \]

assignments per expert.

Expert 0 now has 520 assignments beyond its capacity. Depending on the routing algorithm, these might be dropped, rerouted, delayed, or handled by a capacity-free (“dropless”) execution scheme. Those choices affect quality, memory use, and latency.

During training, systems may add a load-balancing auxiliary loss or adjust routing biases to discourage pathological concentration. This introduces a tension: we want the router to choose useful experts, but we also want enough balance to keep hardware busy.

What MoE changes during training¶

Sparse activation does not mean only a tiny fraction of the model can learn. Different tokens activate different experts, so across a sufficiently diverse training batch, many or all experts may receive gradient updates.

But the number of training examples reaching an expert matters. An expert that rarely gets selected can learn slowly. Routing decisions can also become self-reinforcing: an expert receives fewer tokens, improves less, and may then be selected even less often.

The router's discrete top-\(k\) selection complicates gradient flow. Gradients normally pass through the selected experts and their continuous routing weights, while the discrete selection itself is not ordinarily differentiated in the same way as a smooth operation.

Training therefore needs to manage expert utilization, numerical stability, and communication, not merely the nominal active FLOPs.

Why MoE does not guarantee faster inference¶

For one token at a time, sparse expert computation may be too small to use a GPU efficiently. If the selected experts are remote, dispatch latency can dominate. Larger batches can improve expert utilization, but the router may send very different numbers of tokens to each expert.

The tradeoff is different in prefill and decode:

Prefill offers many tokens to group by expert, often making expert GEMMs more efficient.
Decode may have only a few new tokens per request, making routing overhead and small expert batches more prominent.

This echoes Days 68–69: a lower theoretical FLOP count does not necessarily produce lower wall-clock latency. The roofline model must now include interconnect traffic and irregular load balance.

Mental model: a company with specialist teams¶

Imagine a company with eight specialist teams. A dispatcher sends each incoming task to the two teams best suited to it.

The company has far more total expertise than any single team, but an individual task uses only two teams. That is the capacity-versus-active-work advantage of MoE.

Now suppose most tasks are routed to the same team. The other teams sit idle while one queue grows. And if teams are in different offices, shipping documents between them takes time.

A better dispatcher must balance expertise, queue lengths, and transfer cost—not simply choose the most popular team every time.

The deeper lesson¶

The previous lessons divided the same computation across devices. MoE changes which computation happens at all.

\[ \boxed{\text{MoE trades dense computation for sparse routing, expert storage, and communication.}} \]

It is especially attractive when extra parameter capacity improves model quality more than activating those parameters on every token would justify. But the true serving cost depends on expert placement, routing balance, batch size, and network speed.

The practical question is: Can we increase useful model capacity without turning token routing and uneven expert workloads into the new bottleneck?