# Mixture of Experts (MoE) Architecture, routing, training, and inference optimization for sparse expert models. --- ## MoE Fundamentals ### Architecture ``` Input → Router → Top-k Experts → Combine → Output │ ├── Expert 1 (FFN) ├── Expert 2 (FFN) ├── ... └── Expert N (FFN) ``` ### Key Equations ``` # Router logits router_logits = x @ W_router^T # Top-k selection top_k_indices, top_k_weights = top_k(router_logits, k) # Expert outputs expert_outputs = [expert_i(x) for i in top_k_indices] # Weighted combination output = Σ(top_k_weights[i] * expert_outputs[i]) ``` ### Sparsity Benefits | Aspect | Dense | MoE (k=2, N=8) | |---|---|---| | **Params (Total)** | 7B | 47B (13B active) | | **FLOPs/Token** | 100% | ~25-30% | | **Training Cost** | 1x | ~1.5-2x (comm) | | **Inference Cost** | 1x | ~30-40% | | **Quality** | Baseline | Better at scale | --- ## Routing Mechanisms ### Top-k Gating ```python class TopKRouter: def __init__(self, d_model: int, num_experts: int, k: int = 2): self.router = nn.Linear(d_model, num_experts, bias=False) self.k = k def forward(self, x): logits = self.router(x) # [batch, seq, num_experts] weights, indices = torch.topk(logits, self.k, dim=-1) weights = F.softmax(weights, dim=-1) return weights, indices ``` ### Load Balancing (Critical) ```python # Auxiliary loss: encourage uniform expert usage def load_balancing_loss(router_logits, top_k_indices): # Fraction of tokens routed to each expert expert_mask = F.one_hot(top_k_indices, num_classes=num_experts).float() tokens_per_expert = expert_mask.mean(dim=[0, 1]) # Router probabilities router_probs = F.softmax(router_logits, dim=-1).mean(dim=[0, 1]) # Loss: minimize coefficient of variation return num_experts * (tokens_per_expert * router_probs).sum() # Z-loss (logit normalization) def z_loss(router_logits): return torch.logsumexp(router_logits, dim=-1).mean() ** 2 ``` ### Router Variants | Variant | Description | Pros | |---|---|---| | **Top-k** | Fixed k experts | Simple, predictable | | **Top-p** | Probability threshold | Adaptive compute | | **Expert Choice** | Experts choose tokens | Perfect load balance | | **Soft MoE** | Soft assignment | Differentiable | | **Hash Routing** | Deterministic hash | No router params | --- ## Expert Architectures ### Standard (FFN Experts) ```python class Expert(nn.Module): def __init__(self, d_model: int, d_ff: int): super().__init__() self.w1 = nn.Linear(d_model, d_ff, bias=False) self.w2 = nn.Linear(d_ff, d_model, bias=False) self.gate = nn.Linear(d_model, d_ff, bias=False) # SwiGLU def forward(self, x): return self.w2(F.silu(self.gate(x)) * self.w1(x)) ``` ### Shared + Specialized ```python class SharedExpertMoE(nn.Module): def __init__(self, d_model, num_experts, shared_experts=1): self.shared = nn.ModuleList([Expert(d_model, d_ff) for _ in range(shared_experts)]) self.routed = nn.ModuleList([Expert(d_model, d_ff) for _ in range(num_experts)]) self.router = Router(d_model, num_experts) def forward(self, x): # Shared always active shared_out = sum(e(x) for e in self.shared) # Routed top-k routed_out = self.moe_forward(x) return shared_out + routed_out ``` ### DeepSeek MoE (Fine-Grained) ```python # DeepSeek-V2/V3: 160-256 experts, 6-8 active # Key innovations: # 1. Fine-grained experts (smaller d_ff) # 2. Shared experts (always active, capture common knowledge) # 3. Auxiliary-loss-free load balancing (device-level) # 4. MLA (Multi-head Latent Attention) for KV compression class DeepSeekMoE(nn.Module): def __init__(self, config): self.num_experts = 160 # Fine-grained self.num_shared = 2 # Always active self.k = 6 # Top-6 routed # ... implementation ``` --- ## Training MoE ### Parallelism Strategy ```python # Expert Parallelism (EP) # Each expert on different device # Data parallel within expert # Tensor Parallelism (TP) within expert # Sequence Parallelism (SP) # Pipeline Parallelism (PP) across layers # 3D Parallelism: DP × TP × EP # DeepSeek-V3: 2048 GPUs, 3D parallelism ``` ### Training Stability ```python # Gradient clipping per expert def clip_expert_gradients(model, max_norm=1.0): for expert in model.experts: nn.utils.clip_grad_norm_(expert.parameters(), max_norm) # Router gradient scaling router_grad_scale = 1.0 / math.sqrt(num_experts) ``` ### Communication Optimization ```python # All-to-all for token dispatch # Overlap computation with communication async def moe_forward(x): # 1. Compute router logits logits = router(x) # 2. Dispatch tokens (async all-to-all) dispatch_future = async_all_to_all(x, indices) # 3. Compute shared experts while waiting shared_out = shared_experts(x) # 4. Wait for routed tokens routed_x = await dispatch_future # 5. Compute routed experts routed_out = routed_experts(routed_x) # 6. Combine return shared_out + combine(routed_out, weights) ``` --- ## Inference Optimization ### Expert Caching ```python class ExpertCache: def __init__(self, max_experts_in_memory=8): self.cache = {} # expert_id → (weights, device) self.max_size = max_experts_in_memory self.access_order = [] def get(self, expert_id): if expert_id in self.cache: self.access_order.remove(expert_id) self.access_order.append(expert_id) return self.cache[expert_id] # Load from disk weights = load_expert(expert_id) self.evict_if_needed() self.cache[expert_id] = weights self.access_order.append(expert_id) return weights ``` ### Speculative Expert Activation ```python # Predict which experts needed next class SpeculativeRouter: def __init__(self, router, predictor): self.router = router self.predictor = predictor # Small model predicting next experts def forward(self, x, prev_experts=None): # Use predictor to pre-load experts if prev_experts: predicted = self.predictor(prev_experts) preload_experts(predicted) return self.router(x) ``` ### Quantization ```python # Expert-specific quantization def quantize_moe(model): for i, expert in enumerate(model.experts): # High-traffic experts: higher precision if expert.traffic > threshold: quantize(expert, bits=8) # INT8 else: quantize(expert, bits=4) # INT4 # Router: keep FP16/BF16 quantize(model.router, bits=16) ``` --- ## MoE vs Dense: When to Use | Scenario | Recommendation | |---|---| | **Training from scratch, >10B params** | MoE (compute-efficient) | | **Fine-tuning dense model** | Keep dense | | **Latency-critical inference** | Dense or small MoE | | **Throughput-critical, batch inference** | MoE (better utilization) | | **Limited GPU memory** | MoE (offload experts) | | **Multilingual** | MoE (language-specific experts) | | **Code + Natural language** | MoE (specialized experts) | --- ## Notable MoE Models | Model | Total Params | Active Params | Experts | Active | Key Features | |---|---|---|---|---|---| | **Mixtral 8x7B** | 47B | 13B | 8 | 2 | Open, strong | | **Mixtral 8x22B** | 141B | 39B | 8 | 2 | Larger | | **DeepSeek-V2** | 236B | 21B | 160 | 6 | Fine-grained, MLA | | **DeepSeek-V3** | 671B | 37B | 256 | 8 | Aux-loss-free, MLA | | **Qwen-MoE** | ? | ? | 64 | 4 | Open | | **Nemotron-4-340B** | 340B | ? | ? | ? | NVIDIA | | **GPT-4** | ~1.8T | ~220B | 16? | 2? | Rumored | | **GLaM** | 1.2T | 97B | 64 | 2 | Google | --- ## Related Notes - [[ModelArchitectures/Model-Routing.md]] — Routing between models - [[ModelArchitectures/Inference-Optimization.md]] — MoE inference - [[ModelArchitectures/Small-Language-Models.md]] — Distillation from MoE - [[ModelArchitectures/Reasoning-Models.md]] — MoE for reasoning - [[Production/Cost-Optimization.md]] — MoE cost benefits - [[Production/Scaling-Strategies.md]] — MoE parallelism - [[Frontiers/Research-Directions.md]] — MoE research frontiers --- *Sources: "GShard" (Lepikhin et al.), "Switch Transformers" (Fedus et al.), "Mixtral" (Mistral AI), "DeepSeek-V2/V3" (DeepSeek), "GLaM" (Google), "MoE" surveys, "Expert Choice Routing", "Soft MoE", "Scaling MoE" papers, vLLM/SGLang MoE support.*