moe communication costs

**Mixture of Experts (MoE)** is the sparse-activation architecture that scales a neural network to trillions of parameters while keeping per-token compute fixed — each input activates only a small subset of "expert" sub-networks selected by a learned router, so total model capacity grows without proportional growth in inference FLOPs. GPT-4, Mixtral 8×7B, Switch Transformer, DeepSeek-V2, and Grok all use MoE layers to achieve frontier accuracy at a fraction of the cost of an equivalently-sized dense model. **The core idea — conditional computation.** In a dense Transformer, every token passes through every FFN parameter. In an MoE Transformer, the standard FFN block is replaced by $N$ parallel expert FFNs plus a lightweight gating (router) network. For each token, the router selects the top-$k$ experts (typically $k = 1$ or $k = 2$), and only those experts run. If $N = 64$ and $k = 2$, the model has 64× the parameters of one expert but only 2× the compute per token — a ~32× parameter-to-FLOP leverage ratio. **Router design.** The router $G(x)$ maps a token embedding $x \in \mathbb{R}^d$ to a probability distribution over experts: $$G(x) = \text{softmax}(W_g \cdot x + \epsilon)$$ where $W_g \in \mathbb{R}^{N \times d}$ is a learned matrix and $\epsilon$ is optional noise for exploration during training. The top-$k$ entries of $G(x)$ select which experts fire; the corresponding softmax weights become the mixture coefficients for combining expert outputs: $$y = \sum_{i \in \text{TopK}(G(x))} G(x)_i \cdot E_i(x)$$ **Load balancing — the critical auxiliary loss.** Without intervention, training collapses: a few popular experts attract most tokens, receive the strongest gradients, and become even more popular (expert collapse). The fix is an auxiliary loss that penalizes uneven load: $$\mathcal{L}_{\text{aux}} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot p_i$$ where $f_i$ is the fraction of tokens actually routed to expert $i$ and $p_i$ is the mean router probability assigned to expert $i$ across the batch. Minimizing $\mathcal{L}_{\text{aux}}$ pushes the router toward uniform dispatch. Typical $\alpha$: 0.01–0.1. **Capacity factor and token dropping.** Each expert can process at most $C = \text{capacity\_factor} \times T/N$ tokens per batch (where $T$ = total tokens). Tokens that overflow are either dropped (Switch Transformer, capacity factor ≈ 1.25) or re-routed to a shared fallback expert. DeepSeek-V2 eliminates dropping entirely with a "shared expert" that all tokens pass through, plus routed experts for specialization. | Architecture | Experts | Top-k | Key innovation | Model capacity | Active params/token | |---|---|---|---|---|---| | Switch Transformer (2022) | 128–2048 | 1 | Simplified to $k$=1, capacity routing | 1.6T params (C variant) | ~1/128 of total | | Mixtral 8×7B (2024) | 8 | 2 | Dense-quality at 7B active cost | 47B total | 13B | | GPT-4 (2023, reported) | ~16 | 2 | Multi-head MoE per layer | ~1.8T total | ~220B | | DeepSeek-V2 (2024) | 160 routed + 2 shared | 6 | Fine-grained experts + shared | 236B total | 21B | | Grok-1 (2024) | 8 | 2 | Open-weight frontier MoE | 314B total | ~86B | | DBRX (Databricks, 2024) | 16 | 4 | Fine-grained 16-expert design | 132B total | 36B | **Training — expert parallelism.** MoE layers require a collective all-to-all communication: tokens are gathered at the GPU hosting their assigned expert, processed, then scattered back. This is the defining bottleneck of MoE training at scale. A typical layout: data-parallel across most of the model, expert-parallel across the MoE FFN. With $P$ GPUs and $N$ experts, each GPU holds $N/P$ experts and receives tokens routed to them from all other GPUs. **Inference — why MoE is hard on hardware.** Although only top-$k$ experts compute per token, all $N$ experts must reside in memory (HBM) because the router's selections are input-dependent and change every token. This means: - **Memory** scales with total parameters (not active parameters). A 1.8T-parameter MoE at fp16 needs ~3.6 TB of HBM — requiring multi-node inference. - **Compute** scales with active parameters ($k$ experts × expert size). The arithmetic intensity is low (small matrix per expert), making MoE decode memory-bandwidth-bound even more severely than dense models. - **Expert offloading** (expert-to-CPU/SSD): exploits the sparsity by keeping only hot experts in HBM and paging cold ones on demand — but latency spikes when a token routes to a cold expert. **Chip-design implications.** An MoE-optimized accelerator needs: (1) massive HBM capacity to hold all experts (HBM3E 6-stack or 8-stack configurations), (2) very high memory bandwidth (the decode bottleneck), (3) fast all-to-all interconnect between chips for expert parallelism (NVLink, UALink, or custom mesh), and (4) a small low-latency router engine that can select experts before launching the main compute — a pattern the CFS Inference Simulator models at /infer. ```svg Mixture of Experts — token routing through a Transformer MoE layer Input tokens token₁ token₂ token₃ token₄ Router G(x) = softmax(Wg·x) select top-k Expert 1 (FFN) W₁ · ReLU · W₂ Expert 2 (FFN) W₁ · ReLU · W₂ Expert 3 (FFN) ★ selected Expert N (FFN) ★ selected route Σ weighted MoE layer output Key insight: N experts in memory, only top-k compute per token • Parameters scale ×N (capacity) but FLOPs scale ×k (active) → ~N/k leverage ratio • All experts must be in HBM (input-dependent routing) → memory-bound inference, demands high BW ``` **The MoE scaling law.** Empirically, an MoE model with $N$ experts and active parameters $A$ performs roughly like a dense model of size $A \cdot N^{0.3}$ in terms of loss — better than $A$ alone, but not as good as a dense model of size $A \cdot N$. The exponent varies (0.2–0.4) depending on routing quality and expert granularity. This makes MoE the dominant architecture for cost-efficient frontier models: you get 80% of the benefit of a model 5–10× larger at only the inference cost of the active slice. **Fine-grained vs coarse-grained experts.** Early MoE (Switch, Mixtral) used 8–128 experts each the size of a full FFN. DeepSeek-V2 and later designs shrink expert size dramatically (e.g. 256 experts, each 1/16 the FFN width) so more experts can be selected per token ($k = 6$–8) without increasing total compute — this gives smoother routing, less load imbalance, and better generalization because each token assembles a more nuanced combination. **What MoE changes for the hardware stack.** The shift from dense to MoE fundamentally re-weights the hardware bottleneck hierarchy: memory capacity and bandwidth matter more than peak FLOPS, inter-chip interconnect bandwidth becomes the training limiter (all-to-all), and the router decision latency is on the critical path for every single token. This is why the CFS platform models MoE workloads across the HBM (/hbm), KV-cache (/kvcache), and inference (/infer) simulators — each captures a different facet of the MoE serving challenge.

Go deeper with CFSGPT

Get AI-powered deep-dives, save terms, and run advanced simulations — free account.

Create Free Account