spb/forge Public MIT
Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.
C++ 61.2%
C 23%
Python 7.6%
TeX 7.2%
CMake 1.1%
1<!-- Author: Simon-Pierre Boucher — contact@spboucher.ai -->23<div align="center">45# 🔥 Forge67### LLM training from scratch — pure C++20 + Metal on Apple Silicon89**No PyTorch. No MLX. No ML dependencies. Just C++, Metal kernels, and a JSON parser.**1011<br/>12131415161718<br/>192021222324252627282930<br/>3132**[📄 Paper](paper/forge.tex)** · **[🔬 Research notes](RESEARCH.md)** · **[📐 Spec](CLAUDE.md)**3334</div>3536---3738## ✨ What this is3940A complete, working transformer training stack built from nothing on Apple Silicon.41Tensors, autograd, Metal compute kernels, flash attention (forward **and** backward),42AdamW, BPE tokenizer, checkpointing, generation — all hand-written, all validated43against a CPU reference.4445Architecture is **entirely config-driven**: the same binary trains a 12M or a 205M46parameter model by changing a JSON file.4748> **Why it matters:** as of July 2026, *no* major open-source framework ships a fused49> attention **backward** kernel for Metal. MLX throws `"NYI"`. llama.cpp doesn't50> support the op. PyTorch MPS and Candle are forward-only. Forge has one, and it's51> 15× faster than the naive version.5253---5455## 🏆 Headline results5657<div align="center">5859| | Metric | Value |60|:--|:--|--:|61| 🚀 | Training throughput (12M model, ctx 512) | **38.2k tok/s** |62| ⚡ | GEMM f32, `simdgroup_matrix` | **10.8 TFLOPS** |63| 🔥 | GEMM f16 via `matmul2d` (M5 neural accelerators) | **51.5 TFLOPS** |64| 🎯 | Flash attention backward speedup | **15.2×** |65| 📉 | Validation perplexity (12M, 1 epoch TinyStories) | **20.27** |66| ✅ | CPU↔Metal parity checks | **85 passing** |6768</div>6970### 🎓 Trained model — 12.2M params, one epoch, ~8 minutes7172<div align="center">7374757677787980</div>8182Validation loss falls monotonically with no sign of overfitting:8384| step | 99 | 199 | 299 | 399 | 499 | final |85|:--|--:|--:|--:|--:|--:|--:|86| **train** | 5.14 | 4.10 | 3.48 | 3.17 | 3.07 | **2.99** |87| **val** | 5.16 | 4.10 | 3.48 | 3.20 | 3.07 | **3.01** |8889Sampled from the trained checkpoint (`temp 0.8`, `top-k 40`):9091> *Once upon a time, there was a little dog named Spot. Spot loved to hop and play in92> the yard with his toys. One day, it was very cold. Spot wanted to play, but it was93> too big for Spot to jump in the yard with the yarn. A big dog saw the yarn and94> wanted them too…*9596Coherent English, consistent characters, story structure — from a 12M-parameter model97trained for eight minutes on one machine.9899---100101## 🔬 Four findings worth your time102103These came out of measurement, and several contradicted what we expected.104105### 1️⃣ `constant constexpr` cost us 12×106107```cpp108constant constexpr uint TM = 4; // ❌ an address-space VARIABLE, not a constant109enum : uint { TM = 4 }; // ✅ a real compile-time constant110```111112In MSL, `constant` is an **address-space qualifier**, and program-scope variables must113live there. So that declaration isn't a compile-time constant — loop bounds built from114it don't unroll, `acc[i][j]` becomes dynamic indexing into opaque `simdgroup_matrix`115values, and the driver spills all 16 accumulators (256 B each) to the stack.116117**0.82 → 10.21 TFLOPS.** Before the fix, the "optimized" matrix kernel was *3× slower118than a naive one*.119120### 2️⃣ The IR can't diagnose it — benchmark instead121122`metal -S -emit-llvm` shows the same 3 `alloca`s and 2 MMA intrinsics at `-O0`, `-O2`123and `-O3` — **and after the fix that made it 12× faster**. Unrolling and fragment124promotion happen in the driver's AIR→ISA back end. Reading the IR looked like125confirmation and sent us the wrong way.126127### 3️⃣ Register pressure, not bandwidth, owns attention backward128129| attention backward, per layer | gpt-10m | gpt-25m |130|:--|--:|--:|131| `k,v,dk,dv` all in registers | 150 ms | — |132| read-only `k,v` from device | 118 ms | 611 ms |133| `dK`/`dV` split into 2 kernels | 105 ms | 549 ms |134| **tiled with `simdgroup_matrix`** | **9.66 ms** | **45.5 ms** |135| **+ `dK`/`dV` split again** | **7.05 ms** | **32.9 ms** |136137Once shader sources were embedded in the metallib, `gpudebug` showed it directly:138139| kernel | temp registers | **spilled bytes** |140|:--|--:|--:|141| flash forward, scalar | 126 | **368** |142| flash forward, MMA | 85 | **0** ✅ |143| flash backward dKV, fused | 111 | **4352** ❌ |144145That 4352-byte spill was a bug we hadn't suspected. Splitting the kernel recovered146another 27%.147148### 4️⃣ The M5 neural accelerators are worth 4.9× — and f16 is the key149150`mpp::tensor_ops::matmul2d` (Metal 4 cooperative tensors) vs. our hand-written kernel:151152| shape | `simdgroup` f32 | `simdgroup` f16 | **`matmul2d` f32** | **`matmul2d` f16** |153|:--|--:|--:|--:|--:|154| 2048³ | 9.3 T | 10.6 T | 14.9 T | **51.5 T** |155| 4096³ | 9.9 T | 11.8 T | 14.6 T | **44.3 T** |156157Verified numerically, not just timed: f32 is **bit-exact** vs the CPU reference, f16158differs by 3.8e-06, all outputs non-zero, and the transposed variants training needs159(`nt`, `tn`) are bit-exact too.160161**The negative result matters as much:** f16 buys only **+18–22%** on the162`simdgroup_matrix` path. So mixed precision looked like a memory-only feature — until163it turned out to be the *entry condition* for a 4.9× path.164165---166167## 🚀 Quick start168169```bash170# Build (macOS + Xcode CLT + CMake ≥ 3.24)171cmake -B build && cmake --build build -j172cd build && ctest --output-on-failure # parity + gradcheck + overfit + tokenizer173174# Data: download TinyStories, train a BPE vocab, tokenize to uint16 .bin175python3 tools/prepare_data.py --out data/tinystories --vocab-size 4096176177# Train · generate · eval178./build/forge train --config configs/gpt-10m-1epoch.json --data data/tinystories --out runs/exp1179./build/forge generate --checkpoint runs/exp1/ckpt_latest.bin \180 --tokenizer data/tinystories/tok4096.model \181 --prompt "Once upon a time" --temp 0.8 --top-k 40182./build/forge eval --checkpoint runs/exp1/ckpt_latest.bin --data data/tinystories/val.bin183./build/forge info --config configs/gpt-200m.json184```185186Add `--backend cpu` to run the same model through the CPU reference path — every op is187bit-comparable with the GPU path, which is what the parity suite checks.188189---190191## 📊 Benchmarks192193**GEMM** (`tests/bench_matmul.cpp`, TFLOPS, f32):194195| kernel | naive | 16×16 tiled | `simdgroup_matrix` |196|:--|--:|--:|--:|197| 4096³ | 1.29 | 2.53 | **9.49** |198| forward MLP `X·W₁ᵀ` | 1.46 | 2.53 | **10.68** |199| backward `dX = dY·W` | 1.51 | 2.54 | **10.58** |200| backward `dW = dYᵀ·X` | 0.76 | 2.15 | 4.69 ⚠️ |201202⚠️ `dW` lags — `K = B·T` is huge with few threadgroups. Split-K would fix it (not done).203204**Scale** — all configs train on one M5 Max:205206| config | params | context | tok/s |207|:--|--:|--:|--:|208| `gpt-10m` | 12.2M | 512 | **38.2k** |209| `gpt-25m` | 29.9M | 1024 | 22.1k |210| `gpt-100m` | 97.5M | 1024 | 9.7k |211| `gpt-200m` | 205.5M | 1024 | 1.8k† |212213† measured before the MMA backward landed — pessimistic.214215---216217## 🏗️ Architecture218219| path | contents |220|:--|:--|221| `src/core/` | `Tensor` (shared-storage views), `Allocator` (bucketed MTLBuffer pool), `Device` (queue + pipeline cache), autograd tape |222| `src/kernels/` | 14 `.metal` files: GEMM (naive→tiled→simdgroup→`matmul2d`), flash attention (scalar + MMA), softmax, norms, elementwise, embedding, cross-entropy, AdamW, fake-quant (QAT), MoE gating |223| `src/ops/` | `cpu/` reference impls · `metal/` dispatch + batched `Stream` · `ops.cpp` autograd layer routing to either backend |224| `src/nn/` | `Module`, `Linear` (pruning mask + int8/ternary QAT), attention (GQA + RoPE behind an interface), SwiGLU/GELU MLP or top-k MoE, `Transformer` |225| `src/train/` | mmap dataloader, AdamW or Muon, warmup+cosine or WSD schedule, trainer, resumable checkpoints |226| `src/tokenizer/` | byte-level BPE, verified identical to the Python encoder |227| `src/core/fmodel.*` | **`.forge` weight format** — Apple-native, git-style (see below) |228| `tools/` | `prepare_data.py` (TinyStories) · `prepare_hf_data.py` (any HF dataset/mixture, streamed) · `fmodel.py` (inspect / history / safetensors export) |229| `tests/` | parity · gradcheck · overfit · tokenizer · 3 benchmarks |230231**Model:** decoder-only transformer · RMSNorm or LayerNorm · SwiGLU or GELU · RoPE232(interleaved pairs) or learned positions · GQA · tied embeddings — all from config.233234**Training modes (config-selected):** optimizer `adamw` | `muon` (Newton–Schulz235orthogonalized momentum on hidden matrices, AdamW on embeddings/head) · schedule236`cosine` | `wsd` (flat plateau + 1−√ cooldown, extendable runs) · `quant`237`none` | `int8` | `ternary` (BitNet-style fake-quant each forward, STE backward,238f32 master weights) · `n_experts`/`moe_top_k` (softmax router, renormalized top-k239gates, differentiable load-balance loss; v1 computes experts densely). See240`configs/gpt-50m-{muon,deep,ternary,moe}.json`.241242**Verified:** 85 CPU↔Metal parity checks (≤1e-4, most bit-exact) · numerical gradient243checks on every parameterized op and a full transformer · single-batch overfit to244loss < 0.05 in 86 steps · exact checkpoint resume · BPE round-trip vs Python.245246**`.forge` weight format** (`forge export --checkpoint ckpt.bin --out model.forge`):247a model *repository* rather than a file. Tensors are 16 KB-page-aligned inside248content-addressed shards capped at 95 MB (GitHub-pushable), and loading is249`mmap` + `newBuffer(bytesNoCopy)` — on unified memory the file-cache pages *are*250the GPU memory, so a multi-GB model loads in milliseconds with zero copies.251Saves are git-style commits: a tiny JSON manifest per save, and **only tensors252whose content hash changed since the parent are written** — repeated exports253cost only the delta. Store as f32 (zero-copy load) or f16/bf16 (half size).254`generate`/`eval` accept a `.forge` repo directly; `tools/fmodel.py` gives255`inspect`, `log` (history), and `to-safetensors` for PyTorch/HF interop.256257**HF data pipeline** (`tools/prepare_hf_data.py`): streams any of 13 registered258Hugging Face datasets (FineWeb-Edu, DCLM, Cosmopedia, FineMath, OpenWebMath,259Wikipedia, C4, SmolTalk, …) or a weighted mixture (`--mix fineweb-edu:0.6,dclm:0.4`,260or presets like `smollm-web`) straight into `train.bin`/`val.bin` — no261full-corpus downloads, pay only for the megabytes you keep.262263---264265## 🗺️ Roadmap266267- [ ] **Mixed precision** (f16/bf16 + f32 master weights + loss scaling) — now the top268 priority, since it gates the 4.9× `matmul2d` path269- [ ] **Wire `matmul2d` into `ops::matmul`** behind runtime macOS-26 detection270- [ ] **Activation checkpointing** — activation memory, not parameter count, bounds271 model size today272- [ ] **Split-K** for the `dW = dYᵀ·X` GEMM (4.7 vs 10.6 TFLOPS)273- [ ] **KV cache** for generation (currently recomputes full context per token)274- [ ] Attention tile sweep — BQ/BK unswept; softmax leaves 96/128 threads idle275276---277278## 👤 Author279280<div align="center">281282**Simon-Pierre Boucher**283284[](mailto:contact@spboucher.ai)285[](https://github.com/spboucher-ai)286287</div>288289---290291## 🙏 Acknowledgements292293GEMM structure follows **MLX**'s STEEL kernels. Attention follows **FlashAttention-2**294as adapted to Apple GPUs by **metal-flash-attention**. Training-loop details — AdamW295with ε outside the sqrt, weight decay on rank-≥2 params only, clip folded into the296optimizer's gradient read, fused classifier writing logit gradients in place — follow297**llm.c** and **nanoGPT**. Built on Apple's **metal-cpp**.298299<div align="center">300<br/>301<sub>Built from scratch on Apple Silicon. Measured, not assumed.</sub>302</div>303