spb/forge Public MIT

1# CLAUDE.md23<!--4Author: Simon-Pierre Boucher5Contact: contact@spboucher.ai6-->78> **Instruction to Claude:** Every source file created in this project (headers, .cpp, .metal, scripts) must begin with a header comment containing:9> `Author: Simon-Pierre Boucher — contact@spboucher.ai`1011---1213## ⚠️ MANDATORY FIRST STEP: Initial Web Research Phase1415**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.1617Research the following topics thoroughly (multiple searches and full-page fetches per topic, prioritizing Apple's official documentation, WWDC session notes, and reputable engineering blogs):18191. **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`.202. **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.213. **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.224. **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.235. **Flash attention on Metal**: how MLX and llama.cpp implement fused/tiled attention on Apple GPUs; online softmax algorithm details.246. **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.257. **Current macOS/Xcode toolchain**: any recent changes to Metal 3.x/4 features relevant to compute, bfloat16 support status on M3, and known pitfalls.2627**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).2829Additionally, **whenever implementing a new kernel or hitting a Metal API uncertainty during the project, search the web again** rather than guessing.3031---3233## Project: Forge — LLM Training Framework from Scratch (C++ / Metal)3435A 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.3637### Target hardware38- Apple Silicon, primary target: **M3 Ultra, 96 GB unified memory** (80-core GPU, ~28 TFLOPs FP16, 800 GB/s bandwidth)39- Must still run correctly (slower) on any M-series Mac4041### Core design principles421. **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.432. **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.443. **Unified memory advantage**: exploit Apple Silicon's shared CPU/GPU memory — use `MTLStorageModeShared` buffers, avoid copies entirely.454. **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).465. **Hackability**: code should be readable and modifiable for research experiments (pruning masks, quantized weights, custom attention variants).4748---4950## Repository layout5152```53forge/54├── CLAUDE.md55├── CMakeLists.txt56├── configs/57│   ├── gpt-10m.json58│   ├── gpt-25m.json59│   ├── gpt-100m.json60│   └── gpt-200m.json61├── src/62│   ├── core/63│   │   ├── tensor.h / tensor.cpp        # Tensor class, shapes, strides, dtype64│   │   ├── allocator.h / allocator.cpp  # MTLBuffer pool, no per-op allocations65│   │   ├── device.h / device.cpp        # Metal device, queues, pipeline cache66│   │   └── autograd.h / autograd.cpp    # Dynamic graph, backward()67│   ├── ops/68│   │   ├── cpu/                         # Reference implementations (all ops)69│   │   └── metal/                       # Metal dispatch wrappers70│   ├── kernels/                         # .metal shader sources71│   │   ├── matmul.metal                 # naive → tiled → simdgroup_matrix72│   │   ├── softmax.metal73│   │   ├── layernorm.metal              # (or rmsnorm)74│   │   ├── attention.metal              # fused attention (flash-style)75│   │   ├── elementwise.metal            # GELU/SiLU, add, mul, fused bias+act76│   │   ├── embedding.metal77│   │   ├── cross_entropy.metal          # fused softmax + CE loss78│   │   └── adamw.metal                  # fused optimizer update79│   ├── nn/80│   │   ├── module.h                     # base Module, parameter registration81│   │   ├── linear.h / linear.cpp82│   │   ├── embedding.h83│   │   ├── attention.h                  # MHA + GQA support84│   │   ├── mlp.h85│   │   ├── transformer.h                # block + full model from config86│   │   └── config.h                     # ModelConfig struct ← JSON87│   ├── train/88│   │   ├── dataloader.h / .cpp          # mmap'd binary token files89│   │   ├── optimizer.h / .cpp           # AdamW (+ weight decay, grad clip)90│   │   ├── scheduler.h                  # warmup + cosine decay91│   │   ├── trainer.h / .cpp             # training loop, checkpointing, logging92│   │   └── checkpoint.h / .cpp          # save/load weights + optimizer state93│   ├── tokenizer/94│   │   └── bpe.h / bpe.cpp              # BPE encode/decode (load pretrained vocab)95│   └── main.cpp                         # CLI: train / generate / eval96├── tools/97│   ├── prepare_data.py                  # HF dataset → tokenized .bin (uint16)98│   └── train_tokenizer.py               # train small BPE vocab on corpus99└── tests/100    ├── test_ops.cpp                     # CPU vs Metal parity for every op101    ├── test_gradcheck.cpp               # numerical gradient checking102    └── test_overfit.cpp                 # sanity: overfit tiny batch to ~0 loss103```104105---106107## Architecture specification108109### Model (decoder-only transformer, configurable)110- Token embedding (optionally tied with output head)111- Learned positional embeddings **or** RoPE (config flag; default RoPE)112- N × TransformerBlock:113  - Pre-norm (LayerNorm or RMSNorm, config flag; default RMSNorm)114  - Multi-head attention with causal mask; support GQA via `n_kv_heads`115  - MLP: SwiGLU (default) or GELU, `d_ff` from config116  - Residual connections117- Final norm → LM head → fused softmax cross-entropy118119### Tensor class requirements120- Dtypes: `f32` (default for training), `f16`/`bf16` (compute/storage where safe), `u16`/`i32` (tokens/indices)121- Backed by `MTLBuffer` with `MTLStorageModeShared`; raw pointer accessible from CPU at all times122- Shape + strides; row-major; support views/reshape without copy123- Reference-counted buffer ownership through the allocator pool124125### Autograd requirements126- Dynamic tape: each forward op records a node with inputs + backward lambda127- `loss.backward()` walks the tape in reverse; gradients accumulate into `.grad` tensors128- `no_grad` scope for inference/generation129- Gradient accumulation across micro-batches must be supported130131### Metal execution model132- One `MTLCommandQueue`; batch many kernel dispatches per `MTLCommandBuffer`133- **Never** call `waitUntilCompleted` per-op; sync only at loss readback / logging boundaries134- Precompile all pipelines at startup; cache `MTLComputePipelineState` by kernel name + specialization constants135- Threadgroup sizes chosen per-kernel; expose via constants for tuning136137### Kernel optimization roadmap (implement in this order)1381. `matmul` naive (correctness baseline)1392. `matmul` tiled with threadgroup memory1403. `matmul` with `simdgroup_matrix` (f16 accumulate f32) — this is the perf-critical path1414. Fused `softmax` (online, single-pass)1425. `rmsnorm` / `layernorm` (parallel reduction)1436. Fused attention kernel (flash-attention style: tiled QK^T + online softmax + PV in one kernel, no materialized attention matrix)1447. Fused `bias + activation`, fused residual add1458. Fused `cross_entropy` (avoid materializing full logits softmax)1469. Fused `AdamW` update (one kernel over all params)147148### Training loop requirements149- Mixed precision: weights/master in f32, compute in f16/bf16 with loss scaling (config flag)150- AdamW with decoupled weight decay; gradient clipping by global norm151- LR schedule: linear warmup → cosine decay (configurable)152- Gradient accumulation for large effective batch sizes153- Checkpoint every N steps (weights + optimizer + RNG state + step); resumable154- Logging: step, loss, tokens/sec, LR, grad norm → stdout + CSV file155- Deterministic mode (fixed seeds) for debugging156157### Data pipeline158- `tools/prepare_data.py`: downloads HF dataset, tokenizes with trained BPE, writes `train.bin` / `val.bin` as flat uint16 token arrays159- C++ dataloader: `mmap()` the .bin file, sample random contiguous windows of `context_length + 1`, build (input, target) batches directly into shared buffers160- No Python at training time161162### CLI163```164forge train    --config configs/gpt-25m.json --data data/tinystories --out runs/exp1165forge generate --checkpoint runs/exp1/ckpt_5000 --prompt "Once upon a time" --temp 0.8 --top-k 40166forge eval     --checkpoint runs/exp1/ckpt_5000 --data data/tinystories/val.bin167```168169---170171## Research extensibility (design for these now, implement later)172- **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).173- **Quantized weights**: leave a clean seam in `Linear` forward to swap in ternary/int8 weight kernels later (BitNet-style).174- **Attention variants**: `attention.h` behind an interface so linear-attention variants can be added without touching the rest.175176---177178## Testing & validation protocol (mandatory)1791. Every Metal kernel: parity test vs CPU reference (`tests/test_ops.cpp`), max abs error ≤ 1e-4 (f32) / 1e-2 (f16)1802. Every op with parameters: numerical gradient check on small tensors (central differences, rel error ≤ 1e-3)1813. End-to-end sanity: overfit a single batch of 64 sequences to loss < 0.05 within 500 steps1824. Throughput benchmark target on M3 Ultra: ≥ 100k tokens/sec for the 25M config, ≥ 50k tokens/sec for 200M (batch ≥ 512 sequences of 512 tokens)183184## Build185- CMake ≥ 3.24, clang from Xcode toolchain, C++20186- `.metal` files compiled to a `.metallib` at build time via `xcrun metal` custom command187- `make test` runs all parity + gradcheck tests; CI-style: tests must pass before any kernel is considered done188189## Coding conventions190- C++20, no exceptions in hot paths, `snake_case` functions, `PascalCase` types191- Every file starts with the author header (see instruction at top)192- Comments explain *why*, not *what*; kernel files document their threadgroup layout and memory access pattern193- No premature abstraction: three concrete uses before generalizing194195## Suggested milestones1961. **M1 — Skeleton**: Tensor, allocator, device init, naive CPU matmul, build system1972. **M2 — CPU training**: full model forward/backward on CPU, overfit tiny batch, gradcheck green1983. **M3 — Metal core**: matmul (tiled), softmax, norm, elementwise on GPU; parity tests green1994. **M4 — Full GPU training**: train 10M model on TinyStories end-to-end, generation working2005. **M5 — Performance**: simdgroup matmul, fused attention, fused CE + AdamW; hit throughput targets2016. **M6 — Scale**: 100M–200M configs, mixed precision, resumable long runs, SmolLM-corpus pipeline202