Architecture Deep-Dive Guide
Table of Contents
- What is a Mixture of Experts (MoE)?
- How Experts Are Defined
- How the Router Works
- Two Architecture Options
- Key Design Decisions Explained
- VRAM and Compute Analysis
What is a Mixture of Experts (MoE)?
A standard Transformer has one FFN per layer that processes every token:
Token → Attention → FFN → Output
↑
(always this one FFN)
An MoE Transformer has many FFNs (experts) per layer, but each token only uses a few:
Token → Attention → Router → [Expert 3, Expert 17, Expert 28, Expert 31] → Output
↑ ↑
(learned selector) (only top-K of N experts activated)
The magic: You get N× more parameters (knowledge capacity) but only K× the compute of a single expert. A 20B MoE model with 32 experts and top-4 routing has 20B total params but only uses ~3.6B per token — comparable compute to a 3.6B dense model, but with the knowledge of a 20B model.
How Experts Are Defined
You do NOT manually define what each expert "knows." Experts are:
- Structurally identical — each expert is a small SwiGLU FFN:
class Expert(nn.Module):
def __init__(self, hidden_size=2880, intermediate_size=2880):
self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
def forward(self, x):
# SwiGLU: gate × value
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
Randomly initialized — all experts start the same (with random weights)
Automatically specialized during training — the router learns to send different token patterns to different experts. Through gradient flow, each expert adapts to the tokens it receives most frequently.
What Do Experts End Up Learning?
Research (DeepSeek-V3 Appendix C, OLMoE Section 5) shows experts naturally specialize in:
- Syntactic roles: One expert handles Python function definitions, another handles imports
- Semantic domains: Math expressions, natural language, code comments
- Structural patterns: Opening brackets, closing statements, docstrings
- Languages: In multilingual models, experts may specialize by programming language
You can visualize expert activation patterns after training to understand specialization, but you don't need to engineer it.
How the Router Works
The router is a simple learned linear layer:
class Router(nn.Module):
def __init__(self, hidden_size, num_experts):
# This single matrix IS the routing mechanism
# Each row is an "expert embedding" — a learned representation
# of what kind of tokens each expert should handle
self.gate = nn.Linear(hidden_size, num_experts, bias=False)
def forward(self, hidden_states):
# Step 1: Score each expert for each token
logits = self.gate(hidden_states) # [num_tokens, num_experts]
# Step 2: Convert to probabilities
# GPT-OSS style:
scores = torch.softmax(logits, dim=-1)
# OR DeepSeek-V3 style:
# scores = torch.sigmoid(logits)
# Step 3: Select top-K experts per token
topk_scores, topk_indices = torch.topk(scores, k=4) # top-4
# Step 4: Normalize gate weights (so they sum to 1)
gate_weights = topk_scores / topk_scores.sum(dim=-1, keepdim=True)
return gate_weights, topk_indices
Load Balancing
Without balancing, the router might send all tokens to the same 2-3 experts (rich-get-richer). This wastes the other experts and creates GPU compute imbalance.
Three approaches:
| Method | Used By | How It Works |
|---|---|---|
| Auxiliary Loss | GPT-OSS (coef=0.9), OLMoE (coef=0.01) | Extra loss term penalizing uneven expert load |
| Bias-Based (aux-loss-free) | DeepSeek-V3, Moonlight | Per-expert bias added to routing scores; adjusted online |
| Z-Loss | OLMoE, Qwen3 | Penalizes large router logits (stabilizes training) |
For your first training run, start with auxiliary loss (coef=0.01) + z-loss (coef=0.001). These are well-understood and supported in all frameworks. Aux-loss-free is better but requires custom implementation in Megatron-LM.
Two Architecture Options
Option A: GPT-OSS-20B Style (RECOMMENDED)
┌──────────────────────────────────────────────────────────┐
│ Token Embedding (vocab=201,088 → dim=2,880) │
├──────────────────────────────────────────────────────────┤
│ Layer 0: Sliding Attention (window=128) + MoE FFN │
│ Layer 1: Full Causal Attention + MoE FFN │
│ Layer 2: Sliding Attention + MoE FFN │
│ Layer 3: Full Causal Attention + MoE FFN │
│ ... (alternating pattern for 24 layers) │
│ Layer 23: Full Causal Attention + MoE FFN │
├──────────────────────────────────────────────────────────┤
│ RMS Norm → LM Head (dim=2,880 → vocab=201,088) │
└──────────────────────────────────────────────────────────┘
Each MoE layer:
Input → RMSNorm → GQA Attention → Residual →
RMSNorm → Router → Top-4 of 32 Experts → Weighted Sum → Residual
Key features:
- Alternating sliding/full attention: Saves 50% attention memory. Sliding window (128 tokens) handles local patterns; full attention handles long-range dependencies.
- GQA (64 query heads, 8 KV heads): Efficient attention with 8:1 head ratio.
- 32 experts, top-4: Each token activates 4 of 32 experts.
- SwiGLU with clamping (limit=7.0): Prevents activation explosions.
Option B: DeepSeek-V2-Lite Style
┌──────────────────────────────────────────────────────────┐
│ Token Embedding (vocab=102,400 → dim=2,048) │
├──────────────────────────────────────────────────────────┤
│ Layer 0: MLA Attention + Dense FFN (10,944 hidden) │ ← Dense!
│ Layer 1: MLA Attention + MoE FFN (64 experts, top-6) │
│ Layer 2: MLA Attention + MoE FFN │
│ ... (all MoE after layer 0) │
│ Layer 26: MLA Attention + MoE FFN │
├──────────────────────────────────────────────────────────┤
│ RMS Norm → LM Head │
└──────────────────────────────────────────────────────────┘
Each MoE layer:
Input → RMSNorm → MLA Attention → Residual →
RMSNorm → [2 Shared Experts (always)] + [Router → Top-6 of 64] → Residual
Key features:
- MLA (Multi-Head Latent Attention): Compresses KV cache from 4096 to 576 floats per token per layer (86% reduction). Enables much longer context at inference.
- Shared + Routed experts: 2 experts always active (shared knowledge) + 6 of 64 routed (specialized). 8 total active per token.
- Fine-grained experts: 64 small experts (1408 hidden dim) instead of fewer large ones. More combinatorial routing flexibility.
- Layer 0 is dense: Early layers don't benefit from MoE routing.
Key Design Decisions Explained
1. How Many Experts?
| Configuration | Total Experts | Active | Sparsity | Quality | Compute |
|---|---|---|---|---|---|
| GPT-OSS-20B | 32 | 4 | 8× | Good | Low |
| DeepSeek-V2-Lite | 64 | 6 | ~11× | Good | Low |
| GPT-OSS-120B | 128 | 4 | 32× | Great | Low |
| DeepSeek-V3 | 256 | 8 | 32× | SOTA | Medium |
| Kimi K2 | 384 | 8 | 48× | SOTA | Medium |
Rule of thumb: More experts = more knowledge capacity = better quality at same compute. But more experts = more memory for weights and more complex routing.
For 4×H100 inference: 32-64 experts is the sweet spot. All weights fit comfortably in BF16.
2. Tokenizer Choice
| Option | Vocab Size | Pros | Cons |
|---|---|---|---|
| o200k_harmony (GPT-OSS) | 201,088 | Best code tokenization, GPT-4o compatible | Large embedding table (579M params) |
| DeepSeek tokenizer | 102,400 | Good code+multilingual | Requires custom setup |
| Train your own (BPE) | 32K-64K | Optimized for your domain, smaller embeddings | Training cost, may miss rare tokens |
Recommendation: Use o200k_harmony (201K vocab) if you want GPT-OSS compatibility. Use a custom 64K vocab tokenizer if you want to minimize embedding params and optimize for your specific code/language distribution.
3. Context Length Strategy
Phase 1a: Pre-train at 4K context → cheap, covers most code files
Phase 1c: Extend to 32K → covers full files, multi-file context
Phase 1c: Extend to 131K with YaRN → repo-level understanding
YaRN (Yet Another RoPE Extension) lets you extend context without retraining from scratch.
GPT-OSS uses rope_type: "yarn" with factor: 32 (extends 4K → 131K).
4. Training Precision
| Precision | Memory | Speed | Quality | H100 Support |
|---|---|---|---|---|
| FP32 | 4 bytes/param | 1× | Best | ✓ |
| BF16 | 2 bytes/param | ~2× | Excellent | ✓ |
| FP8 | 1 byte/param | ~4× | Very Good | ✓ (native) |
Recommendation: Start with BF16. Switch to FP8 if you need faster training — H100 has native FP8 tensor cores. DeepSeek-V3 proved FP8 training works at scale.
VRAM and Compute Analysis
Training (16×H100, 1.28TB total VRAM)
For GPT-OSS-20B style (21B params):
Model weights (BF16): 21B × 2 bytes = 42GB
Optimizer states (AdamW): 21B × 8 bytes = 168GB (fp32 master + momentum + variance)
Gradients: 21B × 2 bytes = 42GB
Activations (per micro-batch): ~5-10GB per GPU (depends on seq length)
Total without parallelism: ~260GB
With parallelism (EP=4, PP=4):
Each PP stage: 6 layers → ~1/4 of model
Each EP rank: 8 experts → ~1/4 of MoE params
Optimizer sharded (ZeRO-1): split across all 16 GPUs
Per-GPU memory: ~20-30GB → comfortable on 80GB H100
Inference (4×H100, 320GB total VRAM)
Model weights (BF16): 42GB → fits on single H100 (80GB)
With tensor parallelism (TP=4): 10.5GB per GPU
KV Cache at 4K context: 24 layers × 2 × 512 × 4K × 2B = 201MB
KV Cache at 131K context: 24 layers × 2 × 512 × 131K × 2B = 6.3GB
Total (131K context): ~48GB on 4 GPUs = 12GB per GPU
Remaining VRAM: 68GB per GPU → large batch inference possible
With MXFP4 quantization: ~13GB model → fits on single 16GB GPU!
References
| Resource | Link |
|---|---|
| GPT-OSS-120B model | openai/gpt-oss-120b |
| GPT-OSS-20B model | openai/gpt-oss-20b |
| GPT-OSS paper | arxiv:2508.10925 |
| DeepSeek-V2 Lite | deepseek-ai/DeepSeek-V2-Lite |
| DeepSeek-V2 paper | arxiv:2405.04434 |
| DeepSeek-V3 paper | arxiv:2412.19437 |
| DeepSeekMoE paper | arxiv:2401.06066 |
| Moonlight (Kimi MoE) | moonshotai/Moonlight-16B-A3B-Instruct |
| Moonlight paper | arxiv:2502.16982 |
| OLMoE paper | arxiv:2409.02060 |
| Megatron-LM MoE paper | arxiv:2603.07685 |
| Megatron-LM repo | github.com/NVIDIA/Megatron-LM |
| The Stack v2 | bigcode/the-stack-v2 |
| FineWeb-Edu | HuggingFaceFW/fineweb-edu |
| OpenR1-Math | open-r1/OpenR1-Math-220k |