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.ai2#include "ops/ops.h"34#include "ops/cpu/cpu_ops.h"5#include "ops/metal/metal_ops.h"67#include <cmath>89namespace forge::ops {1011namespace {1213Backend g_backend = Backend::CPU;1415bool gpu() { return g_backend == Backend::Metal; }1617bool grad_needed(std::initializer_list<const Var*> inputs) {18 if (!Tape::get().enabled()) return false;19 for (const Var* v : inputs)20 if (v->requires_grad()) return true;21 return false;22}2324// dst += src (backend-routed)25void accumulate(Tensor& dst, const Tensor& src) {26 if (gpu()) {27 metal::accumulate(dst, src);28 return;29 }30 float* pd = dst.data<float>();31 const float* ps = src.data<float>();32 for (int64_t i = 0; i < dst.numel(); ++i) pd[i] += ps[i];33}3435// dst += src * s, s known on CPU36void axpy_const(Tensor& dst, const Tensor& src, float s) {37 if (gpu()) {38 Tensor tmp = Tensor::empty(src.shape());39 metal::scale(src, s, tmp);40 metal::accumulate(dst, tmp);41 return;42 }43 float* pd = dst.data<float>();44 const float* ps = src.data<float>();45 for (int64_t i = 0; i < dst.numel(); ++i) pd[i] += ps[i] * s;46}4748// dst += src * s[0], s produced on-GPU (never read on CPU)49void axpy_tensor(Tensor& dst, const Tensor& src, const Tensor& s) {50 if (gpu()) {51 metal::axpy(dst, src, s);52 return;53 }54 const float sv = s.data<float>()[0];55 float* pd = dst.data<float>();56 const float* ps = src.data<float>();57 for (int64_t i = 0; i < dst.numel(); ++i) pd[i] += ps[i] * sv;58}5960} // namespace6162void set_backend(Backend b) { g_backend = b; }63Backend backend() { return g_backend; }6465Var matmul(const Var& a, const Var& b, bool ta, bool tb) {66 const int64_t M = ta ? a.value().size(1) : a.value().size(0);67 const int64_t N = tb ? b.value().size(0) : b.value().size(1);68 Tensor out = Tensor::empty({M, N});69 if (gpu()) metal::matmul(a.value(), b.value(), out, ta, tb);70 else cpu::matmul(a.value(), b.value(), out, ta, tb);7172 const bool needs = grad_needed({&a, &b});73 Var result(std::move(out), needs);74 if (needs) {75 Tape::get().record([a, b, result, ta, tb]() {76 const Tensor& dc = result.grad();77 if (a.requires_grad()) {78 Tensor& da = a.grad();79 if (gpu()) {80 if (!ta) metal::matmul(dc, b.value(), da, false, !tb, true);81 else metal::matmul(b.value(), dc, da, tb, true, true);82 } else {83 if (!ta) cpu::matmul(dc, b.value(), da, false, !tb, true);84 else cpu::matmul(b.value(), dc, da, tb, true, true);85 }86 }87 if (b.requires_grad()) {88 Tensor& db = b.grad();89 if (gpu()) {90 if (!tb) metal::matmul(a.value(), dc, db, !ta, false, true);91 else metal::matmul(dc, a.value(), db, true, ta, true);92 } else {93 if (!tb) cpu::matmul(a.value(), dc, db, !ta, false, true);94 else cpu::matmul(dc, a.value(), db, true, ta, true);95 }96 }97 });98 }99 return result;100}101102Var add(const Var& a, const Var& b) {103 Tensor out = Tensor::empty(a.value().shape());104 if (gpu()) metal::add(a.value(), b.value(), out);105 else cpu::add(a.value(), b.value(), out);106107 const bool needs = grad_needed({&a, &b});108 Var result(std::move(out), needs);109 if (needs) {110 Tape::get().record([a, b, result]() {111 if (a.requires_grad()) accumulate(a.grad(), result.grad());112 if (b.requires_grad()) accumulate(b.grad(), result.grad());113 });114 }115 return result;116}117118Var add_bias(const Var& x, const Var& bias) {119 Tensor out = Tensor::empty(x.value().shape());120 if (gpu()) metal::add_bias(x.value(), bias.value(), out);121 else cpu::add_bias(x.value(), bias.value(), out);122123 const bool needs = grad_needed({&x, &bias});124 Var result(std::move(out), needs);125 if (needs) {126 Tape::get().record([x, bias, result]() {127 if (x.requires_grad()) accumulate(x.grad(), result.grad());128 if (bias.requires_grad()) {129 if (gpu()) metal::add_bias_backward(result.grad(), bias.grad());130 else cpu::add_bias_backward(result.grad(), bias.grad());131 }132 });133 }134 return result;135}136137Var mul(const Var& a, const Var& b) {138 Tensor out = Tensor::empty(a.value().shape());139 if (gpu()) metal::mul(a.value(), b.value(), out);140 else cpu::mul(a.value(), b.value(), out);141142 const bool needs = grad_needed({&a, &b});143 Var result(std::move(out), needs);144 if (needs) {145 Tape::get().record([a, b, result]() {146 const Tensor& dout = result.grad();147 if (a.requires_grad()) {148 Tensor tmp = Tensor::empty(dout.shape());149 if (gpu()) metal::mul(dout, b.value(), tmp);150 else cpu::mul(dout, b.value(), tmp);151 accumulate(a.grad(), tmp);152 }153 if (b.requires_grad()) {154 Tensor tmp = Tensor::empty(dout.shape());155 if (gpu()) metal::mul(dout, a.value(), tmp);156 else cpu::mul(dout, a.value(), tmp);157 accumulate(b.grad(), tmp);158 }159 });160 }161 return result;162}163164Var scale(const Var& a, float s) {165 Tensor out = Tensor::empty(a.value().shape());166 if (gpu()) metal::scale(a.value(), s, out);167 else cpu::scale(a.value(), s, out);168169 const bool needs = grad_needed({&a});170 Var result(std::move(out), needs);171 if (needs) {172 Tape::get().record([a, result, s]() {173 if (a.requires_grad()) axpy_const(a.grad(), result.grad(), s);174 });175 }176 return result;177}178179Var silu(const Var& x) {180 Tensor out = Tensor::empty(x.value().shape());181 if (gpu()) metal::silu(x.value(), out);182 else cpu::silu(x.value(), out);183184 const bool needs = grad_needed({&x});185 Var result(std::move(out), needs);186 if (needs) {187 Tape::get().record([x, result]() {188 if (!x.requires_grad()) return;189 if (gpu()) metal::silu_backward(x.value(), result.grad(), x.grad());190 else cpu::silu_backward(x.value(), result.grad(), x.grad());191 });192 }193 return result;194}195196Var gelu(const Var& x) {197 Tensor out = Tensor::empty(x.value().shape());198 if (gpu()) metal::gelu(x.value(), out);199 else cpu::gelu(x.value(), out);200201 const bool needs = grad_needed({&x});202 Var result(std::move(out), needs);203 if (needs) {204 Tape::get().record([x, result]() {205 if (!x.requires_grad()) return;206 if (gpu()) metal::gelu_backward(x.value(), result.grad(), x.grad());207 else cpu::gelu_backward(x.value(), result.grad(), x.grad());208 });209 }210 return result;211}212213Var relu2(const Var& x) {214 Tensor out = Tensor::empty(x.value().shape());215 if (gpu()) metal::relu2(x.value(), out);216 else cpu::relu2(x.value(), out);217218 const bool needs = grad_needed({&x});219 Var result(std::move(out), needs);220 if (needs) {221 Tape::get().record([x, result]() {222 if (!x.requires_grad()) return;223 if (gpu()) metal::relu2_backward(x.value(), result.grad(), x.grad());224 else cpu::relu2_backward(x.value(), result.grad(), x.grad());225 });226 }227 return result;228}229230Var softcap(const Var& x, float cap) {231 Tensor out = Tensor::empty(x.value().shape());232 if (gpu()) metal::softcap(x.value(), cap, out);233 else cpu::softcap(x.value(), cap, out);234235 const bool needs = grad_needed({&x});236 Var result(std::move(out), needs);237 if (needs) {238 Tape::get().record([x, cap, result]() {239 if (!x.requires_grad()) return;240 if (gpu()) metal::softcap_backward(result.value(), result.grad(), cap, x.grad());241 else cpu::softcap_backward(result.value(), result.grad(), cap, x.grad());242 });243 }244 return result;245}246247Var softmax(const Var& x) {248 Tensor out = Tensor::empty(x.value().shape());249 if (gpu()) metal::softmax(x.value(), out);250 else cpu::softmax(x.value(), out);251252 const bool needs = grad_needed({&x});253 Var result(std::move(out), needs);254 if (needs) {255 Tape::get().record([x, result]() {256 if (!x.requires_grad()) return;257 if (gpu()) metal::softmax_backward(result.value(), result.grad(), x.grad());258 else cpu::softmax_backward(result.value(), result.grad(), x.grad());259 });260 }261 return result;262}263264Var fake_quant(const Var& w, QuantMode mode) {265 if (mode == QuantMode::None) return w;266 const int m = mode == QuantMode::Int8 ? 0 : 1;267 Tensor out = Tensor::empty(w.value().shape());268 if (gpu()) metal::fake_quant(w.value(), m, out);269 else cpu::fake_quant(w.value(), m, out);270271 const bool needs = grad_needed({&w});272 Var result(std::move(out), needs);273 if (needs) {274 // Straight-through estimator: d(quant(w))/dw ≈ I.275 Tape::get().record([w, result]() {276 if (w.requires_grad()) accumulate(w.grad(), result.grad());277 });278 }279 return result;280}281282Var sigmoid(const Var& x) {283 Tensor out = Tensor::empty(x.value().shape());284 if (gpu()) metal::sigmoid(x.value(), out);285 else cpu::sigmoid(x.value(), out);286287 const bool needs = grad_needed({&x});288 Var result(std::move(out), needs);289 if (needs) {290 Tape::get().record([x, result]() {291 if (!x.requires_grad()) return;292 if (gpu()) metal::sigmoid_backward(result.value(), result.grad(), x.grad());293 else cpu::sigmoid_backward(result.value(), result.grad(), x.grad());294 });295 }296 return result;297}298299Var topk_renorm(const Var& probs, const Tensor& bias, int64_t k, bool norm) {300 Tensor out = Tensor::empty(probs.value().shape());301 if (gpu()) metal::topk_renorm(probs.value(), bias, k, norm, out);302 else cpu::topk_renorm(probs.value(), bias, k, norm, out);303304 const bool needs = grad_needed({&probs});305 Var result(std::move(out), needs);306 if (needs) {307 Tape::get().record([probs, bias, k, norm, result]() {308 if (!probs.requires_grad()) return;309 if (gpu())310 metal::topk_renorm_backward(probs.value(), bias, result.grad(), k,311 norm, probs.grad());312 else313 cpu::topk_renorm_backward(probs.value(), bias, result.grad(), k,314 norm, probs.grad());315 });316 }317 return result;318}319320Var row_scale(const Var& x, const Var& gates, int64_t e) {321 Tensor out = Tensor::empty(x.value().shape());322 if (gpu()) metal::row_scale(x.value(), gates.value(), e, out);323 else cpu::row_scale(x.value(), gates.value(), e, out);324325 const bool needs = grad_needed({&x, &gates});326 Var result(std::move(out), needs);327 if (needs) {328 Tape::get().record([x, gates, e, result]() {329 const Tensor& dout = result.grad();330 if (x.requires_grad()) {331 if (gpu()) metal::row_scale_accumulate(dout, gates.value(), e, x.grad());332 else cpu::row_scale_accumulate(dout, gates.value(), e, x.grad());333 }334 if (gates.requires_grad()) {335 if (gpu())336 metal::row_scale_gate_backward(dout, x.value(), e, gates.grad());337 else338 cpu::row_scale_gate_backward(dout, x.value(), e, gates.grad());339 }340 });341 }342 return result;343}344345Var rmsnorm(const Var& x, const Var& w, float eps) {346 Tensor out = Tensor::empty(x.value().shape());347 if (gpu()) metal::rmsnorm(x.value(), w.value(), eps, out);348 else cpu::rmsnorm(x.value(), w.value(), eps, out);349350 const bool needs = grad_needed({&x, &w});351 Var result(std::move(out), needs);352 if (needs) {353 Tape::get().record([x, w, eps, result]() {354 Tensor dx_scratch, dw_scratch;355 Tensor& dx = x.requires_grad() ? x.grad()356 : (dx_scratch = Tensor::zeros(x.value().shape()));357 Tensor& dw = w.requires_grad() ? w.grad()358 : (dw_scratch = Tensor::zeros(w.value().shape()));359 if (gpu()) metal::rmsnorm_backward(x.value(), w.value(), eps, result.grad(), dx, dw);360 else cpu::rmsnorm_backward(x.value(), w.value(), eps, result.grad(), dx, dw);361 });362 }363 return result;364}365366Var layernorm(const Var& x, const Var& w, const Var& b, float eps) {367 Tensor out = Tensor::empty(x.value().shape());368 if (gpu()) metal::layernorm(x.value(), w.value(), b.value(), eps, out);369 else cpu::layernorm(x.value(), w.value(), b.value(), eps, out);370371 const bool needs = grad_needed({&x, &w, &b});372 Var result(std::move(out), needs);373 if (needs) {374 Tape::get().record([x, w, b, eps, result]() {375 Tensor dx_scratch, dw_scratch, db_scratch;376 Tensor& dx = x.requires_grad() ? x.grad()377 : (dx_scratch = Tensor::zeros(x.value().shape()));378 Tensor& dw = w.requires_grad() ? w.grad()379 : (dw_scratch = Tensor::zeros(w.value().shape()));380 Tensor& db = b.requires_grad() ? b.grad()381 : (db_scratch = Tensor::zeros(b.value().shape()));382 if (gpu())383 metal::layernorm_backward(x.value(), w.value(), eps, result.grad(), dx, dw, db);384 else385 cpu::layernorm_backward(x.value(), w.value(), eps, result.grad(), dx, dw, db);386 });387 }388 return result;389}390391Var embedding(const Var& weight, const Tensor& ids) {392 const int64_t C = weight.value().size(1);393 std::vector<int64_t> out_shape = ids.shape();394 out_shape.push_back(C);395 Tensor out = Tensor::empty(std::move(out_shape));396 if (gpu()) metal::embedding(weight.value(), ids, out);397 else cpu::embedding(weight.value(), ids, out);398399 const bool needs = grad_needed({&weight});400 Var result(std::move(out), needs);401 if (needs) {402 Tape::get().record([weight, ids, result]() {403 if (!weight.requires_grad()) return;404 if (gpu()) metal::embedding_backward(ids, result.grad(), weight.grad());405 else cpu::embedding_backward(ids, result.grad(), weight.grad());406 });407 }408 return result;409}410411Tensor rope_freqs(int64_t head_dim, float theta, float scale_factor,412 float low_freq_factor, float high_freq_factor, int64_t original_ctx) {413 Tensor freqs = Tensor::empty({head_dim / 2});414 float* f = freqs.data<float>();415 constexpr float kTwoPi = 6.28318530717958647692f;416 for (int64_t k = 0; k < head_dim / 2; ++k) {417 float inv = std::pow(theta, -2.0f * float(k) / float(head_dim));418 if (scale_factor > 0.0f) {419 // HF "llama3" rope scaling (modeling_rope_utils.py).420 const float wavelen = kTwoPi / inv;421 const float lo = float(original_ctx) / low_freq_factor; // long waves422 const float hi = float(original_ctx) / high_freq_factor; // short waves423 if (wavelen > lo) {424 inv /= scale_factor;425 } else if (wavelen > hi) {426 const float s = (float(original_ctx) / wavelen - low_freq_factor) /427 (high_freq_factor - low_freq_factor);428 inv = (1.0f - s) * inv / scale_factor + s * inv;429 }430 }431 f[k] = inv;432 }433 return freqs;434}435436Var rope(const Var& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset) {437 Tensor out = Tensor::empty(x.value().shape());438 if (gpu()) metal::rope(x.value(), n_heads, freqs, pos_offset, out);439 else cpu::rope(x.value(), n_heads, freqs, pos_offset, out);440441 const bool needs = grad_needed({&x});442 Var result(std::move(out), needs);443 if (needs) {444 Tape::get().record([x, n_heads, freqs, pos_offset, result]() {445 if (!x.requires_grad()) return;446 if (gpu()) metal::rope_backward(result.grad(), n_heads, freqs, pos_offset, x.grad());447 else cpu::rope_backward(result.grad(), n_heads, freqs, pos_offset, x.grad());448 });449 }450 return result;451}452453Var rope(const Var& x, int64_t n_heads, float theta, int64_t pos_offset) {454 const int64_t hd = x.value().size(2) / n_heads;455 return rope(x, n_heads, rope_freqs(hd, theta), pos_offset);456}457458Var attention(const Var& q, const Var& k, const Var& v,459 int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,460 int64_t window, float attn_softcap) {461 const int64_t B = q.value().size(0), T = q.value().size(1);462 const int64_t head_dim = q.value().size(2) / n_heads;463 Tensor out = Tensor::empty(q.value().shape());464 const bool needs = grad_needed({&q, &k, &v});465466 // Fused path: keeps only the per-row logsumexp instead of a [B,H,T,T]467 // probability tensor, so memory is linear in T rather than quadratic.468 // Softcap needs the materialized-probs path (the fused kernels' online469 // softmax has no cap hook yet).470 if (gpu() && metal::flash_supported(head_dim) && attn_softcap <= 0.0f) {471 Tensor lse = Tensor::empty({B, n_heads, T});472 metal::flash_attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads,473 causal, scale, out, lse, metal::FlashKernel::Auto,474 window);475 Var result(std::move(out), needs);476 if (needs) {477 Tape::get().record(478 [q, k, v, lse, n_heads, n_kv_heads, causal, scale, window, result]() {479 Tensor dq_scratch, dk_scratch, dv_scratch;480 Tensor& dq = q.requires_grad() ? q.grad()481 : (dq_scratch = Tensor::zeros(q.value().shape()));482 Tensor& dk = k.requires_grad() ? k.grad()483 : (dk_scratch = Tensor::zeros(k.value().shape()));484 Tensor& dv = v.requires_grad() ? v.grad()485 : (dv_scratch = Tensor::zeros(v.value().shape()));486 metal::flash_attention_backward(q.value(), k.value(), v.value(),487 result.value(), lse, result.grad(),488 n_heads, n_kv_heads, causal, scale,489 dq, dk, dv,490 metal::FlashKernel::Auto, window);491 });492 }493 return result;494 }495496 Tensor probs = Tensor::empty({B, n_heads, T, T});497 if (gpu())498 metal::attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads, causal,499 scale, out, &probs, window, attn_softcap);500 else501 cpu::attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads, causal,502 scale, out, &probs, window, attn_softcap);503504 Var result(std::move(out), needs);505 if (needs) {506 Tape::get().record(507 [q, k, v, probs, n_heads, n_kv_heads, scale, attn_softcap, result]() {508 Tensor dq_scratch, dk_scratch, dv_scratch;509 Tensor& dq = q.requires_grad() ? q.grad()510 : (dq_scratch = Tensor::zeros(q.value().shape()));511 Tensor& dk = k.requires_grad() ? k.grad()512 : (dk_scratch = Tensor::zeros(k.value().shape()));513 Tensor& dv = v.requires_grad() ? v.grad()514 : (dv_scratch = Tensor::zeros(v.value().shape()));515 if (gpu())516 metal::attention_backward(q.value(), k.value(), v.value(), probs,517 result.value(), result.grad(), n_heads,518 n_kv_heads, scale, dq, dk, dv,519 attn_softcap);520 else521 cpu::attention_backward(q.value(), k.value(), v.value(), probs,522 result.value(), result.grad(), n_heads,523 n_kv_heads, scale, dq, dk, dv,524 attn_softcap);525 });526 }527 return result;528}529530Var cross_entropy(const Var& logits, const Tensor& targets) {531 const bool needs = grad_needed({&logits});532 const int64_t N = logits.value().size(0);533534 // n_valid comes from the CPU-resident targets (written by the data535 // loader, never touched by the GPU).536 int64_t n_valid = 0;537 for (int64_t i = 0; i < N; ++i) {538 const int32_t t = targets.dtype() == DType::I32539 ? targets.data<int32_t>()[i]540 : int32_t(targets.data<uint16_t>()[i]);541 if (t >= 0) ++n_valid;542 }543544 // llm.c pattern: the logit gradient is a byproduct of the forward pass;545 // save it and scale by d(loss) at backward time.546 Tensor dlogits;547 if (needs) dlogits = Tensor::zeros(logits.value().shape());548549 Tensor loss_out;550 if (gpu()) {551 Tensor losses = Tensor::empty({N});552 loss_out = Tensor::empty({1});553 metal::cross_entropy(logits.value(), targets, n_valid, losses, loss_out,554 needs ? &dlogits : nullptr);555 } else {556 const float loss =557 cpu::cross_entropy(logits.value(), targets, needs ? &dlogits : nullptr);558 loss_out = Tensor::full({1}, loss);559 }560561 Var result(std::move(loss_out), needs);562 if (needs) {563 Tape::get().record([logits, dlogits, result]() {564 if (logits.requires_grad())565 axpy_tensor(logits.grad(), dlogits, result.grad());566 });567 }568 return result;569}570571Var reshape(const Var& x, std::vector<int64_t> shape) {572 return x.reshaped(std::move(shape));573}574575} // namespace forge::ops576