# CLAUDE.md

Instruction to Claude: Every source file created in this project (headers, .cpp, .metal, scripts) must begin with a header comment containing: Author: Simon-Pierre Boucher — contact@spboucher.ai

# ⚠️ MANDATORY FIRST STEP: Initial Web Research Phase

Before writing any code, Claude must perform a substantial web research session to acquire up-to-date, accurate knowledge. Metal compute programming is niche and evolves with each OS/hardware generation — training data may be outdated or incomplete. Do NOT rely on memory for Metal API details, shader syntax, or performance characteristics.

Research the following topics thoroughly (multiple searches and full-page fetches per topic, prioritizing Apple's official documentation, WWDC session notes, and reputable engineering blogs):

1. Metal-cpp: current setup, headers, memory management (NS::SharedPtr, autorelease pools in C++), how to integrate into a CMake project, compiling .metal → .metallib with xcrun metal.
2. Metal Shading Language (latest version): compute kernel syntax, threadgroup memory, simdgroup_matrix / simdgroup_multiply_accumulate API (exact types, supported shapes like 8x8, f16/f32 rules), simd shuffle/reduction intrinsics, atomics.
3. Metal compute best practices for Apple Silicon / M3 family: threadgroup sizing, occupancy, memory coalescing, MTLStorageModeShared unified memory patterns, command buffer batching, avoiding sync stalls, GPU timestamps/counters for profiling, dynamic caching on M3.
4. State-of-the-art matmul on Metal: study existing high-performance implementations — llama.cpp/ggml Metal kernels, MLX source code (github.com/ml-explore/mlx, especially mlx/backend/metal/kernels/), and any published tiling strategies for Apple GPUs. Extract concrete tile sizes and techniques that work on M-series.
5. Flash attention on Metal: how MLX and llama.cpp implement fused/tiled attention on Apple GPUs; online softmax algorithm details.
6. Training-specific references: nanoGPT / llm.c (Karpathy) for training loop structure, mixed-precision training with loss scaling, AdamW implementation details, gradient clipping, warmup+cosine schedules.
7. Current macOS/Xcode toolchain: any recent changes to Metal 3.x/4 features relevant to compute, bfloat16 support status on M3, and known pitfalls.

Deliverable of this phase: write a RESEARCH.md file at the repo root summarizing key findings — exact API signatures to use, chosen tile sizes with justification, links to all sources, and a list of pitfalls to avoid. Update it whenever new research is done during the project. Only after RESEARCH.md exists should implementation begin (Milestone M1).

Additionally, whenever implementing a new kernel or hitting a Metal API uncertainty during the project, search the web again rather than guessing.

# Project: Forge — LLM Training Framework from Scratch (C++ / Metal)

A minimal, high-performance framework for training small-to-medium LLMs from scratch on Apple Silicon, written in pure C++20 with Metal compute kernels for GPU acceleration. No PyTorch, no MLX, no external ML dependencies. The framework must be model-agnostic and configurable so the same codebase can train models of different sizes (10M → 500M+ parameters) by changing a config file only.

# Target hardware

• Apple Silicon, primary target: M3 Ultra, 96 GB unified memory (80-core GPU, ~28 TFLOPs FP16, 800 GB/s bandwidth)
• Must still run correctly (slower) on any M-series Mac

# Core design principles

1. Correctness first, speed second: every op gets a CPU reference implementation; Metal kernels are validated against CPU outputs (tolerance ≤ 1e-4) and gradients are validated with numerical gradient checking.
2. Model-agnostic: architecture defined entirely by a JSON/TOML config (n_layers, d_model, n_heads, n_kv_heads, d_ff, vocab_size, context_length, tied embeddings, etc.). No hardcoded model sizes anywhere.
3. Unified memory advantage: exploit Apple Silicon's shared CPU/GPU memory — use MTLStorageModeShared buffers, avoid copies entirely.
4. Minimal dependencies: C++20 standard library, Metal-cpp (Apple's official C++ bindings), and nothing else. JSON parsing may use a single-header library (nlohmann/json).
5. Hackability: code should be readable and modifiable for research experiments (pruning masks, quantized weights, custom attention variants).

# Repository layout

forge/
├── CLAUDE.md
├── CMakeLists.txt
├── configs/
│   ├── gpt-10m.json
│   ├── gpt-25m.json
│   ├── gpt-100m.json
│   └── gpt-200m.json
├── src/
│   ├── core/
│   │   ├── tensor.h / tensor.cpp        # Tensor class, shapes, strides, dtype
│   │   ├── allocator.h / allocator.cpp  # MTLBuffer pool, no per-op allocations
│   │   ├── device.h / device.cpp        # Metal device, queues, pipeline cache
│   │   └── autograd.h / autograd.cpp    # Dynamic graph, backward()
│   ├── ops/
│   │   ├── cpu/                         # Reference implementations (all ops)
│   │   └── metal/                       # Metal dispatch wrappers
│   ├── kernels/                         # .metal shader sources
│   │   ├── matmul.metal                 # naive → tiled → simdgroup_matrix
│   │   ├── softmax.metal
│   │   ├── layernorm.metal              # (or rmsnorm)
│   │   ├── attention.metal              # fused attention (flash-style)
│   │   ├── elementwise.metal            # GELU/SiLU, add, mul, fused bias+act
│   │   ├── embedding.metal
│   │   ├── cross_entropy.metal          # fused softmax + CE loss
│   │   └── adamw.metal                  # fused optimizer update
│   ├── nn/
│   │   ├── module.h                     # base Module, parameter registration
│   │   ├── linear.h / linear.cpp
│   │   ├── embedding.h
│   │   ├── attention.h                  # MHA + GQA support
│   │   ├── mlp.h
│   │   ├── transformer.h                # block + full model from config
│   │   └── config.h                     # ModelConfig struct ← JSON
│   ├── train/
│   │   ├── dataloader.h / .cpp          # mmap'd binary token files
│   │   ├── optimizer.h / .cpp           # AdamW (+ weight decay, grad clip)
│   │   ├── scheduler.h                  # warmup + cosine decay
│   │   ├── trainer.h / .cpp             # training loop, checkpointing, logging
│   │   └── checkpoint.h / .cpp          # save/load weights + optimizer state
│   ├── tokenizer/
│   │   └── bpe.h / bpe.cpp              # BPE encode/decode (load pretrained vocab)
│   └── main.cpp                         # CLI: train / generate / eval
├── tools/
│   ├── prepare_data.py                  # HF dataset → tokenized .bin (uint16)
│   └── train_tokenizer.py               # train small BPE vocab on corpus
└── tests/
    ├── test_ops.cpp                     # CPU vs Metal parity for every op
    ├── test_gradcheck.cpp               # numerical gradient checking
    └── test_overfit.cpp                 # sanity: overfit tiny batch to ~0 loss

# Architecture specification

# Model (decoder-only transformer, configurable)

• Token embedding (optionally tied with output head)
• Learned positional embeddings or RoPE (config flag; default RoPE)
• N × TransformerBlock:
  • Pre-norm (LayerNorm or RMSNorm, config flag; default RMSNorm)
  • Multi-head attention with causal mask; support GQA via n_kv_heads
  • MLP: SwiGLU (default) or GELU, d_ff from config
  • Residual connections
• Final norm → LM head → fused softmax cross-entropy

# Tensor class requirements

• Dtypes: f32 (default for training), f16/bf16 (compute/storage where safe), u16/i32 (tokens/indices)
• Backed by MTLBuffer with MTLStorageModeShared; raw pointer accessible from CPU at all times
• Shape + strides; row-major; support views/reshape without copy
• Reference-counted buffer ownership through the allocator pool

# Autograd requirements

• Dynamic tape: each forward op records a node with inputs + backward lambda
• loss.backward() walks the tape in reverse; gradients accumulate into .grad tensors
• no_grad scope for inference/generation
• Gradient accumulation across micro-batches must be supported

# Metal execution model

• One MTLCommandQueue; batch many kernel dispatches per MTLCommandBuffer
• Never call waitUntilCompleted per-op; sync only at loss readback / logging boundaries
• Precompile all pipelines at startup; cache MTLComputePipelineState by kernel name + specialization constants
• Threadgroup sizes chosen per-kernel; expose via constants for tuning

# Kernel optimization roadmap (implement in this order)

1. matmul naive (correctness baseline)
2. matmul tiled with threadgroup memory
3. matmul with simdgroup_matrix (f16 accumulate f32) — this is the perf-critical path
4. Fused softmax (online, single-pass)
5. rmsnorm / layernorm (parallel reduction)
6. Fused attention kernel (flash-attention style: tiled QK^T + online softmax + PV in one kernel, no materialized attention matrix)
7. Fused bias + activation, fused residual add
8. Fused cross_entropy (avoid materializing full logits softmax)
9. Fused AdamW update (one kernel over all params)

# Training loop requirements

• Mixed precision: weights/master in f32, compute in f16/bf16 with loss scaling (config flag)
• AdamW with decoupled weight decay; gradient clipping by global norm
• LR schedule: linear warmup → cosine decay (configurable)
• Gradient accumulation for large effective batch sizes
• Checkpoint every N steps (weights + optimizer + RNG state + step); resumable
• Logging: step, loss, tokens/sec, LR, grad norm → stdout + CSV file
• Deterministic mode (fixed seeds) for debugging

# Data pipeline

• tools/prepare_data.py: downloads HF dataset, tokenizes with trained BPE, writes train.bin / val.bin as flat uint16 token arrays
• C++ dataloader: mmap() the .bin file, sample random contiguous windows of context_length + 1, build (input, target) batches directly into shared buffers
• No Python at training time

# CLI

forge train    --config configs/gpt-25m.json --data data/tinystories --out runs/exp1
forge generate --checkpoint runs/exp1/ckpt_5000 --prompt "Once upon a time" --temp 0.8 --top-k 40
forge eval     --checkpoint runs/exp1/ckpt_5000 --data data/tinystories/val.bin

# Research extensibility (design for these now, implement later)

• Sparsity masks: Linear must support an optional binary mask applied to weights (for lottery-ticket / pruning experiments). Checkpoints must be able to store initial weights (for rewind-to-init experiments).
• Quantized weights: leave a clean seam in Linear forward to swap in ternary/int8 weight kernels later (BitNet-style).
• Attention variants: attention.h behind an interface so linear-attention variants can be added without touching the rest.

# Testing & validation protocol (mandatory)

1. Every Metal kernel: parity test vs CPU reference (tests/test_ops.cpp), max abs error ≤ 1e-4 (f32) / 1e-2 (f16)
2. Every op with parameters: numerical gradient check on small tensors (central differences, rel error ≤ 1e-3)
3. End-to-end sanity: overfit a single batch of 64 sequences to loss < 0.05 within 500 steps
4. Throughput benchmark target on M3 Ultra: ≥ 100k tokens/sec for the 25M config, ≥ 50k tokens/sec for 200M (batch ≥ 512 sequences of 512 tokens)

# Build

• CMake ≥ 3.24, clang from Xcode toolchain, C++20
• .metal files compiled to a .metallib at build time via xcrun metal custom command
• make test runs all parity + gradcheck tests; CI-style: tests must pass before any kernel is considered done

# Coding conventions

• C++20, no exceptions in hot paths, snake_case functions, PascalCase types
• Every file starts with the author header (see instruction at top)
• Comments explain why, not what; kernel files document their threadgroup layout and memory access pattern
• No premature abstraction: three concrete uses before generalizing

# Suggested milestones

1. M1 — Skeleton: Tensor, allocator, device init, naive CPU matmul, build system
2. M2 — CPU training: full model forward/backward on CPU, overfit tiny batch, gradcheck green
3. M3 — Metal core: matmul (tiled), softmax, norm, elementwise on GPU; parity tests green
4. M4 — Full GPU training: train 10M model on TinyStories end-to-end, generation working
5. M5 — Performance: simdgroup matmul, fused attention, fused CE + AdamW; hit throughput targets
6. M6 — Scale: 100M–200M configs, mixed precision, resumable long runs, SmolLM-corpus pipeline