diff --git a/.trinity/seals/ml_train_hslm.json b/.trinity/seals/ml_train_hslm.json index 1f6c0c9c8d..ae1d379241 100644 --- a/.trinity/seals/ml_train_hslm.json +++ b/.trinity/seals/ml_train_hslm.json @@ -19,23 +19,23 @@ } ], "config": "gen=c,rust,verilog,zig;tests=zig", - "gen_hash_c": "sha256:9261d24381b2c2b5506fad987fc13fbfd5e7142a00c5e67fb878e56862102c7f", - "gen_hash_rust": "sha256:78d0634dbb7f15a4fb098e71779d47a69149fe776703d6a791632f608e3315f2", - "gen_hash_verilog": "sha256:b9743c3e95cd4c48765383ea8c8d7a1ffc5551c925e23be9bbae74d7f30e0550", - "gen_hash_zig": "sha256:00a2ae9af248ea9bae037596bf77b759303d460f8eb283cef3b8d6284b879ea5", + "gen_hash_c": "sha256:9d92144a94228e46d3131ea7f5fe1bc86ead3a721add372e0051fee58317bc46", + "gen_hash_rust": "sha256:fa4751c35bca1feefaad3cd55dc1112fd7df2db4e6c5d7e21c34a024275d4285", + "gen_hash_verilog": "sha256:cf563313b4759cb371ab1a6f4aa5ad22e9770984d31d27bb2b7b448c48a54817", + "gen_hash_zig": "sha256:bbc14e0e851557974231ac8219526c2b75127e280d325f8c725f2131661c5874", "module": "train_hslm", "ring": 12, "seal_schema": 2, - "sealed_at": "2026-10-19T04:28:22Z", + "sealed_at": "2026-10-19T05:13:11Z", "sealed_by": "t27c-bootstrap@0.5.2", - "spec_hash": "sha256:8fd4733d3ff131fc85c947e7e6146666d0f06fdaf3a5e563d7ce2fa70eb61b8d", + "spec_hash": "sha256:8218a7d2003b4b78582badf5e5b3bb42117455e722a4f4805b18963c2c092859", "spec_path": "specs/ml/train_hslm.t27", "tests": { "failed": 0, "failing": [], "forced": false, - "passed": 21, - "total": 21, + "passed": 27, + "total": 27, "vacuous": 0 }, "toolchain": { diff --git a/specs/ml/train_hslm.t27 b/specs/ml/train_hslm.t27 index a23d717e14..30ce0fb131 100644 --- a/specs/ml/train_hslm.t27 +++ b/specs/ml/train_hslm.t27 @@ -57,9 +57,11 @@ ; kernel that is not bit-identical; it only runs in backward). speed=2 (the default) also skips work ; whose result nobody reads: the head at positions 0..25 and, in the last block at those positions, ; everything but the key and value that the last position attends to. Under loss_positions = k only -; positions 0..26-k are unscored, so k = 27 runs the full forward everywhere and skips nothing. The -; tests check that speed 1 and speed 2 produce the same logits as speed 0 to the bit. The model is -; not reduced. +; positions 0..26-k are unscored, so k = 27 runs the full forward everywhere and skips nothing. speed=3 (#10849) +; computes the speed 2 window layer by layer over all 27 positions, so each ternary matrix and each +; gradient row is read once per window instead of once per position, with every sum in speed 2's +; order. The tests check that speed 1, 2 and 3 produce the same logits as speed 0 to the bit, and that +; speed 3 produces the same logits, loss and gradients as speed 2 to the bit. The model is not reduced. ; WHAT DIFFERS FROM THE ZIG TRAINER (trinity-training @ 3e9648c, run R5): f64, not f32; LAMB, not ; AdamW; one global clip, not per tensor; linear warmup + cosine to 1e-6, not the phi schedule; the ; embedding table and RMSNorm gains are trained (the trainer froze the table); the gradient is that of @@ -448,6 +450,11 @@ fn tr_prefix_forward(token: usize, position: usize) -> void { /// One 27-token window: every position in order (each attends to the cache of the earlier ones); /// the logits of the last position land in `logits`, its activations stay for hslm_backward fn tr_window_forward(tokens: []const u16, start: usize, logits: []f64) -> void { + if (speed_level == 3) { + tr3_forward(tokens, start, CONTEXT_LEN - 1, CONTEXT_LEN - 1); + hslm_forward(@as(usize, @intCast(tokens[start + CONTEXT_LEN - 1])), CONTEXT_LEN - 1, logits); + return; + } var p : usize = 0; while (p < CONTEXT_LEN) { const tok : usize = @intCast(tokens[start + p]); @@ -488,6 +495,9 @@ fn tr_window_loss(tokens: []const u16, start: usize) -> f64 { /// Forward, loss and backward of one window under the k-position objective (1 <= k <= 27); /// gradients accumulate in nn/hslm's gw_* and g_*; returns the mean of the k losses fn tr_window_grad_k(tokens: []const u16, start: usize, k: usize) -> f64 { + if (speed_level == 3) { + return tr3_window_grad(tokens, start, k); + } const inv_k : f64 = 1.0 / @as(f64, @floatFromInt(k)); var sum : f64 = 0.0; var p : usize = 0; @@ -518,6 +528,21 @@ fn tr_window_grad_k(tokens: []const u16, start: usize, k: usize) -> f64 { /// cross-entropy of each of the last k against its next token, their mean fn tr_window_loss_k(tokens: []const u16, start: usize, k: usize) -> f64 { var sum : f64 = 0.0; + if (speed_level == 3) { + tr3_forward(tokens, start, CONTEXT_LEN, 0); + var q : usize = CONTEXT_LEN - k; + while (q < CONTEXT_LEN) { + const lo2 : usize = q * VOCAB_SIZE; + sum = sum + cross_entropy(b3_logits[lo2..lo2 + VOCAB_SIZE], @as(usize, @intCast(tokens[start + q + 1]))); + q = q + 1; + } + var j : usize = 0; + while (j < VOCAB_SIZE) { + ref_logits[j] = b3_logits[(CONTEXT_LEN - 1) * VOCAB_SIZE + j]; + j = j + 1; + } + return sum * (1.0 / @as(f64, @floatFromInt(k))); + } var p : usize = 0; while (p < CONTEXT_LEN) { const tok : usize = @intCast(tokens[start + p]); @@ -531,6 +556,617 @@ fn tr_window_loss_k(tokens: []const u16, start: usize, k: usize) -> f64 { return sum * (1.0 / @as(f64, @floatFromInt(k))); } +; SPEED 3: ONE MATRIX OVER THE 27 POSITIONS (#10849). Speed 2 runs a window one position at a time, +; so every position reads every ternary matrix again, and every scored position's backward streams +; the whole 17 MB of f64 weight gradients through outer_add. Speed 3 runs the same window layer by +; layer: the embedding, the norms, Q/K/V, O, the MLP and the head of all positions in one pass over +; each matrix, then the backward of all scored positions the same way, so each gradient row stays +; in L1 while every position adds to it. Nothing else moves: attention stays causal (position p +; reads the keys and values of 0..p, which the pass over K and V has already put in the cache), the +; VSA stream runs per block in position order, the last block at an unscored position still makes +; only its key and value, and earlier keys and values stay constants of a position's backward (no +; position's backward reads anything a later position's forward writes). Every output element is +; summed in the order speed 2 sums it -- j = 0, 1, ... from +0.0 in the forward, the same four +; partial sums in the transpose, positions in ascending order into every gradient -- so speed 3 is +; speed 2 to the bit, and the tests compare the bits. +; MEASURED (Railway lab, 2026-10-19, 48 cores under a load average of 59-66 from other sessions): +; train.cfg loss_positions=27 steps=10 total=32000 seed=1729 eval_every=10 eval_windows=1, the +; ReleaseFast binary of this spec, speed=2 and speed=3 started at the same moment, twice. Both runs +; printed the same losses (first 6.6079, last 6.6012, valid 6.5947) and wrote byte-equal ckpt_10.bin +; and ckpt_10.state. tok/s is 66 windows x 27 positions x 10 steps over train_s. +const TR3_LAB_TOKENS : f64 = 17820.0; +const TR3_LAB_S2_TRAIN_S : f64 = 55.09; +const TR3_LAB_S3_TRAIN_S : f64 = 37.27; +const TR3_LAB_S2_TOK_S : f64 = 323.5; +const TR3_LAB_S3_TOK_S : f64 = 478.1; +const TR3_LAB_S2_TRAIN_S_AGAIN : f64 = 55.62; +const TR3_LAB_S3_TRAIN_S_AGAIN : f64 = 37.27; +const TR3_LAB_SPEEDUP : f64 = 1.478; +const TR3_TILE : usize = 81; +var b3_tok : [CONTEXT_LEN]usize = undefined; +var b3_x : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_n1 : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_rms1 : [NUM_BLOCKS * CONTEXT_LEN]f64 = undefined; +var b3_q : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_attn : [NUM_BLOCKS * CONTEXT_LEN * NUM_HEADS * CONTEXT_LEN]f64 = undefined; +var b3_ctx : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_r : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_pre : [NUM_BLOCKS * CONTEXT_LEN * FF_DIM]f64 = undefined; +var b3_h : [NUM_BLOCKS * CONTEXT_LEN * FF_DIM]f64 = undefined; +var b3_out : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_hn : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_hrms : [CONTEXT_LEN]f64 = undefined; +var b3_logits : [CONTEXT_LEN * VOCAB_SIZE]f64 = undefined; +var b3_y : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_wt : [VOCAB_SIZE]f64 = undefined; +var b3_dl : [CONTEXT_LEN * VOCAB_SIZE]f64 = undefined; +var b3_dz : [CONTEXT_LEN * VOCAB_SIZE]f64 = undefined; +var b3_bo : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_br : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_bctx : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_bq : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_bk : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_bv : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_bn : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_tmp : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_se : [CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var b3_bh : [CONTEXT_LEN * FF_DIM]f64 = undefined; +var b3_bpre : [CONTEXT_LEN * FF_DIM]f64 = undefined; +var b3_sf : [CONTEXT_LEN * FF_DIM]f64 = undefined; +var b3_da : [CONTEXT_LEN]f64 = undefined; + +/// y[p] = x[p] W for rows p_lo..p_hi (x rows in_dim wide, y rows out_dim wide). Each weight row is +/// read once for all rows; a column tile of TR3_TILE keeps that tile of every row in L1. Per output +/// the terms are added in the order j = 0, 1, ... from +0.0, as ternary_matmul adds them. +fn tr3_matmul(x: []const f64, weights: []const Trit, y: []f64, in_dim: usize, out_dim: usize, p_lo: usize, p_hi: usize) -> void { + if (p_lo >= p_hi) { + return; + } + var p : usize = p_lo; + while (p < p_hi) { + var c : usize = 0; + while (c < out_dim) { + y[p * out_dim + c] = 0.0; + c = c + 1; + } + p = p + 1; + } + var c0 : usize = 0; + while (c0 < out_dim) { + var c1 : usize = c0 + TR3_TILE; + if (c1 > out_dim) { + c1 = out_dim; + } + const width : usize = c1 - c0; + var j : usize = 0; + while (j < in_dim) { + const row : usize = j * out_dim + c0; + var c : usize = 0; + while (c < width) { + const wf : f64 = @floatFromInt(@intFromEnum(weights[row + c])); + b3_wt[c] = wf; + c = c + 1; + } + p = p_lo; + while (p < p_hi) { + const xj : f64 = x[p * in_dim + j]; + if (xj != 0.0) { + const yo : usize = p * out_dim + c0; + c = 0; + while (c < width) { + y[yo + c] = y[yo + c] + xj * b3_wt[c]; + c = c + 1; + } + } + p = p + 1; + } + j = j + 1; + } + c0 = c1; + } +} + +/// dx[p] += W dy[p] for rows p_lo..p_hi: the weight row j is turned into floats once, then each row +/// takes the four partial sums of ternary_matmul_t_add +fn tr3_t_add(dy: []const f64, weights: []const Trit, dx: []f64, in_dim: usize, out_dim: usize, p_lo: usize, p_hi: usize) -> void { + if (p_lo >= p_hi) { + return; + } + var j : usize = 0; + while (j < in_dim) { + const row : usize = j * out_dim; + var i : usize = 0; + while (i < out_dim) { + const wf : f64 = @floatFromInt(@intFromEnum(weights[row + i])); + b3_wt[i] = wf; + i = i + 1; + } + var p : usize = p_lo; + while (p < p_hi) { + const yo : usize = p * out_dim; + var a0 : f64 = 0.0; + var a1 : f64 = 0.0; + var a2 : f64 = 0.0; + var a3 : f64 = 0.0; + i = 0; + while (i + 4 <= out_dim) { + a0 = a0 + b3_wt[i] * dy[yo + i]; + a1 = a1 + b3_wt[i + 1] * dy[yo + i + 1]; + a2 = a2 + b3_wt[i + 2] * dy[yo + i + 2]; + a3 = a3 + b3_wt[i + 3] * dy[yo + i + 3]; + i = i + 4; + } + while (i < out_dim) { + a0 = a0 + b3_wt[i] * dy[yo + i]; + i = i + 1; + } + dx[p * in_dim + j] = dx[p * in_dim + j] + ((a0 + a1) + (a2 + a3)); + p = p + 1; + } + j = j + 1; + } +} + +/// gw[j * out_dim + i] += x[p][j] * dy[p][i] over rows p_lo..p_hi in ascending order, the gradient +/// row held while every row adds to it; a row whose x[p][j] is 0 is skipped, as outer_add skips it +fn tr3_outer(x: []const f64, dy: []const f64, gw: []f64, in_dim: usize, out_dim: usize, p_lo: usize, p_hi: usize) -> void { + if (p_lo >= p_hi) { + return; + } + var j : usize = 0; + while (j < in_dim) { + const row : usize = j * out_dim; + var p : usize = p_lo; + while (p < p_hi) { + const xj : f64 = x[p * in_dim + j]; + if (xj != 0.0) { + const yo : usize = p * out_dim; + var i : usize = 0; + while (i < out_dim) { + gw[row + i] = gw[row + i] + xj * dy[yo + i]; + i = i + 1; + } + } + p = p + 1; + } + j = j + 1; + } +} + +/// Rows p_lo..p_hi of a buffer of `width`-wide rows set to 0 +fn tr3_zero(buf: []f64, width: usize, p_lo: usize, p_hi: usize) -> void { + var i : usize = p_lo * width; + while (i < p_hi * width) { + buf[i] = 0.0; + i = i + 1; + } +} + +/// attention_forward's softmax and context for block b at position p, from the query of p and the +/// cached keys and values of positions 0..p +fn tr3_attend(b: usize, p: usize) -> void { + const qo : usize = (b * CONTEXT_LEN + p) * EMBED_DIM; + var head : usize = 0; + while (head < NUM_HEADS) { + const ho : usize = head * HEAD_DIM; + const ao : usize = ((b * CONTEXT_LEN + p) * NUM_HEADS + head) * CONTEXT_LEN; + var j : usize = 0; + while (j <= p) { + const kj : usize = (b * CONTEXT_LEN + j) * EMBED_DIM + ho; + var dot : f64 = 0.0; + var d : usize = 0; + while (d < HEAD_DIM) { + dot = dot + b3_q[qo + ho + d] * kv_k[kj + d]; + d = d + 1; + } + b3_attn[ao + j] = dot * SACRED_ATTN_SCALE; + j = j + 1; + } + var max_score : f64 = b3_attn[ao]; + j = 1; + while (j <= p) { + if (b3_attn[ao + j] > max_score) { + max_score = b3_attn[ao + j]; + } + j = j + 1; + } + var sum : f64 = 0.0; + j = 0; + while (j <= p) { + b3_attn[ao + j] = exp(b3_attn[ao + j] - max_score); + sum = sum + b3_attn[ao + j]; + j = j + 1; + } + j = 0; + while (j <= p) { + b3_attn[ao + j] = b3_attn[ao + j] / sum; + j = j + 1; + } + var c : usize = 0; + while (c < HEAD_DIM) { + var acc : f64 = 0.0; + var m : usize = 0; + while (m <= p) { + acc = acc + b3_attn[ao + m] * kv_v[(b * CONTEXT_LEN + m) * EMBED_DIM + ho + c]; + m = m + 1; + } + b3_ctx[qo + ho + c] = acc; + c = c + 1; + } + head = head + 1; + } +} + +/// hslm_block_forward of block b at positions 0..npos-1 at once; positions below `lo` (the last +/// block's unscored ones) make only their key and value, as tr_kv_only does +fn tr3_block_forward(b: usize, npos: usize, lo: usize) -> void { + const bw = &model.blocks[b]; + const eb : usize = b * CONTEXT_LEN * EMBED_DIM; + const fb : usize = b * CONTEXT_LEN * FF_DIM; + const span : usize = CONTEXT_LEN * EMBED_DIM; + const fspan : usize = CONTEXT_LEN * FF_DIM; + var p : usize = 0; + while (p < npos) { + const o : usize = eb + p * EMBED_DIM; + if (b > 0) { + var i : usize = 0; + while (i < EMBED_DIM) { + b3_x[o + i] = b3_out[o - span + i]; + i = i + 1; + } + } + b3_rms1[b * CONTEXT_LEN + p] = rms_norm_into(b3_x[o..o + EMBED_DIM], &bw.norm1_gamma, b3_n1[o..o + EMBED_DIM]); + p = p + 1; + } + const ns = b3_n1[eb..eb + span]; + tr3_matmul(ns, &bw.w_k, kv_k[eb..eb + span], EMBED_DIM, EMBED_DIM, 0, npos); + tr3_matmul(ns, &bw.w_v, kv_v[eb..eb + span], EMBED_DIM, EMBED_DIM, 0, npos); + tr3_matmul(ns, &bw.w_q, b3_q[eb..eb + span], EMBED_DIM, EMBED_DIM, lo, npos); + p = 0; + while (p < npos) { + const o : usize = eb + p * EMBED_DIM; + const n1 = b3_n1[o..o + EMBED_DIM]; + const kk = kv_k[o..o + EMBED_DIM]; + const vv = kv_v[o..o + EMBED_DIM]; + nudge(MAT_K, b, n1, kk, EMBED_DIM); + scale_into(kk, bw.alpha_k, kk); + nudge(MAT_V, b, n1, vv, EMBED_DIM); + scale_into(vv, bw.alpha_v, vv); + rope(kk, p, false); + if (p >= lo) { + const qq = b3_q[o..o + EMBED_DIM]; + nudge(MAT_Q, b, n1, qq, EMBED_DIM); + scale_into(qq, bw.alpha_q, qq); + rope(qq, p, false); + } + p = p + 1; + } + p = lo; + while (p < npos) { + tr3_attend(b, p); + p = p + 1; + } + tr3_matmul(b3_ctx[eb..eb + span], &bw.w_o, &b3_y, EMBED_DIM, EMBED_DIM, lo, npos); + p = lo; + while (p < npos) { + const o : usize = eb + p * EMBED_DIM; + const yo : usize = p * EMBED_DIM; + const yy = b3_y[yo..yo + EMBED_DIM]; + nudge(MAT_O, b, b3_ctx[o..o + EMBED_DIM], yy, EMBED_DIM); + scale_into(yy, bw.alpha_o, yy); + var i : usize = 0; + while (i < EMBED_DIM) { + b3_r[o + i] = b3_x[o + i] + b3_y[yo + i]; + i = i + 1; + } + p = p + 1; + } + tr3_matmul(b3_r[eb..eb + span], &bw.w_up, b3_pre[fb..fb + fspan], EMBED_DIM, FF_DIM, lo, npos); + p = lo; + while (p < npos) { + const o : usize = eb + p * EMBED_DIM; + const f : usize = fb + p * FF_DIM; + nudge(MAT_UP, b, b3_r[o..o + EMBED_DIM], b3_pre[f..f + FF_DIM], FF_DIM); + var i : usize = 0; + while (i < FF_DIM) { + b3_pre[f + i] = bw.alpha_up * b3_pre[f + i] + bw.b_up[i]; + b3_h[f + i] = relu(b3_pre[f + i]); + i = i + 1; + } + p = p + 1; + } + tr3_matmul(b3_h[fb..fb + fspan], &bw.w_down, &b3_y, FF_DIM, EMBED_DIM, lo, npos); + p = lo; + while (p < npos) { + const o : usize = eb + p * EMBED_DIM; + const f : usize = fb + p * FF_DIM; + const yo : usize = p * EMBED_DIM; + nudge(MAT_DOWN, b, b3_h[f..f + FF_DIM], b3_y[yo..yo + EMBED_DIM], EMBED_DIM); + var i : usize = 0; + while (i < EMBED_DIM) { + b3_y[yo + i] = bw.alpha_down * b3_y[yo + i] + bw.b_down[i]; + b3_out[o + i] = b3_r[o + i] + b3_y[yo + i]; + i = i + 1; + } + p = p + 1; + } + if (vsa_on) { + p = lo; + while (p < npos) { + const o : usize = eb + p * EMBED_DIM; + vsa_block_forward(b, p, b3_out[o..o + EMBED_DIM]); + p = p + 1; + } + } +} + +/// Positions 0..npos-1 of a window, layer by layer; the last block runs in full and the head runs +/// at positions lo_full..npos-1, the logits of position p in row p of b3_logits +fn tr3_forward(tokens: []const u16, start: usize, npos: usize, lo_full: usize) -> void { + var p : usize = 0; + while (p < npos) { + const tok : usize = @intCast(tokens[start + p]); + b3_tok[p] = tok; + embed_into(tok, p, b3_x[p * EMBED_DIM..(p + 1) * EMBED_DIM]); + if (vsa_on) { + vsa_embed(tok, p, vsa_in[p * VSA_DIM..(p + 1) * VSA_DIM]); + } + p = p + 1; + } + var b : usize = 0; + while (b < NUM_BLOCKS) { + var lo : usize = 0; + if (b + 1 == NUM_BLOCKS) { + lo = lo_full; + } + tr3_block_forward(b, npos, lo); + b = b + 1; + } + const lb : usize = (NUM_BLOCKS - 1) * CONTEXT_LEN * EMBED_DIM; + p = lo_full; + while (p < npos) { + const xo : usize = lb + p * EMBED_DIM; + const hrms : f64 = rms_of(b3_out[xo..xo + EMBED_DIM]); + b3_hrms[p] = hrms; + var i : usize = 0; + while (i < EMBED_DIM) { + b3_hn[p * EMBED_DIM + i] = b3_out[xo + i] / hrms; + i = i + 1; + } + p = p + 1; + } + tr3_matmul(&b3_hn, &model.w_out, &b3_logits, EMBED_DIM, VOCAB_SIZE, lo_full, npos); + p = lo_full; + while (p < npos) { + const lo2 : usize = p * VOCAB_SIZE; + nudge(MAT_OUT, 0, b3_hn[p * EMBED_DIM..(p + 1) * EMBED_DIM], b3_logits[lo2..lo2 + VOCAB_SIZE], VOCAB_SIZE); + var j : usize = 0; + while (j < VOCAB_SIZE) { + b3_logits[lo2 + j] = model.alpha_out * b3_logits[lo2 + j] * SACRED_LOGIT_SCALE + model.b_out[j]; + j = j + 1; + } + p = p + 1; + } +} + +/// attention_backward's softmax part for block b at position p: from d context (b3_bctx row p) to +/// the gradients of p's query, key and value, through RoPE and the alphas +fn tr3_attend_backward(b: usize, p: usize) -> void { + const bw = &model.blocks[b]; + const qo : usize = (b * CONTEXT_LEN + p) * EMBED_DIM; + const ro : usize = p * EMBED_DIM; + var head : usize = 0; + while (head < NUM_HEADS) { + const ho : usize = head * HEAD_DIM; + const ao : usize = ((b * CONTEXT_LEN + p) * NUM_HEADS + head) * CONTEXT_LEN; + var mean_da : f64 = 0.0; + var j : usize = 0; + while (j <= p) { + const vj : usize = (b * CONTEXT_LEN + j) * EMBED_DIM + ho; + var dot : f64 = 0.0; + var d : usize = 0; + while (d < HEAD_DIM) { + dot = dot + b3_bctx[ro + ho + d] * kv_v[vj + d]; + d = d + 1; + } + b3_da[j] = dot; + mean_da = mean_da + b3_attn[ao + j] * dot; + j = j + 1; + } + j = 0; + while (j <= p) { + b3_da[j] = b3_attn[ao + j] * (b3_da[j] - mean_da); + j = j + 1; + } + var c : usize = 0; + while (c < HEAD_DIM) { + var acc : f64 = 0.0; + var m : usize = 0; + while (m <= p) { + acc = acc + b3_da[m] * kv_k[(b * CONTEXT_LEN + m) * EMBED_DIM + ho + c]; + m = m + 1; + } + b3_bq[ro + ho + c] = SACRED_ATTN_SCALE * acc; + b3_bk[ro + ho + c] = SACRED_ATTN_SCALE * b3_da[p] * b3_q[qo + ho + c]; + b3_bv[ro + ho + c] = b3_attn[ao + p] * b3_bctx[ro + ho + c]; + c = c + 1; + } + head = head + 1; + } + const bq = b3_bq[ro..ro + EMBED_DIM]; + const bk = b3_bk[ro..ro + EMBED_DIM]; + const bv = b3_bv[ro..ro + EMBED_DIM]; + rope(bq, p, true); + rope(bk, p, true); + scale_into(bq, bw.alpha_q, bq); + scale_into(bk, bw.alpha_k, bk); + scale_into(bv, bw.alpha_v, bv); +} + +/// hslm_block_backward of block b at the scored positions lo..26 at once: b3_bo (d loss / d block +/// output) in, b3_bo (d loss / d block input) out +fn tr3_block_backward(b: usize, lo: usize) -> void { + const bw = &model.blocks[b]; + const eo : usize = b * EMBED_DIM; + const fo : usize = b * FF_DIM; + const eb : usize = b * CONTEXT_LEN * EMBED_DIM; + const fb : usize = b * CONTEXT_LEN * FF_DIM; + const span : usize = CONTEXT_LEN * EMBED_DIM; + const fspan : usize = CONTEXT_LEN * FF_DIM; + const ge : usize = b * EMBED_DIM * EMBED_DIM; + const gf : usize = b * EMBED_DIM * FF_DIM; + const hi : usize = CONTEXT_LEN; + var p : usize = lo; + while (p < hi) { + const ro : usize = p * EMBED_DIM; + var i : usize = 0; + while (i < EMBED_DIM) { + g_b_down[eo + i] = g_b_down[eo + i] + b3_bo[ro + i]; + b3_se[ro + i] = bw.alpha_down * b3_bo[ro + i]; + i = i + 1; + } + p = p + 1; + } + tr3_outer(b3_h[fb..fb + fspan], &b3_se, gw_down[gf..gf + FF_DIM * EMBED_DIM], FF_DIM, EMBED_DIM, lo, hi); + tr3_zero(&b3_bh, FF_DIM, lo, hi); + tr3_t_add(&b3_se, &bw.w_down, &b3_bh, FF_DIM, EMBED_DIM, lo, hi); + p = lo; + while (p < hi) { + const rf : usize = p * FF_DIM; + var i : usize = 0; + while (i < FF_DIM) { + b3_bpre[rf + i] = relu_derivative(b3_pre[fb + rf + i]) * b3_bh[rf + i]; + g_b_up[fo + i] = g_b_up[fo + i] + b3_bpre[rf + i]; + b3_sf[rf + i] = bw.alpha_up * b3_bpre[rf + i]; + i = i + 1; + } + p = p + 1; + } + tr3_outer(b3_r[eb..eb + span], &b3_sf, gw_up[gf..gf + EMBED_DIM * FF_DIM], EMBED_DIM, FF_DIM, lo, hi); + p = lo; + while (p < hi) { + const ro : usize = p * EMBED_DIM; + var i : usize = 0; + while (i < EMBED_DIM) { + b3_br[ro + i] = b3_bo[ro + i]; + i = i + 1; + } + p = p + 1; + } + tr3_t_add(&b3_sf, &bw.w_up, &b3_br, EMBED_DIM, FF_DIM, lo, hi); + p = lo; + while (p < hi) { + const ro : usize = p * EMBED_DIM; + var i : usize = 0; + while (i < EMBED_DIM) { + b3_se[ro + i] = bw.alpha_o * b3_br[ro + i]; + i = i + 1; + } + p = p + 1; + } + tr3_outer(b3_ctx[eb..eb + span], &b3_se, gw_o[ge..ge + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM, lo, hi); + tr3_zero(&b3_bctx, EMBED_DIM, lo, hi); + tr3_t_add(&b3_se, &bw.w_o, &b3_bctx, EMBED_DIM, EMBED_DIM, lo, hi); + p = lo; + while (p < hi) { + tr3_attend_backward(b, p); + p = p + 1; + } + const ns = b3_n1[eb..eb + span]; + tr3_outer(ns, &b3_bq, gw_q[ge..ge + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM, lo, hi); + tr3_outer(ns, &b3_bk, gw_k[ge..ge + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM, lo, hi); + tr3_outer(ns, &b3_bv, gw_v[ge..ge + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM, lo, hi); + tr3_zero(&b3_bn, EMBED_DIM, lo, hi); + tr3_t_add(&b3_bq, &bw.w_q, &b3_bn, EMBED_DIM, EMBED_DIM, lo, hi); + tr3_t_add(&b3_bk, &bw.w_k, &b3_bn, EMBED_DIM, EMBED_DIM, lo, hi); + tr3_t_add(&b3_bv, &bw.w_v, &b3_bn, EMBED_DIM, EMBED_DIM, lo, hi); + p = lo; + while (p < hi) { + const ro : usize = p * EMBED_DIM; + const o : usize = eb + ro; + rms_norm_backward(b3_x[o..o + EMBED_DIM], &bw.norm1_gamma, b3_rms1[b * CONTEXT_LEN + p], b3_bn[ro..ro + EMBED_DIM], b3_tmp[ro..ro + EMBED_DIM], g_gamma1[eo..eo + EMBED_DIM]); + var i : usize = 0; + while (i < EMBED_DIM) { + b3_bo[ro + i] = b3_br[ro + i] + b3_tmp[ro + i]; + i = i + 1; + } + p = p + 1; + } +} + +/// tr_window_grad_k at speed 3: the forward of the whole window, then the cross-entropy of each +/// scored position in order, then one backward over all of them; the last position's logits, +/// d loss / d logits and input gradient land where speed 2 leaves them +fn tr3_window_grad(tokens: []const u16, start: usize, k: usize) -> f64 { + const inv_k : f64 = 1.0 / @as(f64, @floatFromInt(k)); + const lo : usize = CONTEXT_LEN - k; + tr3_forward(tokens, start, CONTEXT_LEN, lo); + var sum : f64 = 0.0; + var p : usize = lo; + while (p < CONTEXT_LEN) { + const target : usize = @intCast(tokens[start + p + 1]); + const lo2 : usize = p * VOCAB_SIZE; + sum = sum + cross_entropy(b3_logits[lo2..lo2 + VOCAB_SIZE], target); + cross_entropy_backward(b3_logits[lo2..lo2 + VOCAB_SIZE], target, b3_dl[lo2..lo2 + VOCAB_SIZE]); + var j : usize = 0; + while (j < VOCAB_SIZE) { + b3_dl[lo2 + j] = b3_dl[lo2 + j] * inv_k; + j = j + 1; + } + p = p + 1; + } + p = lo; + while (p < CONTEXT_LEN) { + const lo2 : usize = p * VOCAB_SIZE; + var j : usize = 0; + while (j < VOCAB_SIZE) { + g_b_out[j] = g_b_out[j] + b3_dl[lo2 + j]; + b3_dz[lo2 + j] = model.alpha_out * SACRED_LOGIT_SCALE * b3_dl[lo2 + j]; + j = j + 1; + } + p = p + 1; + } + tr3_outer(&b3_hn, &b3_dz, &gw_out, EMBED_DIM, VOCAB_SIZE, lo, CONTEXT_LEN); + tr3_zero(&b3_bn, EMBED_DIM, lo, CONTEXT_LEN); + tr3_t_add(&b3_dz, &model.w_out, &b3_bn, EMBED_DIM, VOCAB_SIZE, lo, CONTEXT_LEN); + const lb : usize = (NUM_BLOCKS - 1) * CONTEXT_LEN * EMBED_DIM; + p = lo; + while (p < CONTEXT_LEN) { + const ro : usize = p * EMBED_DIM; + rms_norm_plain_backward(b3_out[lb + ro..lb + ro + EMBED_DIM], b3_hrms[p], b3_bn[ro..ro + EMBED_DIM], b3_bo[ro..ro + EMBED_DIM]); + p = p + 1; + } + var b : usize = NUM_BLOCKS; + while (b > 0) { + b = b - 1; + tr3_block_backward(b, lo); + } + p = lo; + while (p < CONTEXT_LEN) { + const ro : usize = p * EMBED_DIM; + const eo : usize = b3_tok[p] * EMBED_DIM; + var i : usize = 0; + while (i < EMBED_DIM) { + g_embed[eo + i] = g_embed[eo + i] + b3_bo[ro + i]; + i = i + 1; + } + p = p + 1; + } + const last : usize = CONTEXT_LEN - 1; + var i : usize = 0; + while (i < EMBED_DIM) { + win_gin[i] = b3_bo[last * EMBED_DIM + i]; + i = i + 1; + } + var j : usize = 0; + while (j < VOCAB_SIZE) { + win_logits[j] = b3_logits[last * VOCAB_SIZE + j]; + win_dl[j] = b3_dl[last * VOCAB_SIZE + j]; + j = j + 1; + } + return sum * inv_k; +} + ; --------------------------------------------------------------------------------------------- ; 2. Parameters: the flat vector, its segments, init, requantize ; --------------------------------------------------------------------------------------------- @@ -1767,7 +2403,7 @@ fn tr_eval_line(step: u64, have_train: bool, avg_loss: f64, valid: f64, valid_fu fn tr_main() -> i32 { const t_begin : f64 = tr_now(); tr_read_cfg(); - speed_level = @intCast(cfg_speed % 3); + speed_level = @intCast(cfg_speed % 4); if (cfg_gradcheck > 0) { return tr_gradcheck(); } @@ -2064,6 +2700,134 @@ fn tr_speed_grad_reldiff() -> f64 { return sqrt(num / den); } +var s3_grad : [NPARAM]f64 = undefined; +var s3_gin : [EMBED_DIM]f64 = undefined; + +/// 1 if the two floats differ in any bit, else 0 +fn tr_bits_differ(a: f64, b: f64) -> usize { + const x : u64 = @bitCast(a); + const y : u64 = @bitCast(b); + if (x != y) { + return 1; + } + return 0; +} + +/// The logits of all 27 positions of one window, at speed 2 (hslm_forward at each position) and +/// at speed 3 (one tr3_forward with the head at every position): the number of logits whose bits +/// differ +fn tr_speed3_logit_bitdiff() -> usize { + tr_tiny_setup(5); + speed_level = 3; + tr3_forward(train_tok[0..ntrain], 700, CONTEXT_LEN, 0); + speed_level = 2; + var diff : usize = 0; + var p : usize = 0; + while (p < CONTEXT_LEN) { + hslm_forward(@as(usize, @intCast(train_tok[700 + p])), p, &ref_logits); + var j : usize = 0; + while (j < VOCAB_SIZE) { + diff = diff + tr_bits_differ(ref_logits[j], b3_logits[p * VOCAB_SIZE + j]); + j = j + 1; + } + p = p + 1; + } + return diff; +} + +/// One window of the k-position objective at speed 2 and at speed 3 from the same model: the +/// number of values whose bits differ among every gradient (all NPARAM of them, gathered unscaled), +/// the input gradient, the last position's logits and the loss +fn tr_speed3_grad_bitdiff(k: usize) -> usize { + tr_tiny_setup(5); + speed_level = 2; + zero_weight_gradients(); + const l2 : f64 = tr_window_grad_k(train_tok[0..ntrain], 700, k); + tr_gather(1.0); + var i : usize = 0; + while (i < NPARAM) { + s3_grad[i] = g_val[i]; + i = i + 1; + } + i = 0; + while (i < EMBED_DIM) { + s3_gin[i] = win_gin[i]; + i = i + 1; + } + var j : usize = 0; + while (j < VOCAB_SIZE) { + ref_logits[j] = win_logits[j]; + j = j + 1; + } + speed_level = 3; + zero_weight_gradients(); + const l3 : f64 = tr_window_grad_k(train_tok[0..ntrain], 700, k); + tr_gather(1.0); + speed_level = 2; + var diff : usize = tr_bits_differ(l2, l3); + i = 0; + while (i < NPARAM) { + diff = diff + tr_bits_differ(s3_grad[i], g_val[i]); + i = i + 1; + } + i = 0; + while (i < EMBED_DIM) { + diff = diff + tr_bits_differ(s3_gin[i], win_gin[i]); + i = i + 1; + } + j = 0; + while (j < VOCAB_SIZE) { + diff = diff + tr_bits_differ(ref_logits[j], win_logits[j]); + j = j + 1; + } + return diff; +} + +/// How many gradient values the speed 2 window of tr_speed3_grad_bitdiff left nonzero (the +/// comparison is not one of zeros) +fn tr_speed3_nonzero() -> usize { + var n : usize = 0; + var i : usize = 0; + while (i < NPARAM) { + if (s3_grad[i] != 0.0) { + n = n + 1; + } + i = i + 1; + } + return n; +} + +/// The forward-only k-position loss of one window at speed 2 and at speed 3: 1 if the bits differ +fn tr_speed3_loss_bitdiff(k: usize) -> usize { + tr_tiny_setup(5); + speed_level = 2; + const l2 : f64 = tr_window_loss_k(train_tok[0..ntrain], 900, k); + speed_level = 3; + const l3 : f64 = tr_window_loss_k(train_tok[0..ntrain], 900, k); + speed_level = 2; + return tr_bits_differ(l2, l3); +} + +/// The largest gap between a recorded speed 2 / speed 3 tok/s and the one its train_s gives +fn tr_speed3_toks_gap() -> f64 { + const a : f64 = abs(TR3_LAB_TOKENS / TR3_LAB_S2_TRAIN_S - TR3_LAB_S2_TOK_S); + const b : f64 = abs(TR3_LAB_TOKENS / TR3_LAB_S3_TRAIN_S - TR3_LAB_S3_TOK_S); + if (a > b) { + return a; + } + return b; +} + +/// |recorded speed-up - the ratio of the recorded tok/s| +fn tr_speed3_ratio_gap() -> f64 { + return abs(TR3_LAB_S3_TOK_S / TR3_LAB_S2_TOK_S - TR3_LAB_SPEEDUP); +} + +/// The speed-up of the second pair of runs, from its train_s +fn tr_speed3_speedup_again() -> f64 { + return TR3_LAB_S2_TRAIN_S_AGAIN / TR3_LAB_S3_TRAIN_S_AGAIN; +} + /// Bytes that differ between an MSLH image and the image packed again from its own decoded floats fn tr_mslh_roundtrip_diff() -> usize { tr_init(3); @@ -2596,6 +3360,40 @@ test speed_two_input_gradient_matches_the_reference given rel = tr_speed_grad_reldiff() then rel < 1e-12 +test speed_three_gives_the_reference_logits_to_the_bit + given diff = tr_speed_logit_diff(3) + then diff == 0.0 + +test speed_three_gives_the_logits_of_every_position_of_speed_two_to_the_bit + given diff = tr_speed3_logit_bitdiff() + then diff == 0 + +test speed_three_gives_every_gradient_of_speed_two_to_the_bit_at_k_27 + given diff = tr_speed3_grad_bitdiff(27) + and nonzero = tr_speed3_nonzero() + then diff == 0 + and nonzero > 1000000 + +test speed_three_gives_every_gradient_of_speed_two_to_the_bit_at_k_1 + given diff = tr_speed3_grad_bitdiff(1) + and nonzero = tr_speed3_nonzero() + then diff == 0 + and nonzero > 100000 + +test speed_three_gives_the_k_position_loss_of_speed_two_to_the_bit + given d27 = tr_speed3_loss_bitdiff(27) + and d1 = tr_speed3_loss_bitdiff(1) + then d27 == 0 + and d1 == 0 + +test speed_three_lab_tok_s_and_its_speed_up_follow_from_the_measured_times + given toks = tr_speed3_toks_gap() + and ratio = tr_speed3_ratio_gap() + and again = tr_speed3_speedup_again() + then toks < 0.05 + and ratio < 0.0005 + and again > 1.4 + test a_few_steps_on_a_tiny_fixed_slice_lower_the_loss given after = tr_tiny_after() and before = tr_tiny_before()