# ✨ 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

# 🎓 Trained model — 12.2M params, one epoch, ~8 minutes

Validation loss falls monotonically with no sign of overfitting:

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×

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 allocas 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

Once shader sources were embedded in the metallib, gpudebug showed it directly:

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:

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

# 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):

⚠️ 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:

† measured before the MMA backward landed — pessimistic.

# 🏗️ Architecture

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

# 🙏 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.