Forge
LLM training from scratch in C++20 and Metal
A complete transformer training stack for Apple Silicon with zero ML dependencies — including the fused flash-attention backward kernel no major framework ships.
38.2k
tokens/s training throughput
10.8
TFLOPS GEMM f32
51.5
TFLOPS matmul2d f16 (M5)
15.2x
flash-attention backward speedup
85
CPU-Metal parity checks
~6.5k
lines of code
Overview
Forge is a complete, working transformer training stack built from nothing on Apple Silicon: tensors, autograd, Metal compute kernels, flash attention forward and backward, AdamW and Muon optimizers, a BPE tokenizer, checkpointing, and generation — all hand-written in roughly 6,500 lines of C++20, with no PyTorch, no MLX, and no ML dependencies.
As of July 2026, no major open-source framework ships a fused attention backward kernel for Metal — MLX throws NYI, llama.cpp lacks the op, PyTorch MPS and Candle are forward-only. Forge has one, and it is 15.2x faster than the naive version. Hand-tuned simdgroup_matrix GEMM reaches 10.8 TFLOPS f32, and Metal 4's matmul2d hits 51.5 TFLOPS f16 on the M5 neural accelerators.
Architecture is entirely config-driven: the same binary trains a 12M or a 205M parameter model by changing a JSON file, with GQA, RoPE, SwiGLU, MoE routing, and BitNet-style ternary quantization-aware training all selectable from config. Every op is validated against a CPU reference through 85 parity checks and numerical gradient checks.
Key Features
Fused flash attention backward
The Metal kernel no major framework ships — MLX, llama.cpp, PyTorch MPS, and Candle all lack it. Forge's runs 15.2x faster than naive.
Hand-written Metal GEMM
Naive, tiled, and simdgroup_matrix variants reaching 10.8 TFLOPS f32, plus Metal 4 matmul2d at 51.5 TFLOPS f16.
Zero ML dependencies
Pure C++20, Metal kernels, and a JSON parser. Tensors, autograd, optimizers, and tokenizer all built from nothing.
Config-driven architecture
The same binary trains 12M to 205M parameter models by editing JSON: GQA, RoPE, SwiGLU or GELU, MoE, tied embeddings.
Measured, not assumed
85 CPU-Metal parity checks, numerical gradient checks on every parameterized op, single-batch overfit, exact checkpoint resume.
The .forge weight format
A git-style model repository: content-addressed 95 MB shards, zero-copy mmap loading on unified memory, delta-only saves.
Modern training modes
Muon orthogonalized momentum, WSD schedules, int8 and BitNet-style ternary QAT, and top-k MoE with load-balance loss — all config-selected.
Streaming HF data pipeline
Streams any of 13 registered Hugging Face datasets or weighted mixtures straight to training binaries — no full-corpus downloads.
Published GPU findings
Documented compiler traps and profiling results, including a constant constexpr pitfall that cost 12x and a 4352-byte register spill.
How It Works
Core tensor layer
Shared-storage tensor views, a bucketed MTLBuffer pool allocator, a device wrapper with pipeline caching, and an autograd tape.
Dual-backend ops
CPU reference implementations and Metal dispatch live side by side; the autograd layer routes to either backend, and the parity suite keeps them bit-comparable.
Metal kernel suite
14 .metal files covering GEMM in four generations, flash attention (scalar and MMA), softmax, norms, embeddings, cross-entropy, AdamW, fake-quant, and MoE gating.
Config-driven transformer
Decoder-only model with RMSNorm or LayerNorm, SwiGLU or GELU, RoPE or learned positions, GQA, and optional top-k MoE — all assembled from JSON.
.forge model repository
Tensors page-aligned inside content-addressed shards; loading is mmap plus bytesNoCopy, so on unified memory a multi-GB model loads in milliseconds. Saves write only changed tensors.
Tech Stack
Core
GPU
Training
Data & formats
Highlights
- A 12.2M-parameter model trains in ~8 minutes to 20.27 validation perplexity and generates coherent English stories
- One MSL fix — replacing constant constexpr with an enum — took the GEMM kernel from 0.82 to 10.21 TFLOPS (12x)
- gpudebug profiling exposed a 4352-byte register spill in the fused backward kernel; splitting it recovered another 27%
- M5 neural accelerators deliver 4.9x via matmul2d — but only through f16, making mixed precision the entry condition, not a memory optimization
- f32 matmul2d output is bit-exact against the CPU reference; f16 differs by 3.8e-06
- Ships with a LaTeX paper and research notes documenting measured findings on MSL compiler behavior