# 🔥 Forge
### LLM training from scratch — pure C++20 + Metal on Apple Silicon
**No PyTorch. No MLX. No ML dependencies. Just C++, Metal kernels, and a JSON parser.**
**[📄 Paper](paper/forge.tex)** · **[🔬 Research notes](RESEARCH.md)** · **[📐 Spec](CLAUDE.md)**
---
## ✨ What this is
A complete, working transformer training stack built from nothing on Apple Silicon.
Tensors, autograd, Metal compute kernels, flash attention (forward **and** backward),
AdamW, BPE tokenizer, checkpointing, generation — all hand-written, all validated
against a CPU reference.
Architecture is **entirely config-driven**: the same binary trains a 12M or a 205M
parameter model by changing a JSON file.
> **Why it matters:** as of July 2026, *no* major open-source framework ships a fused
> attention **backward** kernel for Metal. MLX throws `"NYI"`. llama.cpp doesn't
> support the op. PyTorch MPS and Candle are forward-only. Forge has one, and it's
> 15× faster than the naive version.
---
## 🏆 Headline results
| | Metric | Value |
|:--|:--|--:|
| 🚀 | Training throughput (12M model, ctx 512) | **38.2k tok/s** |
| ⚡ | GEMM f32, `simdgroup_matrix` | **10.8 TFLOPS** |
| 🔥 | GEMM f16 via `matmul2d` (M5 neural accelerators) | **51.5 TFLOPS** |
| 🎯 | Flash attention backward speedup | **15.2×** |
| 📉 | Validation perplexity (12M, 1 epoch TinyStories) | **20.27** |
| ✅ | CPU↔Metal parity checks | **85 passing** |
### 🎓 Trained model — 12.2M params, one epoch, ~8 minutes
Validation loss falls monotonically with no sign of overfitting:
| step | 99 | 199 | 299 | 399 | 499 | final |
|:--|--:|--:|--:|--:|--:|--:|
| **train** | 5.14 | 4.10 | 3.48 | 3.17 | 3.07 | **2.99** |
| **val** | 5.16 | 4.10 | 3.48 | 3.20 | 3.07 | **3.01** |
Sampled from the trained checkpoint (`temp 0.8`, `top-k 40`):
> *Once upon a time, there was a little dog named Spot. Spot loved to hop and play in
> the yard with his toys. One day, it was very cold. Spot wanted to play, but it was
> too big for Spot to jump in the yard with the yarn. A big dog saw the yarn and
> wanted them too…*
Coherent English, consistent characters, story structure — from a 12M-parameter model
trained for eight minutes on one machine.
---
## 🔬 Four findings worth your time
These came out of measurement, and several contradicted what we expected.
### 1️⃣ `constant constexpr` cost us 12×
```cpp
constant constexpr uint TM = 4; // ❌ an address-space VARIABLE, not a constant
enum : uint { TM = 4 }; // ✅ a real compile-time constant
```
In MSL, `constant` is an **address-space qualifier**, and program-scope variables must
live there. So that declaration isn't a compile-time constant — loop bounds built from
it don't unroll, `acc[i][j]` becomes dynamic indexing into opaque `simdgroup_matrix`
values, and the driver spills all 16 accumulators (256 B each) to the stack.
**0.82 → 10.21 TFLOPS.** Before the fix, the "optimized" matrix kernel was *3× slower
than a naive one*.
### 2️⃣ The IR can't diagnose it — benchmark instead
`metal -S -emit-llvm` shows the same 3 `alloca`s and 2 MMA intrinsics at `-O0`, `-O2`
and `-O3` — **and after the fix that made it 12× faster**. Unrolling and fragment
promotion happen in the driver's AIR→ISA back end. Reading the IR looked like
confirmation and sent us the wrong way.
### 3️⃣ Register pressure, not bandwidth, owns attention backward
| attention backward, per layer | gpt-10m | gpt-25m |
|:--|--:|--:|
| `k,v,dk,dv` all in registers | 150 ms | — |
| read-only `k,v` from device | 118 ms | 611 ms |
| `dK`/`dV` split into 2 kernels | 105 ms | 549 ms |
| **tiled with `simdgroup_matrix`** | **9.66 ms** | **45.5 ms** |
| **+ `dK`/`dV` split again** | **7.05 ms** | **32.9 ms** |
Once shader sources were embedded in the metallib, `gpudebug` showed it directly:
| kernel | temp registers | **spilled bytes** |
|:--|--:|--:|
| flash forward, scalar | 126 | **368** |
| flash forward, MMA | 85 | **0** ✅ |
| flash backward dKV, fused | 111 | **4352** ❌ |
That 4352-byte spill was a bug we hadn't suspected. Splitting the kernel recovered
another 27%.
### 4️⃣ The M5 neural accelerators are worth 4.9× — and f16 is the key
`mpp::tensor_ops::matmul2d` (Metal 4 cooperative tensors) vs. our hand-written kernel:
| shape | `simdgroup` f32 | `simdgroup` f16 | **`matmul2d` f32** | **`matmul2d` f16** |
|:--|--:|--:|--:|--:|
| 2048³ | 9.3 T | 10.6 T | 14.9 T | **51.5 T** |
| 4096³ | 9.9 T | 11.8 T | 14.6 T | **44.3 T** |
Verified numerically, not just timed: f32 is **bit-exact** vs the CPU reference, f16
differs by 3.8e-06, all outputs non-zero, and the transposed variants training needs
(`nt`, `tn`) are bit-exact too.
**The negative result matters as much:** f16 buys only **+18–22%** on the
`simdgroup_matrix` path. So mixed precision looked like a memory-only feature — until
it turned out to be the *entry condition* for a 4.9× path.
---
## 🚀 Quick start
```bash
# Build (macOS + Xcode CLT + CMake ≥ 3.24)
cmake -B build && cmake --build build -j
cd build && ctest --output-on-failure # parity + gradcheck + overfit + tokenizer
# Data: download TinyStories, train a BPE vocab, tokenize to uint16 .bin
python3 tools/prepare_data.py --out data/tinystories --vocab-size 4096
# Train · generate · eval
./build/forge train --config configs/gpt-10m-1epoch.json --data data/tinystories --out runs/exp1
./build/forge generate --checkpoint runs/exp1/ckpt_latest.bin \
--tokenizer data/tinystories/tok4096.model \
--prompt "Once upon a time" --temp 0.8 --top-k 40
./build/forge eval --checkpoint runs/exp1/ckpt_latest.bin --data data/tinystories/val.bin
./build/forge info --config configs/gpt-200m.json
```
Add `--backend cpu` to run the same model through the CPU reference path — every op is
bit-comparable with the GPU path, which is what the parity suite checks.
---
## 📊 Benchmarks
**GEMM** (`tests/bench_matmul.cpp`, TFLOPS, f32):
| kernel | naive | 16×16 tiled | `simdgroup_matrix` |
|:--|--:|--:|--:|
| 4096³ | 1.29 | 2.53 | **9.49** |
| forward MLP `X·W₁ᵀ` | 1.46 | 2.53 | **10.68** |
| backward `dX = dY·W` | 1.51 | 2.54 | **10.58** |
| backward `dW = dYᵀ·X` | 0.76 | 2.15 | 4.69 ⚠️ |
⚠️ `dW` lags — `K = B·T` is huge with few threadgroups. Split-K would fix it (not done).
**Scale** — all configs train on one M5 Max:
| config | params | context | tok/s |
|:--|--:|--:|--:|
| `gpt-10m` | 12.2M | 512 | **38.2k** |
| `gpt-25m` | 29.9M | 1024 | 22.1k |
| `gpt-100m` | 97.5M | 1024 | 9.7k |
| `gpt-200m` | 205.5M | 1024 | 1.8k† |
† measured before the MMA backward landed — pessimistic.
---
## 🏗️ Architecture
| path | contents |
|:--|:--|
| `src/core/` | `Tensor` (shared-storage views), `Allocator` (bucketed MTLBuffer pool), `Device` (queue + pipeline cache), autograd tape |
| `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 |
| `src/ops/` | `cpu/` reference impls · `metal/` dispatch + batched `Stream` · `ops.cpp` autograd layer routing to either backend |
| `src/nn/` | `Module`, `Linear` (pruning mask + int8/ternary QAT), attention (GQA + RoPE behind an interface), SwiGLU/GELU MLP or top-k MoE, `Transformer` |
| `src/train/` | mmap dataloader, AdamW or Muon, warmup+cosine or WSD schedule, trainer, resumable checkpoints |
| `src/tokenizer/` | byte-level BPE, verified identical to the Python encoder |
| `src/core/fmodel.*` | **`.forge` weight format** — Apple-native, git-style (see below) |
| `tools/` | `prepare_data.py` (TinyStories) · `prepare_hf_data.py` (any HF dataset/mixture, streamed) · `fmodel.py` (inspect / history / safetensors export) |
| `tests/` | parity · gradcheck · overfit · tokenizer · 3 benchmarks |
**Model:** decoder-only transformer · RMSNorm or LayerNorm · SwiGLU or GELU · RoPE
(interleaved pairs) or learned positions · GQA · tied embeddings — all from config.
**Training modes (config-selected):** optimizer `adamw` | `muon` (Newton–Schulz
orthogonalized momentum on hidden matrices, AdamW on embeddings/head) · schedule
`cosine` | `wsd` (flat plateau + 1−√ cooldown, extendable runs) · `quant`
`none` | `int8` | `ternary` (BitNet-style fake-quant each forward, STE backward,
f32 master weights) · `n_experts`/`moe_top_k` (softmax router, renormalized top-k
gates, differentiable load-balance loss; v1 computes experts densely). See
`configs/gpt-50m-{muon,deep,ternary,moe}.json`.
**Verified:** 85 CPU↔Metal parity checks (≤1e-4, most bit-exact) · numerical gradient
checks on every parameterized op and a full transformer · single-batch overfit to
loss < 0.05 in 86 steps · exact checkpoint resume · BPE round-trip vs Python.
**`.forge` weight format** (`forge export --checkpoint ckpt.bin --out model.forge`):
a model *repository* rather than a file. Tensors are 16 KB-page-aligned inside
content-addressed shards capped at 95 MB (GitHub-pushable), and loading is
`mmap` + `newBuffer(bytesNoCopy)` — on unified memory the file-cache pages *are*
the GPU memory, so a multi-GB model loads in milliseconds with zero copies.
Saves are git-style commits: a tiny JSON manifest per save, and **only tensors
whose content hash changed since the parent are written** — repeated exports
cost only the delta. Store as f32 (zero-copy load) or f16/bf16 (half size).
`generate`/`eval` accept a `.forge` repo directly; `tools/fmodel.py` gives
`inspect`, `log` (history), and `to-safetensors` for PyTorch/HF interop.
**HF data pipeline** (`tools/prepare_hf_data.py`): streams any of 13 registered
Hugging Face datasets (FineWeb-Edu, DCLM, Cosmopedia, FineMath, OpenWebMath,
Wikipedia, C4, SmolTalk, …) or a weighted mixture (`--mix fineweb-edu:0.6,dclm:0.4`,
or presets like `smollm-web`) straight into `train.bin`/`val.bin` — no
full-corpus downloads, pay only for the megabytes you keep.
---
## 🗺️ Roadmap
- [ ] **Mixed precision** (f16/bf16 + f32 master weights + loss scaling) — now the top
priority, since it gates the 4.9× `matmul2d` path
- [ ] **Wire `matmul2d` into `ops::matmul`** behind runtime macOS-26 detection
- [ ] **Activation checkpointing** — activation memory, not parameter count, bounds
model size today
- [ ] **Split-K** for the `dW = dYᵀ·X` GEMM (4.7 vs 10.6 TFLOPS)
- [ ] **KV cache** for generation (currently recomputes full context per token)
- [ ] Attention tile sweep — BQ/BK unswept; softmax leaves 96/128 threads idle
---
## 👤 Author
**Simon-Pierre Boucher**
[](mailto:contact@spboucher.ai)
[](https://github.com/spboucher-ai)
---
## 🙏 Acknowledgements
GEMM structure follows **MLX**'s STEEL kernels. Attention follows **FlashAttention-2**
as adapted to Apple GPUs by **metal-flash-attention**. Training-loop details — AdamW
with ε outside the sqrt, weight decay on rank-≥2 params only, clip folded into the
optimizer's gradient read, fused classifier writing logit gradients in place — follow
**llm.c** and **nanoGPT**. Built on Apple's **metal-cpp**.
Built from scratch on Apple Silicon. Measured, not assumed.