Nearly every frontier-class open-weight model shipped in the last two years is a Mixture-of-Experts (MoE): DeepSeek-V3, the Qwen3 MoE family, Mixtral before them, gpt-oss. The appeal is simple arithmetic. A model with 200B total parameters that activates 15B per token costs roughly 15B parameters' worth of FLOPs per token but carries 200B parameters' worth of knowledge. You buy capacity with memory instead of compute.
That trade is only a bargain if your infrastructure is built for it. Teams that treat an MoE checkpoint as "a dense model with more weights" hit the same three walls every time: the GPU memory bill is set by total parameters while the speed-up they budgeted for was based on active parameters; fine-tuning silently collapses the router; and serving throughput falls off a cliff at exactly the batch sizes they benchmarked at. This tutorial covers what is different about MoE in PyTorch, how to fine-tune one without breaking it, and how to size a serving deployment honestly.
How an MoE layer actually works
In a dense transformer, every token goes through the same feed-forward network. In an MoE transformer, the FFN is replaced by E parallel expert FFNs plus a small linear router. For each token, the router scores the experts, takes the top k (usually 1, 2, or 8), and sends the token only to those.
A minimal, readable version — this is not the fast path, but it is the mental model:
import torch
import torch.nn as nn
import torch.nn.functional as F
class MoEFeedForward(nn.Module):
def __init__(self, d_model, d_ff, num_experts=8, top_k=2):
super().__init__()
self.top_k = top_k
self.num_experts = num_experts
self.router = nn.Linear(d_model, num_experts, bias=False)
self.experts = nn.ModuleList([
nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))
for _ in range(num_experts)
])
def forward(self, x): # x: [tokens, d_model]
logits = self.router(x) # [tokens, E]
weights, idx = torch.topk(logits, self.top_k, dim=-1)
weights = F.softmax(weights, dim=-1, dtype=torch.float32).to(x.dtype)
out = torch.zeros_like(x)
# Loop over experts, not tokens: each expert sees one batched matmul.
for e in range(self.num_experts):
token_ids, slot = (idx == e).nonzero(as_tuple=True)
if token_ids.numel() == 0:
continue
out.index_add_(
0, token_ids,
self.experts[e](x[token_ids]) * weights[token_ids, slot].unsqueeze(-1),
)
return out, logits
Three details in that snippet matter more than they look:
- The router softmax runs in FP32. Router logits are small and close together; computing the gate in BF16 makes expert selection jitter between steps and destabilizes training. Every production implementation upcasts here.
- Iteration is over experts, not tokens. The gather/scatter formulation turns a per-token branch into
Edense matmuls. Real kernels go further and use grouped GEMM (torch._grouped_mmin recent PyTorch, or the Triton grouped-GEMM kernels shipped with vLLM and TorchTitan) so all experts run in a single launch. - The layer returns its logits. You need them for the load-balancing loss below.
Why MoE breaks naive fine-tuning
The router collapses
Left alone, routers are self-reinforcing: an expert that gets slightly more tokens gets better, so it gets more tokens. Within a few hundred steps you can be running a top_k=2 model where two experts serve 80% of traffic and the rest are dead weight you are still paying VRAM for. Pretraining fixes this with an auxiliary load-balancing loss:
def load_balancing_loss(router_logits, top_k, num_experts):
# router_logits: [tokens, E] collected from every MoE layer
probs = F.softmax(router_logits.float(), dim=-1)
_, idx = torch.topk(router_logits, top_k, dim=-1)
mask = F.one_hot(idx, num_experts).sum(dim=1).float() # [tokens, E]
f = mask.mean(dim=0) # fraction of tokens dispatched per expert
p = probs.mean(dim=0) # mean router probability per expert
return num_experts * torch.sum(f * p)
loss = ce_loss + 0.01 * sum(load_balancing_loss(l, top_k, E) for l in all_router_logits)
The coefficient is conventionally 0.01. Too high and the router is pushed toward uniformity, hurting quality; too low and it collapses. Note that some newer models (DeepSeek-V3 among them) replace the auxiliary loss with a per-expert bias term updated outside the gradient — if you are fine-tuning one of those, do not bolt an aux loss back on.
In practice, the safest fine-tuning recipe is to not train the router at all. Freeze it, freeze the experts, and attach LoRA adapters to the attention projections only. You keep the pretrained routing distribution intact and the adapter is a few dozen MB:
from peft import LoraConfig, get_peft_model
cfg = LoraConfig(
r=32, lora_alpha=64, lora_dropout=0.05, task_type="CAUSAL_LM",
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], # attention only
)
model = get_peft_model(base_model, cfg)
If domain adaptation demands touching the FFNs, add the expert projections to target_modules but keep gate/router modules out of it, and watch expert utilization as a training metric.
Utilization is a first-class metric
Log it. A simple hook per MoE layer, accumulated over a few hundred steps:
counts = torch.zeros(num_experts, device="cuda")
def tally(module, args, output):
_, logits = output
_, idx = torch.topk(logits, module.top_k, dim=-1)
counts.index_add_(0, idx.flatten(), torch.ones(idx.numel(), device=logits.device))
for m in model.modules():
if isinstance(m, MoEFeedForward):
m.register_forward_hook(tally)
Healthy top_k=2 models sit within roughly 2x of uniform across experts. A max/min ratio above 10x means the router has collapsed and your effective parameter count is a fraction of what you are paying to host.
Memory: budget for total, not active, parameters
This is where most cost models go wrong. Take a 120B-total / 5B-active MoE in BF16:
| Item | Bytes | 120B total |
|---|---|---|
| Weights (BF16) | 2 B/param | ~240 GB |
| Optimizer state (AdamW FP32 m, v) | 8 B/param | ~960 GB |
| FP32 master weights + grads | 8 B/param | ~960 GB |
Full fine-tuning is a multi-node problem before you have loaded a single token. LoRA on frozen BF16 experts drops it to weights plus a rounding error — ~240 GB, i.e. four H100s or two H200s — and that is why nearly all MoE fine-tuning in the field is adapter-based.
For sharding, the relevant primitive in modern PyTorch is expert parallelism layered on top of FSDP2: shard the non-expert parameters with fully_shard, and place whole experts on specific ranks with a DTensor mesh so each rank owns a distinct expert group.
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard
mesh = init_device_mesh("cuda", (2, 4), mesh_dim_names=("dp", "ep"))
for block in model.layers:
fully_shard(block, mesh=mesh["dp"]) # attention + norms sharded data-parallel
# experts distributed across mesh["ep"]; dispatch becomes an all-to-all
The cost of expert parallelism is that token dispatch turns into two all-to-all collectives per MoE layer. Across NVLink that is cheap; across a slow inter-node fabric it dominates. Keep the expert-parallel dimension inside a node whenever the model fits.
Serving: throughput does not scale the way you expect
At batch size 1, an MoE is memory-bandwidth-bound like any decoder, and only the active experts' weights get read — which is exactly why MoEs decode fast. As batch size grows, tokens in the batch route to different experts, so the union of experts touched approaches all of them. Past a few dozen concurrent sequences you are streaming the entire weight matrix per step, and your effective arithmetic intensity looks dense again.
Consequences for capacity planning:
- Benchmark at production concurrency. A single-stream benchmark of an MoE overstates throughput per GPU more than it does for a dense model.
- Expect a flatter throughput curve. Batching pays off less than for a dense model of equal active size.
- Quantize aggressively. Because memory is the binding constraint, weight-only INT4/FP8 on experts buys more here than on a dense model.
torchao's FP8 and INT4 weight-only paths, or an FP8 checkpoint served by vLLM, are the standard moves. - Watch KV cache. MoE affects the FFN, not attention: KV-cache math is unchanged, and once experts are quantized the cache often becomes the new limit.
A representative vLLM launch for a large MoE across a node:
vllm serve Qwen/Qwen3-235B-A22B-FP8 \
--tensor-parallel-size 4 \
--enable-expert-parallel \
--max-model-len 32768 \
--gpu-memory-utilization 0.90
--enable-expert-parallel splits experts across ranks instead of splitting each expert (tensor-parallel style). On fast intra-node interconnect it is usually the better layout at high concurrency; on anything slower, measure both.
Serving LoRA adapters on top works the same as for dense models — --enable-lora with per-request adapter selection — which is another argument for keeping fine-tunes adapter-shaped.
When an MoE is the wrong answer
MoEs win when you need broad capability and have VRAM but not compute headroom. They are a poor fit when:
- You are deploying to a single GPU, an edge device, or anywhere memory is the scarce resource. A dense 14B beats a 120B-A5B you cannot load.
- Your workload is one narrow domain. Most of the experts contribute little and you are renting memory for capability you never invoke; distilling a dense small model is cheaper to run and simpler to operate.
- Your team cannot own a multi-GPU serving stack. MoE operations — expert-parallel layouts, utilization monitoring, quantized checkpoints — carry real complexity.
A checklist before you commit
- Compute the VRAM bill from total parameters at your target precision, then add KV cache at your target context and concurrency.
- Decide adapter-based fine-tuning first; reach for full or expert-only fine-tuning only when adapters demonstrably fall short.
- Freeze the router unless you have a deliberate plan and load-balancing metrics in your training dashboard.
- Log expert utilization from step one; alert if max/min utilization exceeds 10x.
- Benchmark serving at production concurrency, with FP8 or INT4 experts, both tensor-parallel and expert-parallel.
- Compare against a dense model of similar active size on your own evaluation set before assuming the MoE wins.
Mixture-of-Experts is the reason open-weight quality kept climbing while inference cost per token did not. Getting the benefit, though, is an infrastructure exercise as much as a modelling one: the memory arithmetic, the router hygiene and the concurrency behaviour all have to be right before the parameter count on the model card means anything.
IntelliSensei's consultants size, fine-tune and deploy MoE models on PyTorch for enterprise teams. If you are weighing an MoE against a dense model, or your fine-tune has quietly collapsed its router, get in touch.