@tryhamster/gerbil 1.0.0-rc.8 → 1.0.0

package/docs/gpu-engine/architectures.md

@@ -0,0 +1,398 @@
+# Adding a New Model Architecture
+
+This guide walks through how to add support for a new model family to the gerbil GPU engine. The reference implementation is the Qwen2 graph generator in `src/gpu/architectures/qwen2.ts`.
+
+---
+
+## How Architecture Support Works
+
+The engine uses a **registry pattern**: HuggingFace models declare their architecture in `config.json` via the `architectures` field (e.g., `["Qwen2ForCausalLM"]`). The engine maps this string to a **graph generator function** that produces a complete `ModelGraph` (IR) from the config.
+
+Adding a new model family requires:
+1. Writing a graph generator function
+2. Registering it in the architecture registry
+3. (Optionally) writing a custom HF key mapper if the model uses non-standard weight naming
+
+No changes to the executor, kernels, or device layer are needed -- the IR is the contract.
+
+---
+
+## Prerequisites
+
+Before writing a generator, you need:
+
+1. **A model's `config.json`**: Download from HuggingFace to understand the architecture dimensions
+2. **Understanding of the model's layer structure**: Read the model's paper or reference implementation to know the computation graph
+3. **A safetensors file**: To verify weight tensor names match your key mapping
+
+---
+
+## Step-by-Step Guide
+
+### Step 1: Study the Config
+
+Download `config.json` from the target model's HuggingFace repo. Example for LLaMA-3:
+
+```json
+{
+  "architectures": ["LlamaForCausalLM"],
+  "hidden_size": 4096,
+  "intermediate_size": 14336,
+  "num_attention_heads": 32,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 500000.0,
+  "vocab_size": 128256,
+  "max_position_embeddings": 8192
+}
+```
+
+Map these to the `ModelArchConfig` fields:
+
+| Config Key | ModelArchConfig Field |
+|-----------|---------------------|
+| `hidden_size` | `hidden_size` |
+| `num_hidden_layers` | `num_layers` |
+| `num_attention_heads` | `num_heads` |
+| `num_key_value_heads` | `num_kv_heads` |
+| `intermediate_size` | `intermediate_size` |
+| `vocab_size` | `vocab_size` |
+| `max_position_embeddings` | `context_length` |
+| `rms_norm_eps` | `rms_norm_eps` |
+| `rope_theta` | `rope_base` |
+
+### Step 2: Understand the Layer Structure
+
+For LLaMA-family models, each transformer layer contains:
+
+```
+Input
+  -> RMSNorm (input_layernorm)
+  -> Q/K/V projections (3x MatMul)
+  -> RoPE (on Q and K)
+  -> GQA Attention
+  -> O projection (MatMul)
+  -> Residual Add
+  -> RMSNorm (post_attention_layernorm)
+  -> Gate projection (MatMul)     ]
+  -> Up projection (MatMul)       ] SwiGLU MLP
+  -> SiLU (on gate)               ]
+  -> Mul (gate * up)              ]
+  -> Down projection (MatMul)     ]
+  -> Residual Add
+Output
+```
+
+This is identical to Qwen2's structure, which is why the Qwen2 generator can often be reused directly for LLaMA-family models.
+
+### Step 3: Check Weight Key Names
+
+Download or inspect the safetensors file to see actual weight names. For LLaMA-3:
+
+```
+model.embed_tokens.weight
+model.layers.0.input_layernorm.weight
+model.layers.0.self_attn.q_proj.weight
+model.layers.0.self_attn.k_proj.weight
+model.layers.0.self_attn.v_proj.weight
+model.layers.0.self_attn.o_proj.weight
+model.layers.0.post_attention_layernorm.weight
+model.layers.0.mlp.gate_proj.weight
+model.layers.0.mlp.up_proj.weight
+model.layers.0.mlp.down_proj.weight
+model.norm.weight
+lm_head.weight
+```
+
+After stripping the `model.` prefix, these are already in canonical form. The default HF key mapper handles this.
+
+### Step 4: Write the Graph Generator
+
+Create a new file, e.g., `src/gpu/architectures/llama.ts`:
+
+```typescript
+import type { ModelGraph, ModelArchConfig, ModelCapabilities, OpNode, TensorDesc } from "../ir.js";
+import { CANONICAL_KEYS } from "../ir.js";
+
+export function generateLlamaGraph(rawConfig: Record<string, unknown>): ModelGraph {
+  // Extract config values
+  const hidden_size = rawConfig.hidden_size as number;
+  const num_layers = rawConfig.num_hidden_layers as number;
+  const num_heads = rawConfig.num_attention_heads as number;
+  const num_kv_heads = (rawConfig.num_key_value_heads as number) ?? num_heads;
+  const intermediate_size = rawConfig.intermediate_size as number;
+  const vocab_size = rawConfig.vocab_size as number;
+  const context_length = (rawConfig.max_position_embeddings as number) ?? 8192;
+  const rms_norm_eps = (rawConfig.rms_norm_eps as number) ?? 1e-5;
+  const rope_base = (rawConfig.rope_theta as number) ?? 10000.0;
+  const head_dim = (rawConfig.head_dim as number) ?? Math.floor(hidden_size / num_heads);
+  const kv_dim = num_kv_heads * head_dim;
+
+  const config: ModelArchConfig = {
+    hidden_size,
+    num_layers,
+    num_heads,
+    num_kv_heads,
+    head_dim,
+    intermediate_size,
+    vocab_size,
+    context_length,
+    rms_norm_eps,
+    norm_type: "rmsnorm",
+    rope_base,
+    rope_dim: head_dim,
+    kv_layout: "LHSd",
+    is_moe: false,
+    has_vision_tower: false,
+  };
+
+  // ... (same tensor/node generation as Qwen2)
+  // For LLaMA, the structure is identical to Qwen2.
+  // The only differences are config values and default parameters.
+}
+```
+
+If the model uses the same layer structure as Qwen2 (which LLaMA does), you can literally reuse `generateQwen2Graph` -- the generator reads all dimensions from the config, so it adapts automatically.
+
+### Step 5: Handle Architectural Differences
+
+For models with different layer structures, you'll need custom node generation. Common differences:
+
+#### Different MLP Structure (e.g., Phi)
+
+Phi models use `fc1`/`fc2` instead of `gate_proj`/`up_proj`/`down_proj`:
+
+```typescript
+// Phi MLP: hidden -> fc1 (with GELU) -> fc2 -> hidden
+// Instead of SwiGLU: hidden -> gate+up (with SiLU) -> down -> hidden
+
+addNode({
+  id: `${prefix}_fc1`,
+  opType: "MatMul",
+  inputs: [norm2Out, fc1Weight],
+  outputs: [fc1Out],
+  attributes: { M_tensor: norm2Out, K: hidden_size, N: intermediate_size },
+});
+
+addNode({
+  id: `${prefix}_gelu`,
+  opType: "GELU",
+  inputs: [fc1Out],
+  outputs: [geluOut],
+  attributes: { count_tensor: fc1Out },
+});
+
+addNode({
+  id: `${prefix}_fc2`,
+  opType: "MatMul",
+  inputs: [geluOut, fc2Weight],
+  outputs: [fc2Out],
+  attributes: { M_tensor: geluOut, K: intermediate_size, N: hidden_size },
+});
+```
+
+#### Different Normalization (LayerNorm vs RMSNorm)
+
+Some models (GPT-2, older architectures) use LayerNorm instead of RMSNorm. Use the `LayerNorm` op type:
+
+```typescript
+addNode({
+  id: `${prefix}_norm`,
+  opType: "LayerNorm",
+  inputs: [prevOutput, normWeight, normBias],  // LayerNorm has a bias term
+  outputs: [normOut],
+  attributes: { hidden_size, eps },
+});
+```
+
+#### Q/K/V Bias
+
+Some models (Qwen2, but not Qwen3 or LLaMA) include bias terms in the Q/K/V projections. When bias is present:
+
+```typescript
+// After MatMul for Q projection:
+addNode({
+  id: `${prefix}_q_bias`,
+  opType: "Add",
+  inputs: [qOut, qProjBias],
+  outputs: [qBiasOut],
+  attributes: { count_tensor: qOut },
+});
+```
+
+### Step 6: Handle Non-Standard Weight Names
+
+If the model uses weight names that don't follow the canonical convention after stripping `model.`, write a custom key mapper:
+
+```typescript
+// For a hypothetical model with unusual naming:
+import type { HFKeyMapper } from "../ir.js";
+
+export function createPhiKeyMapper(): HFKeyMapper {
+  return (hfKey: string): string | null => {
+    let key = hfKey;
+    if (key.startsWith("model.")) key = key.slice(6);
+
+    // Phi uses "fc1" and "fc2" instead of "gate_proj"/"up_proj"/"down_proj"
+    key = key.replace(/\.mlp\.fc1\./, ".mlp.gate_proj.");
+    key = key.replace(/\.mlp\.fc2\./, ".mlp.down_proj.");
+
+    // Phi uses "dense" instead of "o_proj"
+    key = key.replace(/\.self_attn\.dense\./, ".self_attn.o_proj.");
+
+    return key;
+  };
+}
+```
+
+### Step 7: Register the Architecture
+
+In `src/gpu/architectures/index.ts`, import and register:
+
+```typescript
+import { generateLlamaGraph } from "./llama.js";
+
+export const ARCHITECTURES: Record<string, GraphGenerator> = {
+  // Existing
+  Qwen2ForCausalLM: generateQwen2Graph,
+  Qwen3ForCausalLM: generateQwen2Graph,
+
+  // New
+  LlamaForCausalLM: generateLlamaGraph,
+  MistralForCausalLM: generateLlamaGraph,  // Same architecture as LLaMA
+};
+```
+
+### Step 8: Test
+
+```typescript
+import { generateGraph } from "./architectures/index.js";
+
+// Load a test config
+const config = {
+  architectures: ["LlamaForCausalLM"],
+  hidden_size: 4096,
+  num_hidden_layers: 32,
+  num_attention_heads: 32,
+  num_key_value_heads: 8,
+  intermediate_size: 14336,
+  vocab_size: 128256,
+  max_position_embeddings: 8192,
+  rms_norm_eps: 1e-5,
+  rope_theta: 500000.0,
+};
+
+const graph = generateGraph("LlamaForCausalLM", config);
+
+// Verify graph structure
+console.log(`Architecture: ${graph.architecture}`);
+console.log(`Nodes: ${graph.nodes.length}`);
+console.log(`Tensors: ${Object.keys(graph.tensors).length}`);
+console.log(`Execution order: ${graph.executionOrder.length} steps`);
+
+// Check that all expected weight tensors exist
+const weightTensors = Object.values(graph.tensors)
+  .filter(t => t.storage === "constant");
+console.log(`Weight tensors: ${weightTensors.length}`);
+
+// Verify execution order starts and ends correctly
+console.log(`First op: ${graph.executionOrder[0]}`);  // "embed"
+console.log(`Last op: ${graph.executionOrder.at(-1)}`);  // "lm_head"
+```
+
+---
+
+## Architecture Checklist
+
+When adding a new model family, verify each item:
+
+- [ ] Config extraction: All required fields read from `config.json` with sensible defaults
+- [ ] Architecture string: Matches exactly what HF `config.architectures[0]` contains
+- [ ] Embedding: `Embedding` node with correct vocab_size and hidden_size
+- [ ] Per-layer structure: Correct sequence of norm -> attention -> residual -> MLP -> residual
+- [ ] Norm type: Using `RMSNorm` or `LayerNorm` as appropriate
+- [ ] Q/K/V projections: Correct shapes (hidden_size -> hidden_size for Q, hidden_size -> kv_dim for K/V)
+- [ ] GQA: `num_kv_heads` correctly extracted (may differ from `num_heads`)
+- [ ] RoPE: Correct `head_dim`, `rope_base`, and `position_offset` handling
+- [ ] MLP structure: SwiGLU (gate/up/silu/mul/down) vs GELU (fc1/gelu/fc2)
+- [ ] Residual connections: Add nodes connecting skip paths correctly
+- [ ] Final norm: After all layers, before LM head
+- [ ] LM head: MatMul from hidden_size to vocab_size
+- [ ] Biases: Included where the model uses them (check safetensors for `*.bias` keys)
+- [ ] Key mapper: Weight names resolve correctly from safetensors to canonical names
+- [ ] Input/output: Graph inputs = `["input_ids"]`, outputs = `["logits"]`
+- [ ] Execution order: Topologically valid (each node's inputs are produced before it runs)
+- [ ] KV cache tensors: Created for each layer with `storage: "kv_cache"`
+
+---
+
+## Model Family Reference
+
+| Family | Architecture String | Layer Structure | MLP | Norm | Notes |
+|--------|-------------------|----------------|-----|------|-------|
+| Qwen2 | `Qwen2ForCausalLM` | Standard | SwiGLU | RMSNorm | Q/K/V bias |
+| Qwen3 | `Qwen3ForCausalLM` | Standard | SwiGLU | RMSNorm | No bias |
+| LLaMA 2 | `LlamaForCausalLM` | Standard | SwiGLU | RMSNorm | rope_theta=10000 |
+| LLaMA 3 | `LlamaForCausalLM` | Standard | SwiGLU | RMSNorm | rope_theta=500000 |
+| Mistral | `MistralForCausalLM` | Standard | SwiGLU | RMSNorm | Sliding window attention |
+| Phi-3 | `Phi3ForCausalLM` | Standard | SwiGLU | RMSNorm | Fused QKV |
+| GPT-2 | `GPT2LMHeadModel` | Different | GELU MLP | LayerNorm | Pre-norm vs post-norm |
+| SmolLM | `LlamaForCausalLM` | Standard | SwiGLU | RMSNorm | Same as LLaMA |
+
+"Standard" layer structure means: pre-norm -> attention -> residual -> post-norm -> MLP -> residual. Most modern models follow this pattern.
+
+---
+
+## Common Pitfalls
+
+### 1. Forgetting to Update prevOutput
+
+Each layer must update the `prevOutput` variable to its final residual output. If you forget, subsequent layers will read stale data:
+
+```typescript
+// WRONG: prevOutput still points to previous layer
+addNode({ id: `${prefix}_resid2`, ... outputs: [`${prefix}_resid2`], ... });
+// prevOutput is not updated!
+
+// RIGHT:
+prevOutput = `${prefix}_resid2`;
+```
+
+### 2. Shape Mismatches in GQA
+
+When `num_kv_heads < num_heads`, K and V projections output `kv_dim = num_kv_heads * head_dim`, not `hidden_size`. The attention kernel handles the head mapping internally, but the tensor shapes must be correct:
+
+```typescript
+// Q: [T, hidden_size] = [T, num_heads * head_dim]
+// K: [T, kv_dim] = [T, num_kv_heads * head_dim]   <-- NOT hidden_size
+// V: [T, kv_dim]
+```
+
+### 3. Missing Safetensors Key Assignment
+
+Weight tensors must have `safetensorsKey` set, or the model loader won't know to upload them:
+
+```typescript
+addTensor({
+  name: CANONICAL_KEYS.qProj(i),
+  shape: [hidden_size, hidden_size],
+  dtype: "f32",
+  storage: "constant",
+  safetensorsKey: CANONICAL_KEYS.qProj(i),  // Don't forget this!
+});
+```
+
+### 4. Activation Tensors with safetensorsKey
+
+Activation tensors should NOT have `safetensorsKey`. They're computed during the forward pass, not loaded from files:
+
+```typescript
+addTensor({
+  name: qOut,
+  shape: ["T", hidden_size],
+  dtype: "f32",
+  storage: "activation",
+  // NO safetensorsKey here
+});
+```