docs/ARCHITECTURE_GUIDE.md

# Architecture Deep-Dive Guide

## Table of Contents
1. [What is a Mixture of Experts (MoE)?](#what-is-moe)
2. [How Experts Are Defined](#how-experts-are-defined)
3. [How the Router Works](#how-the-router-works)
4. [Two Architecture Options](#two-architecture-options)
5. [Key Design Decisions Explained](#key-design-decisions)
6. [VRAM and Compute Analysis](#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:

1. **Structurally identical** — each expert is a small SwiGLU FFN:
```python
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))
```

2. **Randomly initialized** — all experts start the same (with random weights)

3. **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:

```python
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](https://hf.co/openai/gpt-oss-120b) |
| GPT-OSS-20B model | [openai/gpt-oss-20b](https://hf.co/openai/gpt-oss-20b) |
| GPT-OSS paper | [arxiv:2508.10925](https://arxiv.org/abs/2508.10925) |
| DeepSeek-V2 Lite | [deepseek-ai/DeepSeek-V2-Lite](https://hf.co/deepseek-ai/DeepSeek-V2-Lite) |
| DeepSeek-V2 paper | [arxiv:2405.04434](https://arxiv.org/abs/2405.04434) |
| DeepSeek-V3 paper | [arxiv:2412.19437](https://arxiv.org/abs/2412.19437) |
| DeepSeekMoE paper | [arxiv:2401.06066](https://arxiv.org/abs/2401.06066) |
| Moonlight (Kimi MoE) | [moonshotai/Moonlight-16B-A3B-Instruct](https://hf.co/moonshotai/Moonlight-16B-A3B-Instruct) |
| Moonlight paper | [arxiv:2502.16982](https://arxiv.org/abs/2502.16982) |
| OLMoE paper | [arxiv:2409.02060](https://arxiv.org/abs/2409.02060) |
| Megatron-LM MoE paper | [arxiv:2603.07685](https://arxiv.org/abs/2603.07685) |
| Megatron-LM repo | [github.com/NVIDIA/Megatron-LM](https://github.com/NVIDIA/Megatron-LM) |
| The Stack v2 | [bigcode/the-stack-v2](https://hf.co/datasets/bigcode/the-stack-v2) |
| FineWeb-Edu | [HuggingFaceFW/fineweb-edu](https://hf.co/datasets/HuggingFaceFW/fineweb-edu) |
| OpenR1-Math | [open-r1/OpenR1-Math-220k](https://hf.co/datasets/open-r1/OpenR1-Math-220k) |