@tryhamster/gerbil 1.1.2 → 1.2.0

package/docs/gpu-engine/architectures.md

@@ -1,398 +0,0 @@
-# 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
-});
-```