Expert parallelism — how Mixture-of-Experts models spend more parameters without spending all the FLOPs

Over the last few lessons, we have learned to split an LLM across GPUs along several axes.

Data parallelism splits examples.

Tensor parallelism splits matrix operations.

Pipeline parallelism splits layers.

Sequence/context parallelism splits tokens.

Today we add a different kind of split:

This is the key idea behind Mixture-of-Experts (MoE) models and expert parallelism.

The surprising promise of MoE is:

A model can contain vastly more parameters than it uses for any one token.

That lets us increase model capacity without increasing per-token computation proportionally.

But the price is a new systems problem: tokens must be routed across the cluster to whichever experts they need.

Start with the dense Transformer MLP¶

In a normal Transformer block, every token passes through the same feed-forward network.

A simplified MLP is:

If the model processes tokens:

"The"
"cat"
"sat"
"on"
"the"
"mat"

then every token uses the same:

Conceptually:

"The" ─┐
"cat" ─┤
"sat" ─┤
"on"  ─┼──→ same MLP ──→ outputs
"the" ─┤
"mat" ─┘

The MLP contains a large fraction of a Transformer's parameters and FLOPs.

Suppose its hidden width is \(d\) and its intermediate width is \(d_{\text{ff}}\).

Ignoring biases, the two main matrices contain roughly:

parameters.

Every token activates all of them.

This is a dense layer.

Replace one MLP with many experts¶

Now imagine we create eight different MLPs:

Each expert has its own parameters.

A router examines each token representation \(x\) and chooses which expert should process it.

For example:

"The" → expert 2
"cat" → expert 5
"sat" → expert 5
"on"  → expert 1
"the" → expert 2
"mat" → expert 7

Instead of every token executing every expert, each token might execute only one or two.

This is sparse activation.

The model may contain:

expert parameters while a token uses only:

of them.

That is the core MoE bargain:

The router¶

How does the model choose an expert?

A common design learns a router matrix:

where \(N_E\) is the number of experts.

For token representation:

the router computes logits:

Then:

If there are eight experts:

Suppose:

With top-1 routing, the token goes to expert 2.

With top-2 routing, it might go to experts 2 and 5.

For top-\(k\) routing, define:

The MoE output can be written schematically as:

where \(\alpha_i\) is the routing weight for selected expert \(i\).

The crucial word is selected.

We do not evaluate every \(E_i(x)\).

A concrete parameter example¶

Suppose one dense MLP contains:

If we replace it with 16 experts of the same size, the expert bank contains:

But suppose each token activates only two experts.

Then expert computation per token is roughly equivalent to:

not 16 billion.

So we have:

total expert parameters: 16B
active expert parameters per token: 2B

This distinction between total parameters and active parameters is essential when reading about MoE models.

A 500B-parameter MoE model does not necessarily perform anything like the computation of a 500B dense model for every token.

The library mental model¶

Imagine a huge research library with 100 specialist librarians.

One knows medicine.

One knows tax law.

One knows Sanskrit.

One knows distributed systems.

One knows medieval history.

A dense model is like asking all 100 librarians every question.

That is expensive.

An MoE system has a receptionist who looks at the question and says:

"CUDA all-reduce latency?"
→ ask the distributed-systems librarian

"Drug interaction?"
→ ask the medical librarian

"Translate this Sanskrit verse?"
→ ask the Sanskrit librarian

The organization can employ 100 specialists while consulting only one or two for each question.

That gives the system much more total expertise without paying for all of it every time.

But now the receptionist becomes important.

If everyone gets routed to the same librarian, the whole system jams.

That is exactly the problem MoE training faces.

Expert parallelism: put different experts on different GPUs¶

Suppose we have eight experts and four GPUs.

We might place them like this:

GPU 0 → experts 0,1
GPU 1 → experts 2,3
GPU 2 → experts 4,5
GPU 3 → experts 6,7

Now consider a batch of tokens initially distributed across the GPUs.

After routing, GPU 0 might discover:

token A → expert 5
token B → expert 1
token C → expert 7

Only token B can remain local.

Token A must travel to GPU 2.

Token C must travel to GPU 3.

Across the entire batch, tokens must be rearranged according to expert destination.

This usually creates an all-to-all communication pattern.

What all-to-all means¶

Before routing:

GPU 0: A B C D
GPU 1: E F G H
GPU 2: I J K L
GPU 3: M N O P

Suppose the router chooses:

expert on GPU 0: A F K P
expert on GPU 1: B E J M
expert on GPU 2: C H I N
expert on GPU 3: D G L O

The system must reshuffle the tokens:

             ALL-TO-ALL

GPU 0 ─────┬────┬────┬────→
GPU 1 ─────┼────┼────┼────→
GPU 2 ─────┼────┼────┼────→
GPU 3 ─────┴────┴────┴────→

after dispatch:

GPU 0: A F K P
GPU 1: B E J M
GPU 2: C H I N
GPU 3: D G L O

Each expert processes its assigned tokens.

Then another communication step sends the resulting activations back to the GPUs that own the tokens' surrounding Transformer computation.

So an MoE layer often looks roughly like:

token activations
      ↓
    router
      ↓
   all-to-all
      ↓
local expert MLPs
      ↓
   all-to-all
      ↓
restore token order

This is why expert parallelism can be extremely sensitive to network bandwidth.

The compute advantage can become a communication problem¶

Suppose a dense layer requires:

FLOPs per token.

An MoE layer with many experts but top-2 routing might require approximately:

even if the total expert bank is enormous.

That is excellent computationally.

But every MoE layer may now require moving token representations across GPUs.

So once again we encounter the recurring law of distributed LLMs:

For tensor parallelism, we communicated partial matrix results.

For pipeline parallelism, activations crossed layer boundaries.

For context parallelism, K/V blocks crossed sequence boundaries.

For expert parallelism:

That is the distinctive mental model.

Why load balancing matters¶

Suppose we have four experts and 1,000 tokens.

Perfect balance would send:

tokens to each expert.

But imagine the router sends:

expert 0 → 700 tokens
expert 1 → 150 tokens
expert 2 → 100 tokens
expert 3 →  50 tokens

Even if all experts have identical compute capacity, the layer finishes only when expert 0 finishes its 700 tokens.

The other GPUs may sit idle.

So the throughput is controlled by the busiest expert.

This is analogous to pipeline imbalance, except now the imbalance is created dynamically by learned routing.

Why the router might collapse¶

The router is trained to choose experts that reduce the model's prediction loss.

Suppose early in training expert 3 happens to become slightly better for many tokens.

The router sends more tokens to expert 3.

Expert 3 then receives more training examples.

It improves faster.

The router now has even more reason to choose expert 3.

This can create a feedback loop:

expert gets slightly better
        ↓
router sends it more tokens
        ↓
expert gets more training
        ↓
expert becomes even better
        ↓
router sends even more tokens

Eventually one expert may receive a huge fraction of traffic while others are barely used.

This is often called expert collapse or routing imbalance.

Load-balancing losses¶

To discourage collapse, MoE training commonly adds an auxiliary objective that encourages more balanced expert usage.

The exact formula varies by architecture.

Conceptually, if:

is the fraction of tokens routed to expert \(i\), we would like something closer to:

A simple conceptual penalty might resemble:

Real implementations often use more carefully designed differentiable statistics involving routing probabilities and token fractions.

The total training objective becomes something like:

So the model is being asked to satisfy two goals:

predict tokens well
+
avoid pathological expert imbalance

The coefficient \(\lambda\) determines how strongly we care about balancing relative to language modeling.

Expert capacity¶

Even with balancing, routing is stochastic and batches are finite.

One expert may receive more tokens than expected.

Systems therefore often define an expert capacity.

If a batch contains \(N\) routed token assignments and there are \(E\) experts, the average is:

A capacity factor \(c>1\) might allocate:

token slots per expert.

For example, with:

each expert gets capacity:

Average load is 512, but each expert can accept up to 640 assignments.

What happens beyond capacity depends on the architecture: tokens may be dropped, rerouted, or handled through another mechanism.

The key systems idea is that dynamic routing creates dynamic memory and compute demand.

Experts can specialize — but not always how humans expect¶

It is tempting to imagine:

expert 1 → mathematics
expert 2 → French
expert 3 → programming
expert 4 → biology

Sometimes analyses do reveal specialization.

But learned experts do not have to align with human semantic categories.

An expert might specialize in something more abstract:

particular syntactic patterns
certain token frequencies
specific hidden-state geometries
particular stages of reasoning

The router operates on learned vectors, not human labels.

So "expert" is a useful systems term, but we should not assume every expert corresponds to a recognizable profession.

MoE does not make attention sparse¶

A common misunderstanding is that an MoE Transformer routes the entire model through one expert.

Usually the sparse experts replace particular dense sublayers, especially the feed-forward/MLP blocks.

A simplified MoE Transformer block might be:

input
  ↓
attention        ← shared dense computation
  ↓
residual
  ↓
router
  ↓
sparse experts   ← only selected MLP experts run
  ↓
residual

So every token may still pass through the same attention machinery.

The sparsity is primarily in the expert layer.

That means total per-token compute is not reduced by the same ratio as total parameter sparsity.

Training and inference have different MoE pressures¶

During training, we usually have many tokens in a batch.

That gives each expert a reasonable amount of work.

If 16,384 token assignments are spread over 64 experts, average load is:

assignments per expert.

Those can form efficient matrix multiplications.

During autoregressive inference, especially at low batch size, there may be very few tokens available at each decoding step.

Suppose only 32 tokens are being decoded across 64 experts.

Many experts may receive:

0 tokens
1 token
2 tokens

Tiny expert batches are poor GPU workloads.

This is one reason MoE serving is a different systems problem from MoE training.

Why batching becomes especially valuable for MoE inference¶

Imagine 1,000 simultaneous users generating tokens.

At one decode step, the serving system may have roughly 1,000 token representations to route.

Now an expert might receive dozens of tokens rather than one.

The system can batch those together into a more efficient GEMM.

So high-throughput serving can turn sparse routing from a liability into an advantage.

This creates an important inference tradeoff:

but waiting to accumulate tokens can increase latency.

Again, throughput and latency pull in different directions.

Expert parallelism can combine with every parallelism axis we already know¶

Suppose a giant MoE model uses:

and:

It is tempting to multiply all four numbers directly, but real layouts can be more subtle because tensor-parallel and expert-parallel groups may overlap or apply to different parts of the model.

Conceptually, however, each axis answers a different question:

DP → which examples?
PP → which layers?
TP → which pieces of a matrix?
EP → which experts?

For an MoE layer, a token may therefore experience something like:

token belongs to data-parallel replica 3
        ↓
reaches pipeline stage 2
        ↓
router selects expert 11
        ↓
token is sent to GPU owning expert 11
        ↓
expert's matrix multiply may itself be tensor-parallel
        ↓
result returns to token's normal execution path

A single token can cross several different logical GPU groupings within one Transformer block.

The deeper connection to model capacity¶

Why go through all this trouble?

Because dense scaling ties two quantities together:

and:

If we make every dense layer twice as large, we generally spend much more compute on every token.

MoE partially separates these quantities.

We can increase:

without increasing:

at the same rate.

In rough conceptual form:

while:

where:

• \(N_E\) is total number of experts,
• \(k\) is experts selected per token.

If:

then the ratio is:

The model can expose roughly 32 times as much expert parameter capacity as it activates for one token.

That is the architectural reason MoE is so attractive.

But parameters that are inactive still have to live somewhere¶

Sparse activation reduces compute.

It does not make the unused expert weights disappear.

If the model contains hundreds of billions of expert parameters, those weights still require memory.

So MoE changes the resource profile:

dense model:
large memory
large compute per token

MoE model:
even larger parameter memory
much smaller compute than an equally sized dense model
extra routing communication

This is why expert parallelism is so natural.

The expert weights are distributed across many GPUs because storing all of them on every GPU would defeat much of the point.

The deeper mental model: move compute to data, or data to compute¶

Most of the parallelism we have studied tries to arrange tensors so the needed data is already near the computation.

MoE makes the relationship unusually explicit.

Each expert is a large block of parameters sitting somewhere in the cluster.

A token arrives and asks:

Which parameters should process me?

The router answers.

Then the system often moves the token to those parameters.

So there are two possible worlds:

dense model:
every token visits the same machinery

MoE model:
the machinery is distributed,
and each token visits only selected machinery

This gives us today's main principle:

And expert parallelism is the systems mechanism that makes that conditional computation practical across many GPUs.

The broader lesson connects directly to everything we have learned about distributed LLMs:

Sparse computation saves FLOPs only if the system can route tokens to sparse parameters cheaply enough.

Once again, the neural-network idea and the distributed-systems idea are inseparable.