# 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