Mixture of experts (MoE)

Deep dive

How it works

In: token (or patch) vectors x [T, d] · Out: y [T, d], each a gate-weighted sum of its top-k expert outputs · State: router W_r [d, E] and E expert FFNs · Cost: router O(T·d·E), experts O(T·k·FFN), two all-to-alls per layer under expert parallelism

Per layer, tokens x: [T, d], router W_r: [d, E]:

1. Score. logits = x·W_r. Gates g = softmax(logits) (Switch, GShard) or softmax over only the top-k logits (Mixtral). DeepSeek-V3 uses sigmoid affinities normalized over the selected experts.
2. Select. Top-k experts per token.
3. Dispatch. Group (token, expert) assignments by expert, a counting sort. Keep at most C per expert. With expert parallelism, an all-to-all sends each token to the GPU owning its expert.
4. Compute. Each expert runs its FFN on its contiguous batch (a grouped GEMM).
5. Combine. All-to-all back, then y_t = Σ_{e ∈ topk(t)} g_{t,e}·FFN_e(x_t). Dropped assignments contribute 0, so the token passes through only the residual.
6. Balance. Add L_aux = α·E·Σᵢ fᵢ·Pᵢ, where fᵢ = fraction of assignments routed to expert i (not differentiable) and Pᵢ = mean router probability for i (differentiable). It is minimized by a uniform split. Optionally add the router z-loss β·mean(logsumexp(logits)²) with β = 1e-3 (ST-MoE) to keep logits small.

tokens ─► router ─► top-k ─► sort by expert ─► all-to-all ─► E grouped FFNs ─► all-to-all ─► weighted sum
                                (cap C)            (send)                        (return)

Code

pub struct Routed { pub expert: Vec<usize>, pub gate: Vec<f32> } // T*k entries, token-major

/// Top-k routing. renorm = true: softmax over the selected logits (Mixtral);
/// false: softmax over all E, keep the selected probabilities (Switch/GShard).
pub fn route(logits: &[f32], t: usize, e: usize, k: usize, renorm: bool) -> (Routed, Vec<f32>) {
    let mut r = Routed { expert: Vec::with_capacity(t * k), gate: Vec::with_capacity(t * k) };
    let mut probs = vec![0.0f32; t * e];
    for tok in 0..t {
        let row = &logits[tok * e..][..e];
        let m = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let z: f32 = row.iter().map(|&v| (v - m).exp()).sum();
        for j in 0..e { probs[tok * e + j] = (row[j] - m).exp() / z; }
        let mut idx: Vec<usize> = (0..e).collect();
        idx.sort_by(|&a, &b| row[b].partial_cmp(&row[a]).unwrap()); // O(E log E); use partial select for big E
        let top = &idx[..k];
        let zk: f32 = top.iter().map(|&j| (row[j] - m).exp()).sum();
        for &j in top {
            r.expert.push(j);
            r.gate.push(if renorm { (row[j] - m).exp() / zk } else { probs[tok * e + j] });
        }
    }
    (r, probs)
}

/// Counting sort of assignments by expert, dropping overflow beyond `cap`.
/// Returns per-expert lists of assignment ids (assignment a = token a / k).
pub fn dispatch(r: &Routed, e: usize, cap: usize) -> (Vec<Vec<usize>>, usize) {
    let mut slots: Vec<Vec<usize>> = vec![Vec::with_capacity(cap); e];
    let mut dropped = 0;
    for (a, &ex) in r.expert.iter().enumerate() {                   // token order = priority
        if slots[ex].len() < cap { slots[ex].push(a) } else { dropped += 1 }
    }
    (slots, dropped)
}

/// Switch load-balancing loss alpha * E * sum_i f_i * P_i (equals alpha when perfectly uniform).
pub fn aux_loss(r: &Routed, probs: &[f32], t: usize, e: usize, alpha: f32) -> f32 {
    let mut f = vec![0.0f32; e];
    for &ex in &r.expert { f[ex] += 1.0 / r.expert.len() as f32; }
    let p: Vec<f32> = (0..e).map(|j| (0..t).map(|tok| probs[tok * e + j]).sum::<f32>() / t as f32).collect();
    alpha * e as f32 * f.iter().zip(&p).map(|(a, b)| a * b).sum::<f32>()
}

pub fn z_loss(logits: &[f32], e: usize) -> f32 {
    let n = logits.len() / e;
    logits.chunks(e).map(|row| {
        let m = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let lse = m + row.iter().map(|&v| (v - m).exp()).sum::<f32>().ln();
        lse * lse
    }).sum::<f32>() / n as f32
}

/// Combine: y[tok] += gate * expert_out[a]. Dropped assignments add nothing (residual carries the token).
pub fn combine(r: &Routed, slots: &[Vec<usize>], expert_out: &dyn Fn(usize, usize) -> Vec<f32>,
               t: usize, k: usize, d: usize) -> Vec<f32> {
    let mut y = vec![0.0f32; t * d];
    for (ex, list) in slots.iter().enumerate() {
        for &a in list {
            let tok = a / k;
            let out = expert_out(ex, tok);
            for c in 0..d { y[tok * d + c] += r.gate[a] * out[c]; }
        }
    }
    y
}

fn main() {
    let (t, e, k, d) = (64, 8, 2, 4);
    let logits: Vec<f32> = (0..t * e).map(|i| (((i * 2654435761usize) >> 7) % 1000) as f32 / 250.0).collect();
    let (r, probs) = route(&logits, t, e, k, true);
    for tok in 0..t { assert!((r.gate[tok * k] + r.gate[tok * k + 1] - 1.0).abs() < 1e-5); }
    let cap = ((1.25 * (t * k) as f32) / e as f32).ceil() as usize;     // capacity factor 1.25
    let (slots, dropped) = dispatch(&r, e, cap);
    let loads: Vec<usize> = slots.iter().map(|s| s.len()).collect();
    println!("capacity {cap}, loads {loads:?}, dropped {dropped}");
    println!("aux loss {:.4} (0.01 = uniform), z-loss {:.2}", aux_loss(&r, &probs, t, e, 0.01), z_loss(&logits, e));
    let x: Vec<f32> = (0..t * d).map(|i| i as f32 * 0.01).collect();
    let expert = |ex: usize, tok: usize| -> Vec<f32> { (0..d).map(|c| x[tok * d + c] * (ex + 1) as f32).collect() };
    let y = combine(&r, &slots, &expert, t, k, d);
    println!("y[0] = {:?}", &y[..d]);
}

Complexity and constants

• Router: 2·T·d·E FLOPs, tiny. Experts: k/E of the dense-equivalent FLOPs.
• All-to-all per MoE layer (forward): each token's d-vector is sent k times out and k times back. For d = 4096 in bf16 that is 8 KiB per assignment per direction. This traffic, not FLOPs, usually bounds MoE training throughput across nodes.
• Capacity factor 1.0-1.25 in training (Switch), larger at evaluation. Every 0.25 of slack costs that fraction of extra padded expert compute in fixed-capacity implementations.
• Memory: all E experts' weights, gradients and optimizer states. Mixtral 8x7B in bf16 is about 93 GB of weights for 12.9B active parameters per token.

Efficiency

• Cost: experts cost k/E of the dense-equivalent FLOPs; Mixtral 8x7B runs 12.9B active of 46.7B parameters per token. Weights for all E experts stay resident (about 93 GB in bf16 for Mixtral). Across nodes the bound is the all-to-all: 8 KiB per assignment per direction at d = 4096 in bf16.
• Levers:
  • Grouped GEMM / dropless block-sparse kernels — up to 40% faster than Tutel and 2.4x faster than Megatron-LM dense models (Gale et al. 2023) — any training with skewed expert loads.
  • Limit nodes per token — DeepSeek-V3 uses at most 4, capping cross-node traffic (DeepSeek-AI 2024) — expert parallelism over many nodes.
  • Overlap all-to-all with compute — hides most of the communication (DeepSeek-V3 report) — large expert-parallel training.
  • Capacity factor 1.0-1.25 — each 0.25 of slack is that fraction of padded compute (Switch) — fixed-capacity implementations only.
  • Balanced routing (aux loss or bias) — avoids stragglers; the slowest expert sets step time — always.
• Naive — loop over experts, run E padded small matmuls, drop overflow · low GPU utilization, dropped tokens
• Standard — capacity-padded batches + all-to-all dispatch/combine (GShard, Switch) · k/E FLOPs plus padding
• State of the art — dropless grouped GEMM, node-limited routing, overlapped all-to-all, fp32 router · k/E FLOPs, communication mostly hidden

Implementation notes

• Use a grouped GEMM over the expert batches (one kernel, variable M per expert) instead of E separate matmuls. Dropless kernels (MegaBlocks) avoid both padding and dropping.
• Compute the router in fp32. Small logit differences decide routing, and bf16 rounding flips top-k choices.
• Shared experts (DeepSeekMoE): one or two experts every token always uses, plus fine-grained routed experts. It reduces redundancy across routed experts.
• Aux-loss-free balancing (DeepSeek-V3): add a per-expert bias to the scores used for top-k selection only (not to the gates) and nudge it up or down after each step by load.
• Limit routing fan-out across nodes (DeepSeek-V3: at most 4 nodes per token) to cap all-to-all traffic.
• Initializing from a dense checkpoint (sparse upcycling): copy the FFN into every expert, train a fresh router.

Pitfalls

• Router collapse: a few experts win early, get all the gradient, and the rest die. Watch per-expert load. Aux loss or bias balancing fixes it.
• Token dropping that depends on batch order: the same token gets different outputs in different batches, and evaluation is non-deterministic. Use a larger capacity or dropless at inference.
• Instability from large router logits in bf16. Use z-loss and fp32 router math.
• Gradient accumulation and expert parallelism: per-expert batches become tiny at small global batch, so GEMMs run far below peak.
• Counting "8x7B" as 56B parameters: attention and embeddings are shared, so Mixtral is 46.7B.

Variants and alternatives

In production

Mixtral 8x7B and 8x22B, DeepSeek-V2/V3, Qwen MoE models, Switch Transformer, GShard (translation), Grok-1, DBRX; Google's V-MoE (vision) and LIMoE (image-text). Systems: Megatron-Core MoE, DeepSpeed-MoE, MegaBlocks, vLLM and SGLang expert-parallel inference, DeepEP all-to-all kernels.

Implementations

Sources

• Riquelme et al., Scaling Vision with Sparse Mixture of Experts, NeurIPS 2021. paper
• Mustafa et al., Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts, NeurIPS 2022. paper
• Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, ICLR 2017. paper
• Lepikhin et al., GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, ICLR 2021. paper
• Fedus, Zoph, Shazeer, Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, JMLR 2022. paper
• Zoph et al., ST-MoE: Designing Stable and Transferable Sparse Expert Models, 2022. paper
• Jiang et al., Mixtral of Experts, 2024. paper
• Dai et al., DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, 2024. paper
• DeepSeek-AI, DeepSeek-V3 Technical Report, 2024. paper · code
• Gale, Narayanan, Young, Zaharia, MegaBlocks: Efficient Sparse Training with Mixture-of-Experts, MLSys 2023. paper · code
• Zhou et al., Mixture-of-Experts with Expert Choice Routing, NeurIPS 2022. paper
• Komatsuzaki et al., Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints, ICLR 2023. paper