Accelerator Architecture advanced 9 min read 7 flashcards

Hardware Cost of Mixture-of-Experts

Sparse MoE models reduce FLOPs per token but introduce all-to-all communication, load-imbalance penalties, and memory pressure that can erase those savings unless the system is carefully co-designed.

Mixtral 8x7B holds 47 billion parameters yet activates only 13 billion per token - roughly the same arithmetic as a 13B dense model. On paper, that sounds like free capacity. In practice, the hardware bill is considerably more complicated.

Why MoE Looks Cheap on Paper

A standard Transformer FFN applies the same weight matrix to every token. An MoE layer replaces one large FFN with E smaller expert FFNs and a lightweight router that picks k of them per token. The FLOP count per token scales with k, not with E, so as you add experts you get more capacity without proportionally more computation.

Let d be the model dimension, d_ff the per-expert hidden size, E the number of experts, and k the top-k selection count. The FLOP ratio relative to a dense FFN of size E * d_ff is approximately:

FLOP ratio ≈ k / E

For Mixtral (k=2, E=8) this is 0.25 - a 4x reduction in FFN compute. Switch Transformer (k=1) pushes this to 1/E. The roofline model says: if your workload is compute-bound, this is a genuine win.

The problem is that modern large-scale inference and training are rarely in the compute-bound regime for the FFN block alone.

The All-to-All Communication Wall

When you distribute experts across devices - which you must at scale, because each expert is a separate weight tensor - routing tokens to their assigned experts requires sending activations across the interconnect. The canonical pattern is a pair of all-to-all collectives per MoE layer:

  1. Dispatch: scatter each token's hidden state to the device hosting its chosen expert.
  2. Combine: gather the expert outputs back to the originating device.

Each all-to-all moves B * k * d * sizeof(dtype) bytes across the fabric, where B is the local batch size. On a 2048-TPU v3 pod (as in GShard), the inter-chip bandwidth is the bottleneck, not the expert FLOPs. The NVLink or ICI bandwidth does not scale linearly with the number of experts - it scales with ring topology and bisection bandwidth.

A rough cost model for a single all-to-all on N devices with bisection bandwidth B_net:

t_comm ≈ (B * k * d * dtype_bytes) * (N - 1) / (N * B_net)

At large N, this asymptotes to B * k * d * dtype_bytes / B_net, independent of N. That floor is non-trivial: for a batch of 512 tokens, d=4096, k=2, bfloat16, and B_net=600 GB/s (NVLink 3.0), you get roughly 14 microseconds per all-to-all, per MoE layer. A 32-layer model with MoE in half its layers accumulates ~224 microseconds of communication overhead per forward pass. Compared to a few hundred microseconds of total compute, that is not negligible.

Expert Capacity and Load Imbalance

Routing is dynamic: tokens self-select experts via a learned router. Nothing guarantees uniform load. If tokens cluster on a popular expert, you have two choices:

  • Drop tokens: set a capacity factor C (e.g. 1.25 of the uniform load), and discard tokens that overflow the expert's buffer. They pass through unchanged, losing that layer's transformation.
  • Allow overflow: dynamically extend buffers, breaking the fixed-shape requirement that TPU/GPU kernels need for efficient execution.

Switch Transformers introduced the capacity factor formulation explicitly:

expert_capacity = floor((tokens_per_batch / num_experts) * capacity_factor)

Setting C=1.0 means any imbalance drops tokens. Setting C=2.0 wastes half the expert's compute budget when load is balanced. The practical sweet spot (often 1.2-1.5) means you are padding memory and compute to insure against variance. That padding cost is proportional to E * capacity_per_expert * d_ff * d in activation memory.

An auxiliary load-balancing loss - a differentiable approximation to the fraction of tokens routed to each expert - pushes the router toward uniformity but never eliminates the problem entirely, especially for long-tail token distributions in multilingual or code-heavy corpora.

Memory Pressure: Parameters You Cannot Share

The entire point of MoE is that most weights are dormant per token. But they must all reside in memory. A dense 13B model and an MoE model with 13B active parameters but 47B total parameters have very different memory footprints:

Model Total params Active per token Weights in VRAM
Dense 13B 13B 13B ~26 GB (bf16)
Mixtral 8x7B 47B 13B ~94 GB (bf16)
Switch-C (1.6T) 1571B ~7B ~3.1 TB (bf16)

This is not a theoretical concern. Serving Mixtral 8x7B in bf16 requires at minimum two 80 GB A100s just for the weights. The routing logic, KV cache, and activations add on top. For inference at batch size 1 - the latency-sensitive case - you pay the full memory cost but get minimal benefit from the FLOP reduction because matrix multiplications over a single token are already memory-bound, not compute-bound. You end up reading 47B parameters from DRAM to compute what a 13B dense model would compute with 13B.

The sweet spot for MoE hardware efficiency is large batch, high throughput inference, where the arithmetic intensity of the expert GEMMs finally exceeds the roofline knee and compute savings are real.

When It Falls Down

Latency-critical, small-batch serving. At batch size 1 or small batches, the expert GEMMs are memory-bandwidth-bound regardless. You pay the full parameter memory cost with no compute savings. A dense model of equivalent active-parameter count is strictly cheaper to serve.

Interconnect-limited clusters. Systems with PCIe-only GPU-to-GPU links (e.g. consumer hardware, some cloud VMs) face severe all-to-all penalties. The communication cost can dwarf the compute savings even at large batch sizes. MoE expert parallelism is effectively off the table without high-bandwidth interconnects (NVLink, NVSwitch, InfiniBand HDR, or on-package fabric).

High expert counts with small batch sizes. If E is large (e.g. 64 or 256 as in some Switch variants) and the batch is modest, many experts receive zero or one token. Their capacity slots are wasted, their weight tensors still consume memory, and the router overhead (a softmax over E logits per token per layer) becomes non-trivial.

Token dropping under distribution shift. A capacity factor tuned during training on balanced data may drop tokens excessively at inference on out-of-distribution inputs where one expert is heavily preferred. The model silently degrades rather than raising an error.

Gradient imbalance in fine-tuning. In supervised fine-tuning on narrow domains, the load-balancing loss may fight the task loss. Experts specialise to the domain's token distribution, leaving others undertrained. Catastrophic forgetting can concentrate into specific experts.

Quantisation and expert heterogeneity. Per-tensor quantisation of expert weights requires separate scale factors per expert. With 64+ experts this can add non-trivial metadata overhead and complicate kernel fusion.

Further Reading

Check yourself

7 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track