From fe01b29fabfebfd070489d4f832955552d697310 Mon Sep 17 00:00:00 2001 From: moderation Date: Sun, 15 Feb 2026 19:49:24 -0500 Subject: [PATCH] Optimize Zig port: remove @constCast, f64 inference fast path, readability improvements - Eliminate all @constCast by changing [][]const *Value to [][]*Value - Add f64 fast-path for inference (doLinearF64, doSoftmaxF64, doRmsnormF64, doGptF64) - Break up dense Mersenne Twister expressions into readable intermediates - Use std.mem.swap in mtShuffle - Inline sum in doLinear to avoid temp allocation per row - Pre-size ArrayLists (tokens, losses) where lengths are known - Shrink Value.gen from u64 to u32 --- microgpt.zig | 764 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 764 insertions(+) create mode 100644 microgpt.zig diff --git a/microgpt.zig b/microgpt.zig new file mode 100644 index 0000000..160b0e7 --- /dev/null +++ b/microgpt.zig @@ -0,0 +1,764 @@ +/// +/// The most atomic way to train and inference 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}\n", .{ 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(); +} -- knotmirror 0.0.1