Something went wrong. Try again.
Fork of https://github.com/xenova/microgpt.js
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765////// The most atomic way to train and run inference for a GPT in pure, dependency-free Zig./// This file is the complete algorithm./// Everything else is just efficiency.////// @karpathy (original Python), @xenova (JavaScript port), Zig port///const std = @import("std");const math = std.math;const Allocator = std.mem.Allocator;const Io = std.Io;const ArrayList = std.ArrayList;
// ─── Mersenne Twister 19937 (matches Python's random module exactly) ───
var mt_state: [624]u32 = undefined;var mt_idx: u32 = 625;var gauss_next: ?f64 = null;
fn mtSeed(n: u64) void { var key_buf: [4]u32 = undefined; var key_len: usize = 0; var v = n; if (v == 0) { key_buf[0] = 0; key_len = 1; } else { while (v > 0) : (v = v / 0x100000000) { key_buf[key_len] = @as(u32, @intCast(v & 0xFFFFFFFF)); key_len += 1; } } const key = key_buf[0..key_len];
mt_state[0] = 19650218; for (1..624) |idx| { const prev = mt_state[idx - 1]; const xor = prev ^ (prev >> 30); const product: u32 = @truncate(@as(u64, 1812433253) *% @as(u64, xor)); mt_state[idx] = product +% @as(u32, @intCast(idx)); } var i: usize = 1; var j: usize = 0; var k: usize = @max(624, key.len); while (k > 0) : (k -= 1) { if (i >= 624) { mt_state[0] = mt_state[623]; i = 1; } if (j >= key.len) j = 0; const prev1 = mt_state[i - 1]; const xor1: u32 = @truncate(@as(u64, 1664525) *% @as(u64, prev1 ^ (prev1 >> 30))); mt_state[i] = (mt_state[i] ^ xor1) +% key[j] +% @as(u32, @intCast(j)); i += 1; j += 1; } k = 623; while (k > 0) : (k -= 1) { if (i >= 624) { mt_state[0] = mt_state[623]; i = 1; } const prev2 = mt_state[i - 1]; const xor2: u32 = @truncate(@as(u64, 1566083941) *% @as(u64, prev2 ^ (prev2 >> 30))); mt_state[i] = (mt_state[i] ^ xor2) -% @as(u32, @intCast(i)); i += 1; } mt_state[0] = 0x80000000; mt_idx = 624; gauss_next = null;}
fn mtInt32() u32 { if (mt_idx >= 624) { for (0..624) |k| { const y = (mt_state[k] & 0x80000000) | (mt_state[(k + 1) % 624] & 0x7FFFFFFF); mt_state[k] = mt_state[(k + 397) % 624] ^ (y >> 1) ^ (if (y & 1 != 0) @as(u32, 0x9908B0DF) else @as(u32, 0)); } mt_idx = 0; } var y = mt_state[mt_idx]; mt_idx += 1; y ^= y >> 11; y ^= (y << 7) & 0x9D2C5680; y ^= (y << 15) & 0xEFC60000; y ^= y >> 18; return y;}
fn mtRandom() f64 { const a: f64 = @floatFromInt(mtInt32() >> 5); const b: f64 = @floatFromInt(mtInt32() >> 6); return (a * 67108864.0 + b) / 9007199254740992.0;}
fn mtGauss(mu: f64, sigma: f64) f64 { var z: f64 = undefined; if (gauss_next) |gn| { z = gn; gauss_next = null; } else { const x2pi = mtRandom() * 2.0 * math.pi; const g2rad = @sqrt(-2.0 * @log(1.0 - mtRandom())); z = @cos(x2pi) * g2rad; gauss_next = @sin(x2pi) * g2rad; } return mu + z * sigma;}
fn mtShuffle(comptime T: type, arr: []T) void { var i: usize = arr.len - 1; while (i > 0) : (i -= 1) { const k: u32 = 32 - @clz(@as(u32, @intCast(i + 1))); const shift: u5 = @intCast(32 -| k); // saturating subtract for the case k=32 var r: usize = @intCast(mtInt32() >> shift); while (r > i) { r = @intCast(mtInt32() >> shift); } std.mem.swap(T, &arr[i], &arr[r]); }}
fn mtChoices(weights: []const f64, alloc: Allocator) usize { const n = weights.len; const cum = alloc.alloc(f64, n) catch unreachable; cum[0] = weights[0]; for (1..n) |i| { cum[i] = cum[i - 1] + weights[i]; } const x = mtRandom() * cum[n - 1]; var lo: usize = 0; var hi: usize = n - 1; while (lo < hi) { const mid = (lo + hi) >> 1; if (x < cum[mid]) { hi = mid; } else { lo = mid + 1; } } return lo;}
// ─── Value (Autograd) ───
const Value = struct { data: f64, grad: f64, c0: ?*Value, c1: ?*Value, lg0: f64, lg1: f64, nch: u2, gen: u32,
fn create(alloc: Allocator, data: f64) *Value { const v = alloc.create(Value) catch unreachable; v.* = .{ .data = data, .grad = 0, .c0 = null, .c1 = null, .lg0 = 0, .lg1 = 0, .nch = 0, .gen = 0 }; return v; }
fn bin(alloc: Allocator, data: f64, a: *Value, b: *Value, lg0: f64, lg1: f64) *Value { const v = alloc.create(Value) catch unreachable; v.* = .{ .data = data, .grad = 0, .c0 = a, .c1 = b, .lg0 = lg0, .lg1 = lg1, .nch = 2, .gen = 0 }; return v; }
fn una(alloc: Allocator, data: f64, a: *Value, lg0: f64) *Value { const v = alloc.create(Value) catch unreachable; v.* = .{ .data = data, .grad = 0, .c0 = a, .c1 = null, .lg0 = lg0, .lg1 = 0, .nch = 1, .gen = 0 }; return v; }
fn addV(self: *Value, other: *Value, alloc: Allocator) *Value { return bin(alloc, self.data + other.data, self, other, 1, 1); }
fn addS(self: *Value, s: f64, alloc: Allocator) *Value { return una(alloc, self.data + s, self, 1); }
fn mulV(self: *Value, other: *Value, alloc: Allocator) *Value { return bin(alloc, self.data * other.data, self, other, other.data, self.data); }
fn mulS(self: *Value, s: f64, alloc: Allocator) *Value { return una(alloc, self.data * s, self, s); }
fn powS(self: *Value, p: f64, alloc: Allocator) *Value { return una(alloc, math.pow(f64, self.data, p), self, p * math.pow(f64, self.data, p - 1)); }
fn logV(self: *Value, alloc: Allocator) *Value { return una(alloc, @log(self.data), self, 1.0 / self.data); }
fn expV(self: *Value, alloc: Allocator) *Value { const e = @exp(self.data); return una(alloc, e, self, e); }
fn relu(self: *Value, alloc: Allocator) *Value { return una(alloc, @max(0.0, self.data), self, if (self.data > 0) @as(f64, 1.0) else @as(f64, 0.0)); }
fn neg(self: *Value, alloc: Allocator) *Value { return una(alloc, -self.data, self, -1.0); }
fn subS(self: *Value, s: f64, alloc: Allocator) *Value { return self.addS(-s, alloc); }
fn divV(self: *Value, other: *Value, alloc: Allocator) *Value { return self.mulV(other.powS(-1, alloc), alloc); }
fn divS(self: *Value, s: f64, alloc: Allocator) *Value { return self.mulS(1.0 / s, alloc); }};
var backward_gen: u32 = 0;
fn doBackward(loss: *Value, alloc: Allocator) void { backward_gen += 1; const gen = backward_gen;
var topo: ArrayList(*Value) = .empty; const Frame = struct { v: *Value, state: u8 }; var stack: ArrayList(Frame) = .empty; stack.append(alloc, .{ .v = loss, .state = 0 }) catch unreachable;
while (stack.items.len > 0) { const frame = &stack.items[stack.items.len - 1]; if (frame.v.gen == gen and frame.state == 0) { _ = stack.pop(); continue; } if (frame.state == 0) { frame.v.gen = gen; if (frame.v.nch >= 1) { frame.state = 1; stack.append(alloc, .{ .v = frame.v.c0.?, .state = 0 }) catch unreachable; continue; } } if (frame.state <= 1) { if (frame.v.nch == 2) { frame.state = 2; stack.append(alloc, .{ .v = frame.v.c1.?, .state = 0 }) catch unreachable; continue; } } topo.append(alloc, frame.v) catch unreachable; _ = stack.pop(); }
loss.grad = 1; var i: usize = topo.items.len; while (i > 0) { i -= 1; const v = topo.items[i]; const g = v.grad; if (v.nch >= 1) v.c0.?.grad += v.lg0 * g; if (v.nch == 2) v.c1.?.grad += v.lg1 * g; }}
// ─── Helper functions ───
fn sumValues(vals: []*Value, alloc: Allocator) *Value { var result = vals[0]; for (vals[1..]) |v| { result = result.addV(v, alloc); } return result;}
fn doLinear(x: []*Value, w: []const []*Value, alloc: Allocator) []*Value { const nout = w.len; const out = alloc.alloc(*Value, nout) catch unreachable; for (0..nout) |oi| { const row = w[oi]; var acc = row[0].mulV(x[0], alloc); for (1..x.len) |j| { acc = acc.addV(row[j].mulV(x[j], alloc), alloc); } out[oi] = acc; } return out;}
fn doSoftmax(logits: []*Value, alloc: Allocator) []*Value { var max_val: f64 = -math.inf(f64); for (logits) |v| { if (v.data > max_val) max_val = v.data; } const n = logits.len; const exps = alloc.alloc(*Value, n) catch unreachable; for (0..n) |i| { exps[i] = logits[i].subS(max_val, alloc).expV(alloc); } const total = sumValues(exps, alloc); const result = alloc.alloc(*Value, n) catch unreachable; for (0..n) |i| { result[i] = exps[i].divV(total, alloc); } return result;}
fn doRmsnorm(x: []*Value, alloc: Allocator) []*Value { const n = x.len; const sq = alloc.alloc(*Value, n) catch unreachable; for (0..n) |i| { sq[i] = x[i].mulV(x[i], alloc); } const ms = sumValues(sq, alloc).mulS(1.0 / @as(f64, @floatFromInt(n)), alloc); const s = ms.addS(1e-5, alloc).powS(-0.5, alloc); const result = alloc.alloc(*Value, n) catch unreachable; for (0..n) |i| { result[i] = x[i].mulV(s, alloc); } return result;}
// ─── Model hyperparameters ───const n_embd = 16;const n_head = 4;const n_layer = 1;const block_size = 16;const head_dim = n_embd / n_head;const attn_scale: f64 = 1.0 / @sqrt(@as(f64, head_dim));
const StateDict = struct { wte: [][]*Value, wpe: [][]*Value, lm_head: [][]*Value, attn_wq: [n_layer][][]*Value, attn_wk: [n_layer][][]*Value, attn_wv: [n_layer][][]*Value, attn_wo: [n_layer][][]*Value, mlp_fc1: [n_layer][][]*Value, mlp_fc2: [n_layer][][]*Value,};
fn makeMatrix(alloc: Allocator, nout: usize, nin: usize, stdev: f64) [][]*Value { const rows = alloc.alloc([]*Value, nout) catch unreachable; for (0..nout) |i| { const row = alloc.alloc(*Value, nin) catch unreachable; for (0..nin) |j| { row[j] = Value.create(alloc, mtGauss(0, stdev)); } rows[i] = row; } return rows;}
fn doGpt( token_id: usize, pos_id: usize, keys: *[n_layer]ArrayList([]*Value), kv_values: *[n_layer]ArrayList([]*Value), sd: *const StateDict, alloc: Allocator,) []*Value { const tok_emb = sd.wte[token_id]; const pos_emb = sd.wpe[pos_id]; var x = alloc.alloc(*Value, n_embd) catch unreachable; for (0..n_embd) |i| { x[i] = tok_emb[i].addV(pos_emb[i], alloc); } x = doRmsnorm(x, alloc);
for (0..n_layer) |li| { // 1) Multi-head attention block const x_residual = x; var xn = doRmsnorm(x, alloc);
const q = doLinear(xn, sd.attn_wq[li], alloc); const k_vec = doLinear(xn, sd.attn_wk[li], alloc); const v_vec = doLinear(xn, sd.attn_wv[li], alloc);
keys[li].append(alloc, k_vec) catch unreachable; kv_values[li].append(alloc, v_vec) catch unreachable;
const x_attn = alloc.alloc(*Value, n_embd) catch unreachable; for (0..n_head) |h| { const hs = h * head_dim; const q_h = q[hs .. hs + head_dim];
const num_t = keys[li].items.len; const attn_logits = alloc.alloc(*Value, num_t) catch unreachable; for (0..num_t) |t| { const kt = keys[li].items[t]; const dot = alloc.alloc(*Value, head_dim) catch unreachable; for (0..head_dim) |j_| { dot[j_] = q_h[j_].mulV(kt[hs + j_], alloc); } attn_logits[t] = sumValues(dot, alloc).mulS(attn_scale, alloc); } const attn_weights = doSoftmax(attn_logits, alloc);
for (0..head_dim) |j_| { const weighted = alloc.alloc(*Value, num_t) catch unreachable; for (0..num_t) |t| { weighted[t] = attn_weights[t].mulV(kv_values[li].items[t][hs + j_], alloc); } x_attn[hs + j_] = sumValues(weighted, alloc); } }
xn = doLinear(x_attn, sd.attn_wo[li], alloc); // residual connection const x_res1 = alloc.alloc(*Value, n_embd) catch unreachable; for (0..n_embd) |i| { x_res1[i] = xn[i].addV(x_residual[i], alloc); }
// 2) MLP block const x_residual2 = x_res1; const xn2 = doRmsnorm(x_res1, alloc); var fc1_out = doLinear(xn2, sd.mlp_fc1[li], alloc); for (0..fc1_out.len) |i| { fc1_out[i] = fc1_out[i].relu(alloc); } const fc2_out = doLinear(fc1_out, sd.mlp_fc2[li], alloc); x = alloc.alloc(*Value, n_embd) catch unreachable; for (0..n_embd) |i| { x[i] = fc2_out[i].addV(x_residual2[i], alloc); } }
return doLinear(x, sd.lm_head, alloc);}
// ─── f64 fast-path for inference (no autograd) ───
fn doLinearF64(x: []const f64, w: []const []*Value, alloc: Allocator) []f64 { const nout = w.len; const out = alloc.alloc(f64, nout) catch unreachable; for (0..nout) |oi| { const row = w[oi]; var acc: f64 = 0; for (0..x.len) |j| { acc += row[j].data * x[j]; } out[oi] = acc; } return out;}
fn doSoftmaxF64(logits: []f64) []f64 { var max_val: f64 = -math.inf(f64); for (logits) |v| { if (v > max_val) max_val = v; } var total: f64 = 0; for (logits) |*v| { v.* = @exp(v.* - max_val); total += v.*; } for (logits) |*v| { v.* /= total; } return logits;}
fn doRmsnormF64(x: []const f64, alloc: Allocator) []f64 { const n = x.len; var ms: f64 = 0; for (x) |xi| { ms += xi * xi; } ms /= @floatFromInt(n); const s = 1.0 / @sqrt(ms + 1e-5); const result = alloc.alloc(f64, n) catch unreachable; for (0..n) |i| { result[i] = x[i] * s; } return result;}
fn doGptF64( token_id: usize, pos_id: usize, keys: *[n_layer]ArrayList([]f64), kv_values: *[n_layer]ArrayList([]f64), sd: *const StateDict, alloc: Allocator,) []f64 { const tok_emb = sd.wte[token_id]; const pos_emb = sd.wpe[pos_id]; var x = alloc.alloc(f64, n_embd) catch unreachable; for (0..n_embd) |i| { x[i] = tok_emb[i].data + pos_emb[i].data; } x = doRmsnormF64(x, alloc);
for (0..n_layer) |li| { // 1) Multi-head attention block const x_residual = x; var xn = doRmsnormF64(x, alloc);
const q = doLinearF64(xn, sd.attn_wq[li], alloc); const k_vec = doLinearF64(xn, sd.attn_wk[li], alloc); const v_vec = doLinearF64(xn, sd.attn_wv[li], alloc);
keys[li].append(alloc, k_vec) catch unreachable; kv_values[li].append(alloc, v_vec) catch unreachable;
const x_attn = alloc.alloc(f64, n_embd) catch unreachable; for (0..n_head) |h| { const hs = h * head_dim; const q_h = q[hs .. hs + head_dim];
const num_t = keys[li].items.len; const attn_logits = alloc.alloc(f64, num_t) catch unreachable; for (0..num_t) |t| { const kt = keys[li].items[t]; var dot: f64 = 0; for (0..head_dim) |j_| { dot += q_h[j_] * kt[hs + j_]; } attn_logits[t] = dot * attn_scale; } const attn_weights = doSoftmaxF64(attn_logits);
for (0..head_dim) |j_| { var sum: f64 = 0; for (0..num_t) |t| { sum += attn_weights[t] * kv_values[li].items[t][hs + j_]; } x_attn[hs + j_] = sum; } }
xn = doLinearF64(x_attn, sd.attn_wo[li], alloc); // residual connection const x_res1 = alloc.alloc(f64, n_embd) catch unreachable; for (0..n_embd) |i| { x_res1[i] = xn[i] + x_residual[i]; }
// 2) MLP block const x_residual2 = x_res1; const xn2 = doRmsnormF64(x_res1, alloc); const fc1_out = doLinearF64(xn2, sd.mlp_fc1[li], alloc); for (fc1_out) |*v| { v.* = @max(0.0, v.*); } const fc2_out = doLinearF64(fc1_out, sd.mlp_fc2[li], alloc); x = alloc.alloc(f64, n_embd) catch unreachable; for (0..n_embd) |i| { x[i] = fc2_out[i] + x_residual2[i]; } }
return doLinearF64(x, sd.lm_head, alloc);}
// ─── Main ───pub fn main(init: std.process.Init) !void { const io = init.io; const alloc = init.gpa;
var stdout_buffer: [4096]u8 = undefined; var stdout_writer = Io.File.stdout().writer(io, &stdout_buffer); const stdout = &stdout_writer.interface;
// Read input.txt const file_data = try Io.Dir.cwd().readFileAlloc(io, "input.txt", alloc, .limited(10 * 1024 * 1024)); defer alloc.free(file_data);
// Split into docs var doc_list: ArrayList([]const u8) = .empty; defer doc_list.deinit(alloc); const trimmed_data = std.mem.trim(u8, file_data, &[_]u8{ '\n', '\r', ' ', '\t' }); var line_iter = std.mem.splitScalar(u8, trimmed_data, '\n'); while (line_iter.next()) |line| { const trimmed = std.mem.trim(u8, line, &[_]u8{ ' ', '\t', '\r' }); if (trimmed.len > 0) { try doc_list.append(alloc, trimmed); } } var docs = doc_list.items;
mtSeed(42); mtShuffle([]const u8, docs); try stdout.print("num docs: {d}\n", .{docs.len});
// Tokenizer: unique characters var char_set = std.AutoHashMap(u8, void).init(alloc); defer char_set.deinit(); for (docs) |doc| { for (doc) |ch| { try char_set.put(ch, {}); } } var uchars_list: ArrayList(u8) = .empty; defer uchars_list.deinit(alloc); var it = char_set.keyIterator(); while (it.next()) |key| { try uchars_list.append(alloc, key.*); } std.mem.sort(u8, uchars_list.items, {}, std.sort.asc(u8)); const uchars = uchars_list.items;
var char_to_id = std.AutoHashMap(u8, usize).init(alloc); defer char_to_id.deinit(); for (uchars, 0..) |ch, i| { try char_to_id.put(ch, i); }
const BOS = uchars.len; const vocab_size = uchars.len + 1; try stdout.print("vocab size: {d}\n", .{vocab_size});
// Initialize parameters const param_alloc = std.heap.page_allocator;
var sd: StateDict = undefined; sd.wte = makeMatrix(param_alloc, vocab_size, n_embd, 0.08); sd.wpe = makeMatrix(param_alloc, block_size, n_embd, 0.08); sd.lm_head = makeMatrix(param_alloc, vocab_size, n_embd, 0.08); for (0..n_layer) |i| { sd.attn_wq[i] = makeMatrix(param_alloc, n_embd, n_embd, 0.08); sd.attn_wk[i] = makeMatrix(param_alloc, n_embd, n_embd, 0.08); sd.attn_wv[i] = makeMatrix(param_alloc, n_embd, n_embd, 0.08); sd.attn_wo[i] = makeMatrix(param_alloc, n_embd, n_embd, 0.08); sd.mlp_fc1[i] = makeMatrix(param_alloc, 4 * n_embd, n_embd, 0.08); sd.mlp_fc2[i] = makeMatrix(param_alloc, n_embd, 4 * n_embd, 0.08); }
// Flatten params var params_list: ArrayList(*Value) = .empty; const matrices = [_][][]*Value{ sd.wte, sd.wpe, sd.lm_head }; for (matrices) |mat_| { for (mat_) |row| { for (row) |p| { try params_list.append(param_alloc, p); } } } for (0..n_layer) |li| { const layer_mats = [_][][]*Value{ sd.attn_wq[li], sd.attn_wk[li], sd.attn_wv[li], sd.attn_wo[li], sd.mlp_fc1[li], sd.mlp_fc2[li], }; for (layer_mats) |mat_| { for (mat_) |row| { for (row) |p| { try params_list.append(param_alloc, p); } } } } const params = params_list.items; try stdout.print("num params: {d}\n", .{params.len}); try stdout_writer.flush();
// Adam optimizer buffers const m_buf = try param_alloc.alloc(f64, params.len); const v_buf = try param_alloc.alloc(f64, params.len); @memset(m_buf, 0); @memset(v_buf, 0);
const learning_rate: f64 = 0.01; const beta1: f64 = 0.85; const beta2: f64 = 0.99; const eps_adam: f64 = 1e-8; const num_steps: usize = 1000;
// Training loop for (0..num_steps) |step| { const doc = docs[step % docs.len]; var tokens: ArrayList(usize) = .empty; defer tokens.deinit(alloc); try tokens.ensureTotalCapacity(alloc, doc.len + 2); tokens.appendAssumeCapacity(BOS); for (doc) |ch| { tokens.appendAssumeCapacity(char_to_id.get(ch).?); } tokens.appendAssumeCapacity(BOS); const n = @min(block_size, tokens.items.len - 1);
var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator); defer arena.deinit(); const step_alloc = arena.allocator();
var keys: [n_layer]ArrayList([]*Value) = undefined; var kv_values: [n_layer]ArrayList([]*Value) = undefined; for (0..n_layer) |li| { keys[li] = .empty; kv_values[li] = .empty; }
var losses: ArrayList(*Value) = .empty; losses.ensureTotalCapacity(step_alloc, n) catch unreachable; for (0..n) |pos_id| { const token_id = tokens.items[pos_id]; const target_id = tokens.items[pos_id + 1]; const logits = doGpt(token_id, pos_id, &keys, &kv_values, &sd, step_alloc); const probs = doSoftmax(logits, step_alloc); const loss_t = probs[target_id].logV(step_alloc).neg(step_alloc); losses.append(step_alloc, loss_t) catch unreachable; } const loss = sumValues(losses.items, step_alloc).mulS(1.0 / @as(f64, @floatFromInt(n)), step_alloc);
doBackward(loss, step_alloc);
const lr_t = learning_rate * (1.0 - @as(f64, @floatFromInt(step)) / @as(f64, @floatFromInt(num_steps))); const step_f: f64 = @floatFromInt(step + 1); const bc1 = 1.0 - math.pow(f64, beta1, step_f); const bc2 = 1.0 - math.pow(f64, beta2, step_f); for (params, 0..) |p, i| { m_buf[i] = beta1 * m_buf[i] + (1.0 - beta1) * p.grad; v_buf[i] = beta2 * v_buf[i] + (1.0 - beta2) * p.grad * p.grad; const m_hat = m_buf[i] / bc1; const v_hat = v_buf[i] / bc2; p.data -= lr_t * m_hat / (@sqrt(v_hat) + eps_adam); p.grad = 0; }
try stdout.print("step {d:>4} / {d:>4} | loss {d:.4}\r", .{ step + 1, num_steps, loss.data }); try stdout_writer.flush(); }
// Inference const temperature: f64 = 0.5; try stdout.print("\n--- inference (new, hallucinated names) ---\n", .{}); try stdout_writer.flush();
for (0..20) |sample_idx| { var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator); defer arena.deinit(); const inf_alloc = arena.allocator();
var f_keys: [n_layer]ArrayList([]f64) = undefined; var f_values: [n_layer]ArrayList([]f64) = undefined; for (0..n_layer) |li| { f_keys[li] = .empty; f_values[li] = .empty; }
var token_id: usize = BOS; var sample: ArrayList(u8) = .empty;
for (0..block_size) |pos_id| { const logits = doGptF64(token_id, pos_id, &f_keys, &f_values, &sd, inf_alloc); const scaled = inf_alloc.alloc(f64, logits.len) catch unreachable; for (0..logits.len) |i| { scaled[i] = logits[i] / temperature; } const probs = doSoftmaxF64(scaled); token_id = mtChoices(probs, inf_alloc); if (token_id == BOS) break; sample.append(inf_alloc, uchars[token_id]) catch unreachable; } try stdout.print("sample {d:>2}: {s}\n", .{ sample_idx + 1, sample.items }); } try stdout_writer.flush();}