Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.
1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// CPU-vs-Metal parity for every GPU kernel (CLAUDE.md protocol #1: max abs4// error <= 1e-4 for f32), plus tensor mechanics and the CPU matmul oracle5// check. GPU ops encode into the batched Stream; comparisons happen after6// metal::sync() — the readback boundary.7#include <Foundation/Foundation.hpp>8#include <Metal/Metal.hpp>910#include "core/tensor.h"11#include "nn/transformer.h"12#include "ops/cpu/cpu_ops.h"13#include "ops/metal/metal_ops.h"14#include "ops/ops.h"1516#include <cmath>17#include <cstdio>18#include <cstdlib>19#include <random>20#include <vector>2122namespace {2324int g_failures = 0;25constexpr float kTol = 1e-4f;2627void expect(bool cond, const char* what) {28 if (cond) {29 std::printf(" ok: %s\n", what);30 } else {31 std::printf(" FAIL: %s\n", what);32 ++g_failures;33 }34}3536void expect_close(const forge::Tensor& got, const forge::Tensor& want, const char* what,37 float tol = kTol) {38 const float* pg = got.data<float>();39 const float* pw = want.data<float>();40 float m = 0.0f;41 for (int64_t i = 0; i < got.numel(); ++i) m = std::max(m, std::fabs(pg[i] - pw[i]));42 if (m <= tol) {43 std::printf(" ok: %s (max abs err %.2e)\n", what, double(m));44 } else {45 std::printf(" FAIL: %s (max abs err %.2e > %.0e)\n", what, double(m), double(tol));46 ++g_failures;47 }48}4950void fill_random(forge::Tensor& t, std::mt19937& rng, float lo = -1.0f, float hi = 1.0f) {51 std::uniform_real_distribution<float> dist(lo, hi);52 float* p = t.data<float>();53 for (int64_t i = 0; i < t.numel(); ++i) p[i] = dist(rng);54}5556// Deliberately dumb oracle: index-arithmetic triple loop with double acc.57void matmul_oracle(const forge::Tensor& a, const forge::Tensor& b, forge::Tensor& c,58 bool ta, bool tb) {59 const int64_t M = c.size(0), N = c.size(1);60 const int64_t K = ta ? a.size(0) : a.size(1);61 const int64_t lda = a.size(1), ldb = b.size(1);62 const float* A = a.data<float>();63 const float* B = b.data<float>();64 float* C = c.data<float>();65 for (int64_t i = 0; i < M; ++i)66 for (int64_t j = 0; j < N; ++j) {67 double acc = 0.0;68 for (int64_t k = 0; k < K; ++k) {69 const float av = ta ? A[k * lda + i] : A[i * lda + k];70 const float bv = tb ? B[j * ldb + k] : B[k * ldb + j];71 acc += double(av) * double(bv);72 }73 C[i * N + j] = float(acc);74 }75}7677void test_tensor_basics() {78 std::printf("tensor basics\n");79 forge::Tensor t = forge::Tensor::zeros({4, 8});80 expect(t.numel() == 32, "numel");81 expect(t.is_contiguous(), "contiguous");82 expect(t.strides()[0] == 8 && t.strides()[1] == 1, "row-major strides");8384 forge::Tensor v = t.view({8, 4});85 v.set_item(0, 42.0f);86 expect(t.item_at(0) == 42.0f, "view shares storage");8788 forge::Tensor s = t.slice0(1, 2);89 s.set_item(0, 7.0f);90 expect(t.item_at(8) == 7.0f, "slice0 shares storage at offset");9192 forge::Tensor h = forge::Tensor::full({3}, 1.5f, forge::DType::F16);93 expect(std::fabs(h.item_at(2) - 1.5f) < 1e-6f, "f16 roundtrip");94 forge::Tensor bf = forge::Tensor::full({3}, 1.5f, forge::DType::BF16);95 expect(std::fabs(bf.item_at(1) - 1.5f) < 1e-6f, "bf16 roundtrip");96}9798void test_cpu_matmul() {99 std::printf("cpu matmul vs oracle (all transpose variants)\n");100 std::mt19937 rng(1234);101 const int64_t M = 17, K = 23, N = 13;102 for (int ta = 0; ta <= 1; ++ta)103 for (int tb = 0; tb <= 1; ++tb) {104 forge::Tensor a = ta ? forge::Tensor::empty({K, M}) : forge::Tensor::empty({M, K});105 forge::Tensor b = tb ? forge::Tensor::empty({N, K}) : forge::Tensor::empty({K, N});106 forge::Tensor c = forge::Tensor::empty({M, N});107 forge::Tensor ref = forge::Tensor::empty({M, N});108 fill_random(a, rng);109 fill_random(b, rng);110 forge::cpu::matmul(a, b, c, ta, tb);111 matmul_oracle(a, b, ref, ta, tb);112 char label[64];113 std::snprintf(label, sizeof(label), "cpu matmul ta=%d tb=%d", ta, tb);114 expect_close(c, ref, label, 1e-5f);115 }116}117118void test_metal_matmul() {119 std::printf("metal matmul parity (all kernels, all transpose variants)\n");120 std::mt19937 rng(4321);121 using K_t = forge::metal::MatmulKernel;122 struct KernelCase { K_t kernel; const char* name; };123 const KernelCase kernels[] = {{K_t::Naive, "naive"},124 {K_t::Tiled, "tiled"},125 {K_t::Simdgroup, "simd"}};126 // Ragged (non-multiple of any tile dim) and exactly-tiled shapes, so the127 // simdgroup kernel's predicated and ALIGNED fast paths both get covered.128 struct Shape { int64_t M, K, N; const char* tag; };129 const Shape shapes[] = {{67, 129, 45, "ragged"}, {128, 64, 192, "aligned"}};130131 for (const auto& kc : kernels)132 for (const auto& sh : shapes)133 for (int ta = 0; ta <= 1; ++ta)134 for (int tb = 0; tb <= 1; ++tb) {135 forge::Tensor a = ta ? forge::Tensor::empty({sh.K, sh.M})136 : forge::Tensor::empty({sh.M, sh.K});137 forge::Tensor b = tb ? forge::Tensor::empty({sh.N, sh.K})138 : forge::Tensor::empty({sh.K, sh.N});139 forge::Tensor gpu = forge::Tensor::empty({sh.M, sh.N});140 forge::Tensor ref = forge::Tensor::empty({sh.M, sh.N});141 fill_random(a, rng);142 fill_random(b, rng);143 forge::cpu::matmul(a, b, ref, ta, tb);144 forge::metal::matmul(a, b, gpu, ta, tb, false, kc.kernel);145 forge::metal::sync();146 char label[80];147 std::snprintf(label, sizeof(label), "%s matmul %s ta=%d tb=%d",148 kc.name, sh.tag, ta, tb);149 expect_close(gpu, ref, label);150 }151152 // accumulate flag (ragged shape: both epilogue paths get predication)153 const int64_t M = 67, K = 129, N = 45;154 forge::Tensor a = forge::Tensor::empty({M, K});155 forge::Tensor b = forge::Tensor::empty({K, N});156 forge::Tensor acc_gpu = forge::Tensor::empty({M, N});157 forge::Tensor acc_ref = forge::Tensor::empty({M, N});158 fill_random(a, rng);159 fill_random(b, rng);160 fill_random(acc_gpu, rng);161 std::memcpy(acc_ref.raw(), acc_gpu.raw(), acc_gpu.nbytes());162 forge::cpu::matmul(a, b, acc_ref, false, false, true);163 forge::metal::matmul(a, b, acc_gpu, false, false, true, K_t::Tiled);164 forge::metal::sync();165 expect_close(acc_gpu, acc_ref, "tiled matmul accumulate=true");166167 // simdgroup accumulate takes the staged (non-fast) epilogue path168 forge::Tensor sacc_gpu = forge::Tensor::empty({M, N});169 forge::Tensor sacc_ref = forge::Tensor::empty({M, N});170 fill_random(sacc_gpu, rng);171 std::memcpy(sacc_ref.raw(), sacc_gpu.raw(), sacc_gpu.nbytes());172 forge::cpu::matmul(a, b, sacc_ref, false, false, true);173 forge::metal::matmul(a, b, sacc_gpu, false, false, true, K_t::Simdgroup);174 forge::metal::sync();175 expect_close(sacc_gpu, sacc_ref, "simd matmul accumulate=true");176}177178void test_metal_elementwise() {179 std::printf("metal elementwise parity\n");180 std::mt19937 rng(777);181 const int64_t n = 65537; // non-multiple of the threadgroup size182 forge::Tensor x = forge::Tensor::empty({n});183 forge::Tensor y = forge::Tensor::empty({n});184 fill_random(x, rng, -4.0f, 4.0f);185 fill_random(y, rng, -4.0f, 4.0f);186187 forge::Tensor gpu = forge::Tensor::empty({n});188 forge::Tensor ref = forge::Tensor::empty({n});189190 forge::cpu::add(x, y, ref);191 forge::metal::add(x, y, gpu);192 forge::metal::sync();193 expect_close(gpu, ref, "add");194195 forge::cpu::mul(x, y, ref);196 forge::metal::mul(x, y, gpu);197 forge::metal::sync();198 expect_close(gpu, ref, "mul");199200 forge::cpu::scale(x, 0.37f, ref);201 forge::metal::scale(x, 0.37f, gpu);202 forge::metal::sync();203 expect_close(gpu, ref, "scale");204205 forge::cpu::silu(x, ref);206 forge::metal::silu(x, gpu);207 forge::metal::sync();208 expect_close(gpu, ref, "silu");209210 forge::cpu::gelu(x, ref);211 forge::metal::gelu(x, gpu);212 forge::metal::sync();213 expect_close(gpu, ref, "gelu");214215 // add_bias: [N, C] + [C]216 const int64_t rows = 513, C = 127;217 forge::Tensor xb = forge::Tensor::empty({rows, C});218 forge::Tensor bias = forge::Tensor::empty({C});219 fill_random(xb, rng);220 fill_random(bias, rng);221 forge::Tensor gpub = forge::Tensor::empty({rows, C});222 forge::Tensor refb = forge::Tensor::empty({rows, C});223 forge::cpu::add_bias(xb, bias, refb);224 forge::metal::add_bias(xb, bias, gpub);225 forge::metal::sync();226 expect_close(gpub, refb, "add_bias");227}228229void test_metal_rowops() {230 std::printf("metal softmax/norm parity\n");231 std::mt19937 rng(31337);232 // Row lengths: tiny (< one simdgroup), odd, large (> threadgroup size,233 // vocab-like).234 for (int64_t C : {7LL, 63LL, 384LL, 4099LL}) {235 const int64_t rows = 129;236 forge::Tensor x = forge::Tensor::empty({rows, C});237 fill_random(x, rng, -8.0f, 8.0f);238 forge::Tensor gpu = forge::Tensor::empty({rows, C});239 forge::Tensor ref = forge::Tensor::empty({rows, C});240241 char label[64];242 forge::cpu::softmax(x, ref);243 forge::metal::softmax(x, gpu);244 forge::metal::sync();245 std::snprintf(label, sizeof(label), "softmax C=%lld", C);246 expect_close(gpu, ref, label);247248 forge::Tensor w = forge::Tensor::empty({C});249 forge::Tensor b = forge::Tensor::empty({C});250 fill_random(w, rng);251 fill_random(b, rng);252253 forge::cpu::rmsnorm(x, w, 1e-6f, ref);254 forge::metal::rmsnorm(x, w, 1e-6f, gpu);255 forge::metal::sync();256 std::snprintf(label, sizeof(label), "rmsnorm C=%lld", C);257 expect_close(gpu, ref, label);258259 forge::cpu::layernorm(x, w, b, 1e-6f, ref);260 forge::metal::layernorm(x, w, b, 1e-6f, gpu);261 forge::metal::sync();262 std::snprintf(label, sizeof(label), "layernorm C=%lld", C);263 expect_close(gpu, ref, label);264 }265}266267void test_batched_encoding() {268 std::printf("batched encoding (many dispatches, one sync)\n");269 std::mt19937 rng(55);270 const int64_t n = 4096;271 forge::Tensor x = forge::Tensor::empty({n});272 fill_random(x, rng);273 // chain of 20 dependent ops in ONE command buffer: out = ((x+x)*0.5) etc.274 forge::Tensor cur = forge::Tensor::empty({n});275 forge::metal::add(x, x, cur); // cur = 2x276 for (int i = 0; i < 19; ++i) {277 forge::Tensor next = forge::Tensor::empty({n});278 forge::metal::scale(cur, 0.9f, next);279 cur = next;280 }281 forge::metal::sync();282 // expected: 2 * 0.9^19 * x283 const float k = 2.0f * std::pow(0.9f, 19.0f);284 forge::Tensor ref = forge::Tensor::empty({n});285 forge::cpu::scale(x, k, ref);286 expect_close(cur, ref, "20-dispatch serial chain", 1e-5f);287}288289// Full-model forward+backward on CPU vs Metal: one test that exercises290// every training kernel (embedding fwd/bwd, rope, attention fwd/dq/dkv,291// matmul all variants, norms fwd/bwd, activations fwd/bwd, CE fwd+grad,292// accumulate/axpy) through the real autograd tape.293void test_backend_parity_model(const char* label, const forge::ModelConfig& cfg) {294 std::printf("backend parity: %s\n", label);295 forge::nn::Transformer model(cfg, 7);296297 std::mt19937_64 rng(21);298 forge::Tensor ids = forge::Tensor::empty({2, 6}, forge::DType::I32);299 forge::Tensor targets = forge::Tensor::empty({2 * 6}, forge::DType::I32);300 forge::cpu::fill_uniform_int(ids, 0, cfg.vocab_size, rng);301 forge::cpu::fill_uniform_int(targets, 0, cfg.vocab_size, rng);302 targets.data<int32_t>()[5] = -1; // exercise ignore_index303304 struct Saved { std::string name; forge::Tensor grad; };305 std::vector<Saved> cpu_grads;306 float cpu_loss = 0.0f;307308 {309 forge::ops::set_backend(forge::ops::Backend::CPU);310 model.zero_grad();311 forge::Var loss = model.loss(ids, targets);312 forge::Tape::get().backward(loss);313 cpu_loss = loss.value().data<float>()[0];314 std::unordered_map<const void*, bool> seen;315 for (const auto& [name, p] : model.named_parameters()) {316 if (!seen.emplace(p.id(), true).second) continue;317 forge::Tensor copy = forge::Tensor::empty(p.grad().shape());318 std::memcpy(copy.raw(), p.grad().raw(), copy.nbytes());319 cpu_grads.push_back({name, copy});320 }321 }322323 float metal_loss = 0.0f;324 {325 forge::ops::set_backend(forge::ops::Backend::Metal);326 model.zero_grad();327 forge::Var loss = model.loss(ids, targets);328 forge::Tape::get().backward(loss);329 forge::metal::sync();330 metal_loss = loss.value().data<float>()[0];331 forge::ops::set_backend(forge::ops::Backend::CPU);332 }333334 char lbl[96];335 std::snprintf(lbl, sizeof(lbl), "%s loss (cpu %.5f vs gpu %.5f)", label,336 double(cpu_loss), double(metal_loss));337 expect(std::fabs(cpu_loss - metal_loss) <= 1e-4f, lbl);338339 size_t idx = 0;340 float worst = 0.0f;341 std::string worst_name;342 std::unordered_map<const void*, bool> seen;343 for (const auto& [name, p] : model.named_parameters()) {344 if (!seen.emplace(p.id(), true).second) continue;345 const forge::Tensor& want = cpu_grads[idx++].grad;346 const float* pg = p.grad().data<float>();347 const float* pw = want.data<float>();348 for (int64_t i = 0; i < want.numel(); ++i) {349 const float e = std::fabs(pg[i] - pw[i]);350 if (e > worst) { worst = e; worst_name = name; }351 }352 }353 std::snprintf(lbl, sizeof(lbl), "%s all grads (worst %.2e @ %s)", label,354 double(worst), worst_name.c_str());355 expect(worst <= kTol, lbl);356}357358// Fused flash attention vs the CPU reference: forward output, the saved359// logsumexp, and all three input gradients. Covers MHA and GQA, causal and360// non-causal, and a head_dim on the supported list (the model-level parity361// tests above use head_dim 8, which deliberately falls back to the unfused362// kernel — so without this the fused path would go untested).363void test_flash_attention() {364 std::printf("flash attention parity (fused vs cpu reference)\n");365 std::mt19937 rng(2024);366367 struct Case { int64_t B, T, H, HKV, HD; bool causal; const char* tag; };368 const Case cases[] = {369 {2, 37, 4, 4, 64, true, "MHA causal hd64 (ragged T)"},370 {2, 64, 4, 2, 64, true, "GQA causal hd64 (2q/1kv)"},371 {1, 33, 2, 2, 32, true, "MHA causal hd32"},372 {2, 24, 2, 2, 64, false, "MHA non-causal hd64"},373 {1, 40, 1, 1, 128, true, "single head hd128"},374 };375376 for (const auto& c : cases) {377 const int64_t Cq = c.H * c.HD, Ckv = c.HKV * c.HD;378 forge::Tensor q = forge::Tensor::empty({c.B, c.T, Cq});379 forge::Tensor k = forge::Tensor::empty({c.B, c.T, Ckv});380 forge::Tensor v = forge::Tensor::empty({c.B, c.T, Ckv});381 fill_random(q, rng, -1.5f, 1.5f);382 fill_random(k, rng, -1.5f, 1.5f);383 fill_random(v, rng, -1.5f, 1.5f);384 const float scale = 1.0f / std::sqrt(float(c.HD));385386 forge::Tensor ref_out = forge::Tensor::empty({c.B, c.T, Cq});387 forge::Tensor probs = forge::Tensor::empty({c.B, c.H, c.T, c.T});388 forge::cpu::attention(q, k, v, c.H, c.HKV, c.causal, scale, ref_out, &probs);389390 forge::Tensor gpu_out = forge::Tensor::empty({c.B, c.T, Cq});391 forge::Tensor lse = forge::Tensor::empty({c.B, c.H, c.T});392 forge::metal::flash_attention(q, k, v, c.H, c.HKV, c.causal, scale, gpu_out, lse);393 forge::metal::sync();394395 char label[96];396 std::snprintf(label, sizeof(label), "%s fwd", c.tag);397 expect_close(gpu_out, ref_out, label);398399 // backward: same upstream gradient into both paths400 forge::Tensor dout = forge::Tensor::empty({c.B, c.T, Cq});401 fill_random(dout, rng);402403 forge::Tensor rdq = forge::Tensor::zeros({c.B, c.T, Cq});404 forge::Tensor rdk = forge::Tensor::zeros({c.B, c.T, Ckv});405 forge::Tensor rdv = forge::Tensor::zeros({c.B, c.T, Ckv});406 forge::cpu::attention_backward(q, k, v, probs, ref_out, dout, c.H, c.HKV, scale,407 rdq, rdk, rdv);408409 forge::Tensor gdq = forge::Tensor::zeros({c.B, c.T, Cq});410 forge::Tensor gdk = forge::Tensor::zeros({c.B, c.T, Ckv});411 forge::Tensor gdv = forge::Tensor::zeros({c.B, c.T, Ckv});412 forge::metal::flash_attention_backward(q, k, v, gpu_out, lse, dout, c.H, c.HKV,413 c.causal, scale, gdq, gdk, gdv);414 forge::metal::sync();415416 std::snprintf(label, sizeof(label), "%s dq", c.tag);417 expect_close(gdq, rdq, label);418 std::snprintf(label, sizeof(label), "%s dk", c.tag);419 expect_close(gdk, rdk, label);420 std::snprintf(label, sizeof(label), "%s dv", c.tag);421 expect_close(gdv, rdv, label);422 }423}424425void test_adamw_kernel() {426 std::printf("adamw kernel parity\n");427 std::mt19937 rng(99);428 const int64_t n = 4097;429 forge::Tensor w_gpu = forge::Tensor::empty({n});430 forge::Tensor g = forge::Tensor::empty({n});431 fill_random(w_gpu, rng);432 fill_random(g, rng);433 forge::Tensor w_cpu = forge::Tensor::empty({n});434 std::memcpy(w_cpu.raw(), w_gpu.raw(), w_gpu.nbytes());435 forge::Tensor m_gpu = forge::Tensor::zeros({n});436 forge::Tensor v_gpu = forge::Tensor::zeros({n});437 forge::Tensor m_cpu = forge::Tensor::zeros({n});438 forge::Tensor v_cpu = forge::Tensor::zeros({n});439440 const float lr = 1e-3f, b1 = 0.9f, b2 = 0.95f, eps = 1e-8f, wd = 0.1f, gs = 0.7f;441 for (int64_t t = 1; t <= 3; ++t) {442 forge::metal::adamw_step(w_gpu, g, m_gpu, v_gpu, lr, b1, b2, t, eps, wd, gs);443 forge::metal::sync();444 const float bc1 = 1.0f - std::pow(b1, float(t));445 const float bc2 = 1.0f - std::pow(b2, float(t));446 float* w = w_cpu.data<float>();447 float* m = m_cpu.data<float>();448 float* v = v_cpu.data<float>();449 const float* gp = g.data<float>();450 for (int64_t i = 0; i < n; ++i) {451 const float grad = gp[i] * gs;452 m[i] = b1 * m[i] + (1 - b1) * grad;453 v[i] = b2 * v[i] + (1 - b2) * grad * grad;454 w[i] -= lr * ((m[i] / bc1) / (std::sqrt(v[i] / bc2) + eps) + wd * w[i]);455 }456 }457 expect_close(w_gpu, w_cpu, "adamw 3 steps", 1e-5f);458459 // sumsq + sum reductions460 forge::Tensor out = forge::Tensor::empty({1});461 forge::metal::sumsq(g, out);462 forge::metal::sync();463 double want = 0.0;464 for (int64_t i = 0; i < n; ++i) {465 const double x = double(g.data<float>()[i]);466 want += x * x;467 }468 expect(std::fabs(out.data<float>()[0] - float(want)) <= 1e-2f, "sumsq");469 forge::metal::sum(g, out, 0.5f);470 forge::metal::sync();471 double s = 0.0;472 for (int64_t i = 0; i < n; ++i) s += double(g.data<float>()[i]);473 expect(std::fabs(out.data<float>()[0] - float(s * 0.5)) <= 1e-2f, "sum");474}475476} // namespace477478int main() {479 NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();480481 test_tensor_basics();482 test_cpu_matmul();483 test_metal_matmul();484 test_metal_elementwise();485 test_metal_rowops();486 test_batched_encoding();487488 {489 forge::ModelConfig cfg;490 cfg.n_layers = 2; cfg.d_model = 16; cfg.n_heads = 2; cfg.n_kv_heads = 1;491 cfg.d_ff = 24; cfg.vocab_size = 11; cfg.context_length = 8;492 cfg.tied_embeddings = true;493 test_backend_parity_model("rmsnorm/swiglu/rope/gqa/tied [unfused hd8]", cfg);494495 // head_dim 64 -> the whole model runs through the fused attention path496 forge::ModelConfig fcfg;497 fcfg.n_layers = 2; fcfg.d_model = 128; fcfg.n_heads = 2; fcfg.n_kv_heads = 1;498 fcfg.d_ff = 96; fcfg.vocab_size = 11; fcfg.context_length = 8;499 fcfg.tied_embeddings = true;500 test_backend_parity_model("full model via fused attention [hd64]", fcfg);501502 cfg.norm = "layernorm"; cfg.activation = "gelu"; cfg.use_rope = false;503 cfg.n_kv_heads = 2; cfg.tied_embeddings = false;504 test_backend_parity_model("layernorm/gelu/pos-emb/mha", cfg);505 }506 test_flash_attention();507 test_adamw_kernel();508509 pool->drain();510511 if (g_failures) {512 std::printf("\n%d test(s) FAILED\n", g_failures);513 return 1;514 }515 std::printf("\nall tests passed\n");516 return 0;517}518