diff --git a/specs/igla/own_model.t27 b/specs/igla/own_model.t27 index 5a4fb11fe1..a65ed59a79 100644 --- a/specs/igla/own_model.t27 +++ b/specs/igla/own_model.t27 @@ -31,13 +31,22 @@ ; vocab-729 head and cross-entropy), its backward reaches the embedding, and every gradient is ; checked against a central difference; a mutant that drops one alpha fails 10 of the 46. ; Not in it: the per-block VSA path, f32, and the trainer's STE update (shadow weights). +; - test on 2026-10-18, same t27c, after the PR that added the training step (STE, backward, data): +; nn/ternary_ste 26/26, nn/hslm_backward 15/15, nn/hslm_data 39/39 pass, 0 vacuous. The STE is +; the TWN quantizer (delta 0.7 * mean abs, alpha = mean abs of the kept weights), the +; pass-through gradient, clip_grad_norm and the SGD step on the shadow weights. hslm_backward +; runs forward, backward, STE, clip, SGD and requantize on one 3-head slice of the model +; (ctx 27, head_dim 3): its trit, bias and input gradients match a central difference to 1e-6 +; and 20 steps cut the loss below half. hslm_data holds the curriculum, next-token batches and +; the MSLH header and layout, whose byte count equals hslm_step_10000.bin (4971796). Not in +; them: the STE on the whole 3-block model, file I/O, heap, the tokenizer, shuffle, AdamW. ; - live: GET huggingface.co/api/models/playra/tern-tc-9m and .../tern-tc-9m-t27 returned 200; ; .../tern-tc-100m returned 401, the same as a made-up name, so no 100M is public. ; GET board-production-2dec.up.railway.app/state returned 200 and a made-up path 404; its ; snapshot is dated 2026-10-05, so the board is live and 12 days stale. ; - review: no PR behind any cell has a review, so no DELIVER cell is VERIFIED. ; The engine asks an engineering LEARN step for a source link, a PRACTICE step for a passing test -; and a DELIVER step for a review. So: 1 published, 12 verified, 7 drafts, 7 gaps. +; and a DELIVER step for a review. So: 1 published, 12 verified, 9 drafts, 5 gaps. ; PERCENT is a second, separate number per cell: an estimate of how much of that part exists, with ; its basis in SOURCES. It never feeds a status. The lens and overall figures are plain means, every @@ -46,8 +55,9 @@ ; THE ANSWER this map gives: HSLM was trained before, in Zig only (64 files); its best runs ; (PPL 2.96 at ctx 27) were wiped by EXP-002 and survive only in the ledger. In t27 the loss, three ; optimizers and the HSLM spec's forward and backward passes pass, and that spec is now the trained -; model but for its VSA path (see MEASURED); ternary STE with shadow weights, data loader, -; checkpoint I/O, heap and file I/O and `t27c --backend ml` do not exist yet. +; model but for its VSA path (see MEASURED); the ternary STE, the training step and the data and +; checkpoint layout now pass too, on a model slice and on bytes; the STE over the whole model, +; heap and file I/O and `t27c --backend ml` do not exist yet. ; Not claimed: no model quality beyond the numbers in SOURCES; no GIF or MP4, stills only. ; phi^2 + 1/phi^2 = 3 | TRINITY @@ -63,6 +73,7 @@ use t27f::topic_engine::STATUS_DRAFT; use t27f::topic_engine::STATUS_VERIFIED; use t27f::topic_engine::STATUS_PUBLISHED; use t27f::topic_engine::STAGE_PRACTICE; +use t27f::topic_engine::STAGE_DELIVER; use t27f::topic_engine::DOMAIN_ENGINEERING; use t27f::topic_engine::CONTENT_WORDS; use t27f::topic_engine::ENGINEERING_WORDS; @@ -144,9 +155,9 @@ pub const SOURCES : [27]str = [ "t27 specs/ml/loss/cross_entropy_loss.t27", "t27 specs/ml/optimizer/{sgd,adam,lamb}.t27", "t27 specs/nn/hslm.t27 and specs/hslm/forward_pass.t27", - "t27 specs/nn/hslm.t27 backward, finite-difference checked", - "", - "", + "t27 specs/nn/hslm.t27 backward and nn/hslm_backward.t27 training step on a head slice, finite-difference checked", + "t27 specs/nn/ternary_ste.t27, 26 tests; whole-model shadow weights not yet", + "t27 specs/nn/hslm_data.t27: batches, curriculum, MSLH layout = 4971796 bytes; no file I/O", "igla-coder-gpu research/MODEL_ON_FPGA_PLAN.md", "trinity-fpga b2bd5a30, 7682304/7682304 receipts", "board-production-2dec.up.railway.app/state", @@ -164,8 +175,8 @@ pub const MISSING : [27]str = [ "a public 100M checkpoint: playra/tern-tc-100m answers 401, like a made-up name", "any run past 100M parameters: none exists", "", "", "", "", "", "", "", - "a ternary straight-through estimator in any .t27 spec: none exists", - "a t27 reader for token streams and MSLH checkpoints: none exists", + "", + "", "", "", "", "", "", "heap allocation and file I/O lowered by t27c: neither exists", "a t27c backend that lowers a model spec to a training binary: none exists", @@ -174,17 +185,17 @@ pub const MISSING : [27]str = [ ]; /// The measurements per cell (see the header for how each was taken). -pub const GROUNDED : [27]bool = [true, true, true, true, true, true, true, false, false, true, true, true, true, true, true, true, false, false, true, true, true, true, true, false, false, false, true]; +pub const GROUNDED : [27]bool = [true, true, true, true, true, true, true, false, false, true, true, true, true, true, true, true, true, true, true, true, true, true, true, false, false, false, true]; pub const TEST_PASS : [27]bool = [false, false, false, false, false, false, false, false, false, false, false, false, true, true, true, false, false, false, false, false, false, true, false, false, false, false, false]; pub const REVIEWED : [27]bool = [false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false]; pub const LIVE : [27]bool = [false, false, false, false, false, false, true, false, false, false, false, false, false, false, false, false, false, false, false, false, true, false, false, false, false, false, false]; /// What a host passes the viewer, per cell: the earned status code and the step number as a label. -pub const STATUS : [27]u8 = [4, 4, 4, 3, 3, 3, 3, 0, 0, 4, 4, 4, 4, 4, 4, 3, 0, 0, 4, 4, 5, 4, 3, 0, 0, 0, 3]; +pub const STATUS : [27]u8 = [4, 4, 4, 3, 3, 3, 3, 0, 0, 4, 4, 4, 4, 4, 4, 3, 3, 3, 4, 4, 5, 4, 3, 0, 0, 0, 3]; pub const LABELS : [27]str = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15", "16", "17", "18", "19", "20", "21", "22", "23", "24", "25", "26", "27"]; /// Percent of each part that exists, estimated from its source (see the header). Never a status. -pub const PCT : [27]u8 = [60, 70, 80, 80, 60, 10, 50, 0, 0, 100, 100, 20, 80, 65, 85, 70, 0, 0, 70, 15, 60, 40, 35, 0, 0, 0, 50]; +pub const PCT : [27]u8 = [60, 70, 80, 80, 60, 10, 50, 0, 0, 100, 100, 20, 80, 65, 85, 75, 55, 40, 70, 15, 60, 40, 35, 0, 0, 0, 50]; /// The t27c test-report counts the PRACTICE cells quote. pub const ML_SPECS_MEASURED : i32 = 17; @@ -209,11 +220,11 @@ pub const ROWS : [20]str = [ "In t27: loss, optimizers, HSLM pass", "6 of 17 ML specs compile on the lab", "0 of 9 IGLA CODER specs compile", - "1 published, 12 verified, 7 drafts", - "7 gaps, about 44 percent overall", + "1 published, 12 verified, 9 drafts", + "5 gaps, about 48 percent overall", "HSLM spec: add the VSA path", - "Ternary STE with shadow weights", - "Loader, checkpoint I/O, --backend ml", + "Shadow weights for the whole model", + "File I/O, heap, --backend ml", "" ]; @@ -266,7 +277,7 @@ test every_step_is_a_cell { test a_gap_is_named_not_filled { assert gaps_are_named(SOURCES, MISSING, GROUNDED); - assert count_true(GROUNDED) == 20; + assert count_true(GROUNDED) == 22; var c : i32 = 0; while (c < CELL_COUNT) { const i : usize = @intCast(c); @@ -280,11 +291,11 @@ test a_gap_is_named_not_filled { test a_status_is_earned_not_chosen { assert statuses_earned(STATUS, DOMAIN, GROUNDED, TEST_PASS, REVIEWED, LIVE); - // the evidence slide's "1 published, 12 verified, 7 drafts" and "7 gaps" + // the evidence slide's "1 published, 12 verified, 9 drafts" and "5 gaps" assert count_status(STATUS, STATUS_PUBLISHED) == 1; assert count_status(STATUS, STATUS_VERIFIED) == 12; - assert count_status(STATUS, STATUS_DRAFT) == 7; - assert count_status(STATUS, STATUS_GAP) == 7; + assert count_status(STATUS, STATUS_DRAFT) == 9; + assert count_status(STATUS, STATUS_GAP) == 5; assert count_true(REVIEWED) == 0; // every passing test is a PRACTICE step: loss, optimizers, forward pass, matvec assert count_true(TEST_PASS) == 4; @@ -319,10 +330,24 @@ test the_compile_counts_are_the_measured_ones { test percent_per_lens_is_a_plain_mean { assert mean_pct(0, 9) == 45; - assert mean_pct(9, 9) == 57; + assert mean_pct(9, 9) == 68; assert mean_pct(18, 9) == 30; - // the evidence slide's "about 44 percent overall" - assert mean_pct(0, CELL_COUNT) == 44; + // the evidence slide's "about 48 percent overall" + assert mean_pct(0, CELL_COUNT) == 48; +} + +test the_training_step_cells_are_drafts_until_reviewed { + // steps 16-18 are DELIVER steps: their specs pass (see MEASURED) but the engine asks a review + // for them, and none exists, so they are grounded drafts, not verified and not gaps + var c : i32 = 15; + while (c < 18) { + const i : usize = @intCast(c); + assert GROUNDED[i]; + assert STATUS[i] == STATUS_DRAFT; + assert REVIEWED[i] == false; + assert step_stage(@intCast(c)) == STAGE_DELIVER; + c = c + 1; + } } test the_camera_is_the_demo_camera { diff --git a/specs/igla/own_model_board.t27 b/specs/igla/own_model_board.t27 index 6a5d0cf671..bf9745483b 100644 --- a/specs/igla/own_model_board.t27 +++ b/specs/igla/own_model_board.t27 @@ -48,23 +48,23 @@ pub const LENS_CELLS : i32 = 9; pub const JOB : str = "own_model"; pub const STAGE : str = "own_model"; -pub const MSG : str = "1 published, 12 verified, 7 drafts, 7 gaps; 44%"; +pub const MSG : str = "1 published, 12 verified, 9 drafts, 5 gaps; 48%"; /// What MSG says once its numbers are taken out. pub const MSG_WORDS : str = " published, verified, drafts, gaps; %"; -pub const OVERALL_PCT : u8 = 44; -pub const PROGRESS : f64 = 0.44; +pub const OVERALL_PCT : u8 = 48; +pub const PROGRESS : f64 = 0.48; pub const KEY_LENS : str = "lens_pct"; pub const KEY_CELLS : str = "cells"; pub const KEY_STATUS : str = "status"; pub const KEY_PCT : str = "pct"; -pub const LENS_PCT : [3]u8 = [45, 57, 30]; -pub const CELL_STATUS : [27]u8 = [4, 4, 4, 3, 3, 3, 3, 0, 0, 4, 4, 4, 4, 4, 4, 3, 0, 0, 4, 4, 5, 4, 3, 0, 0, 0, 3]; -pub const CELL_PCT : [27]u8 = [60, 70, 80, 80, 60, 10, 50, 0, 0, 100, 100, 20, 80, 65, 85, 70, 0, 0, 70, 15, 60, 40, 35, 0, 0, 0, 50]; +pub const LENS_PCT : [3]u8 = [45, 68, 30]; +pub const CELL_STATUS : [27]u8 = [4, 4, 4, 3, 3, 3, 3, 0, 0, 4, 4, 4, 4, 4, 4, 3, 3, 3, 4, 4, 5, 4, 3, 0, 0, 0, 3]; +pub const CELL_PCT : [27]u8 = [60, 70, 80, 80, 60, 10, 50, 0, 0, 100, 100, 20, 80, 65, 85, 75, 55, 40, 70, 15, 60, 40, 35, 0, 0, 0, 50]; /// The size of the compact metrics object, in bytes, computed from the keys and values above. -pub const METRICS_BYTES : i32 = 622; +pub const METRICS_BYTES : i32 = 624; pub fn is_digit(ch: u8) -> bool { return ch >= '0' and ch <= '9'; @@ -262,8 +262,8 @@ test the_tally_sums_to_twenty_seven { assert published + verified + drafts + gaps == CELL_COUNT; assert published == 1; assert verified == 12; - assert drafts == 7; - assert gaps == 7; + assert drafts == 9; + assert gaps == 5; } test the_message_is_the_tally { @@ -273,12 +273,12 @@ test the_message_is_the_tally { const n : i32 = @intCast(MSG.len); assert n <= MSG_MAX; // one number wrong, a word changed, or a number added: each refused - assert msg_matches("1 published, 11 verified, 7 drafts, 7 gaps; 44%", STATUS, overall) == false; - assert msg_matches("1 published, 12 verified, 7 drafts, 7 holes; 44%", STATUS, overall) == false; - assert msg_matches("1 published, 12 verified, 7 drafts, 7 gaps; 44% of 100", STATUS, overall) == false; + assert msg_matches("1 published, 11 verified, 9 drafts, 5 gaps; 48%", STATUS, overall) == false; + assert msg_matches("1 published, 12 verified, 9 drafts, 5 holes; 48%", STATUS, overall) == false; + assert msg_matches("1 published, 12 verified, 9 drafts, 5 gaps; 48% of 100", STATUS, overall) == false; assert msg_matches(MSG, STATUS, 45) == false; assert nth_number(MSG, 1) == 12; - assert nth_number(MSG, 4) == 44; + assert nth_number(MSG, 4) == 48; assert nth_number(MSG, 5) == -1; } diff --git a/specs/nn/hslm_backward.t27 b/specs/nn/hslm_backward.t27 new file mode 100644 index 0000000000..ecbc771229 --- /dev/null +++ b/specs/nn/hslm_backward.t27 @@ -0,0 +1,583 @@ +// SPDX-License-Identifier: Apache-2.0 +// t27/specs/nn/hslm_backward.t27 +// One HSLM training step on shadow weights: the output head of specs/nn/hslm.t27 (a ternary +// matrix with a TWN alpha, a bias and cross-entropy), trained the way trinity-training trains +// it. Forward on trits, backward straight-through to the float shadow weights, gradient clip, +// SGD step on the shadow weights, requantize. The backward is checked against a central finite +// difference, and the loss is shown to fall. +// phi^2 + 1/phi^2 = 3 | TRINITY +// +// Source of truth: gHashTag/trinity-training @ 3e9648c1817a97c5f9b8c727b0f8e2291928eead +// src/hslm/ste.zig the quantizers (here: nn/ternary_ste.t27) +// src/hslm/autograd.zig clipGradNorm +// src/hslm/trainer.zig train step: forward, loss, backward, clip, optimizer on the shadow +// weights, requantize +// +// What is in it and what is not, so nobody reads it as more: +// - It is the head only: 9 inputs, 27 classes, 243 trits. The three transformer blocks +// already have their backward in nn/hslm.t27 (46/46 pass, every gradient checked there). A +// step through all of them needs a shadow copy of ~2 M weights next to the trits; this spec +// does not claim that. It claims the step is right on the layer shape the blocks are built +// of: ternary matmul with alpha, a bias, softmax cross-entropy. +// - The kernels (ternary_matmul, ternary_matmul_t_add, outer_add, cross_entropy) are the ones +// in nn/hslm.t27, copied here because that module carries 2 M trits of state that a tiny +// step does not need to load. If one changes, the other must. +// - The optimizer is plain SGD with no momentum, the trainer uses AdamW. +// - The gradient is d loss / d trit, as in nn/hslm.t27 (gw = alpha * x (x) dlogits), and it is +// handed to the shadow weight unchanged: the straight-through estimator. +// - Resonance Law: ctx 27 = 3^3, three heads of head_dim 3, which divides ctx; the model width +// is 3 * 3 = 9. + +module HslmBackward { + use base::types; + use nn::ternary_ste; + + // ═════════════════════════════════════════════════════════════════ + // 1. Shape + // ═════════════════════════════════════════════════════════════════ + + const CTX : usize = 27; // 3^3 + const NUM_HEADS : usize = 3; + const HEAD_DIM : usize = 3; // a divisor of CTX + const IN_DIM : usize = 9; // NUM_HEADS * HEAD_DIM + const OUT_DIM : usize = 27; // classes, one per context position + const N_EXAMPLES : usize = 9; // a batch: tokens 0..8 + const LOGIT_SCALE : f64 = 1.0; // nn/hslm.t27 uses 0.2734 (SACRED_LOGIT_SCALE); 1.0 keeps a 9-wide head learnable + const CLIP_NORM : f64 = 5.0; + const LEARNING_RATE : f64 = 2.0; + + // ═════════════════════════════════════════════════════════════════ + // 2. State + // ═════════════════════════════════════════════════════════════════ + + var shadow : [243]f64 = undefined; // IN_DIM * OUT_DIM float weights, row-major j * OUT_DIM + i + var trits : [243]Trit = undefined; // what the forward pass sees + var alpha : f64 = 1.0; + var bias : [27]f64 = undefined; + var gw : [243]f64 = undefined; // d loss / d trit + var gs : [243]f64 = undefined; // the gradient the shadow weights receive + var gb : [27]f64 = undefined; + var act_x : [9]f64 = undefined; + var logits : [27]f64 = undefined; + var dl : [27]f64 = undefined; // d loss / d logits + var dz : [27]f64 = undefined; // d loss / d (x @ W) + var dx : [9]f64 = undefined; // d loss / d x + + // Real-valued nudges for the gradient checks: a trit, a bias entry, an input entry + var nudge_index : usize = 0; + var nudge_delta : f64 = 0.0; + var x_index : usize = 0; + var x_delta : f64 = 0.0; + + // ═════════════════════════════════════════════════════════════════ + // 3. Kernels (as in nn/hslm.t27) + // ═════════════════════════════════════════════════════════════════ + + fn fill(x: []f64, value: f64) -> void { + var i : usize = 0; + while (i < x.len) { + x[i] = value; + i = i + 1; + } + } + + // output[i] = sum_j input[j] * w[j * out_dim + i], trits as -1 / 0 / +1 + fn ternary_matmul(input: []const f64, weights: []const Trit, output: []f64, in_dim: usize, out_dim: usize) -> void { + var i : usize = 0; + while (i < out_dim) { + var acc : f64 = 0.0; + var j : usize = 0; + while (j < in_dim) { + const w = weights[j * out_dim + i]; + if (w == Trit.pos) { + acc = acc + input[j]; + } else if (w == Trit.neg) { + acc = acc - input[j]; + } + j = j + 1; + } + output[i] = acc; + i = i + 1; + } + } + + // The transpose, accumulated: dx[j] += sum_i w[j * out_dim + i] * dy[i] + fn ternary_matmul_t_add(dy: []const f64, weights: []const Trit, dx_out: []f64, in_dim: usize, out_dim: usize) -> void { + var j : usize = 0; + while (j < in_dim) { + var acc : f64 = 0.0; + var i : usize = 0; + while (i < out_dim) { + const w = weights[j * out_dim + i]; + if (w == Trit.pos) { + acc = acc + dy[i]; + } else if (w == Trit.neg) { + acc = acc - dy[i]; + } + i = i + 1; + } + dx_out[j] = dx_out[j] + acc; + j = j + 1; + } + } + + // Weight gradient of a matmul: gw[j * out_dim + i] += x[j] * dy[i] + fn outer_add(x: []const f64, dy: []const f64, g: []f64, in_dim: usize, out_dim: usize) -> void { + var j : usize = 0; + while (j < in_dim) { + var i : usize = 0; + while (i < out_dim) { + g[j * out_dim + i] = g[j * out_dim + i] + x[j] * dy[i]; + i = i + 1; + } + j = j + 1; + } + } + + // loss = log(sum_j exp(l_j)) - l_target + fn cross_entropy(l: []const f64, target: usize) -> f64 { + var max_logit : f64 = l[0]; + var j : usize = 1; + while (j < l.len) { + if (l[j] > max_logit) { + max_logit = l[j]; + } + j = j + 1; + } + var sum : f64 = 0.0; + j = 0; + while (j < l.len) { + sum = sum + exp(l[j] - max_logit); + j = j + 1; + } + return log(sum) + max_logit - l[target]; + } + + // d loss / d l_j = softmax_j - [j == target] + fn cross_entropy_backward(l: []const f64, target: usize, d: []f64) -> void { + var max_logit : f64 = l[0]; + var j : usize = 1; + while (j < l.len) { + if (l[j] > max_logit) { + max_logit = l[j]; + } + j = j + 1; + } + var sum : f64 = 0.0; + j = 0; + while (j < l.len) { + d[j] = exp(l[j] - max_logit); + sum = sum + d[j]; + j = j + 1; + } + j = 0; + while (j < l.len) { + d[j] = d[j] / sum; + j = j + 1; + } + d[target] = d[target] - 1.0; + } + + // ═════════════════════════════════════════════════════════════════ + // 4. Model: init, forward, backward, train step + // ═════════════════════════════════════════════════════════════════ + + // A small signed pattern in [-scale/2, scale/2); a multiplicative hash mod a prime, as in + // nn/hslm.t27, so no row or column repeats + fn small(seed: usize, scale: f64) -> f64 { + const m = ((seed * 2654435761) % 1000003) % 13; + return scale * (@as(f64, @floatFromInt(m)) / 13.0 - 0.5); + } + + // A token's input vector: nine distinct small floats, a stand-in for an embedding row + fn embed(token: usize, out: []f64) -> void { + var j : usize = 0; + while (j < IN_DIM) { + out[j] = small(token * IN_DIM + j + 1, 4.0); + j = j + 1; + } + } + + // The class a token has to map to: a permutation of 0..26 (5 is coprime to 27) + fn target_of(token: usize) -> usize { + return (token * 5 + 2) % OUT_DIM; + } + + fn requantize_all() -> void { + alpha = ternary_ste::requantize(&shadow, &trits); + } + + fn init_model() -> void { + var k : usize = 0; + while (k < IN_DIM * OUT_DIM) { + shadow[k] = small(k + 7, 0.6); + k = k + 1; + } + fill(&bias, 0.0); + nudge_delta = 0.0; + x_delta = 0.0; + requantize_all(); + } + + // logits = alpha * (x @ W_trits) * LOGIT_SCALE + bias + fn forward(token: usize) -> void { + embed(token, &act_x); + act_x[x_index] = act_x[x_index] + x_delta; + ternary_matmul(&act_x, &trits, &logits, IN_DIM, OUT_DIM); + // the trit nudge: the output moves as if trit `nudge_index` were one plus nudge_delta + const nj = nudge_index / OUT_DIM; + const ni = nudge_index % OUT_DIM; + logits[ni] = logits[ni] + nudge_delta * act_x[nj]; + var i : usize = 0; + while (i < OUT_DIM) { + logits[i] = alpha * logits[i] * LOGIT_SCALE + bias[i]; + i = i + 1; + } + } + + fn zero_gradients() -> void { + fill(&gw, 0.0); + fill(&gs, 0.0); + fill(&gb, 0.0); + } + + // Adds one example's gradients to gw, gb; dx is this example's input gradient + fn backward(target: usize) -> void { + cross_entropy_backward(&logits, target, &dl); + var i : usize = 0; + while (i < OUT_DIM) { + gb[i] = gb[i] + dl[i]; + dz[i] = alpha * LOGIT_SCALE * dl[i]; + i = i + 1; + } + outer_add(&act_x, &dz, &gw, IN_DIM, OUT_DIM); + fill(&dx, 0.0); + ternary_matmul_t_add(&dz, &trits, &dx, IN_DIM, OUT_DIM); + } + + // Mean loss over the batch, no gradients + fn mean_loss() -> f64 { + var total : f64 = 0.0; + var e : usize = 0; + while (e < N_EXAMPLES) { + forward(e); + total = total + cross_entropy(&logits, target_of(e)); + e = e + 1; + } + return total / @as(f64, @floatFromInt(N_EXAMPLES)); + } + + fn correct_count() -> usize { + var hits : usize = 0; + var e : usize = 0; + while (e < N_EXAMPLES) { + forward(e); + var best : usize = 0; + var j : usize = 1; + while (j < OUT_DIM) { + if (logits[j] > logits[best]) { + best = j; + } + j = j + 1; + } + if (best == target_of(e)) { + hits = hits + 1; + } + e = e + 1; + } + return hits; + } + + // One training step: returns the mean loss before it. Forward and backward over the batch, + // average, straight-through to the shadow weights, clip, SGD, requantize. + fn train_step(lr: f64) -> f64 { + zero_gradients(); + var total : f64 = 0.0; + var e : usize = 0; + while (e < N_EXAMPLES) { + forward(e); + total = total + cross_entropy(&logits, target_of(e)); + backward(target_of(e)); + e = e + 1; + } + const inv = 1.0 / @as(f64, @floatFromInt(N_EXAMPLES)); + var k : usize = 0; + while (k < IN_DIM * OUT_DIM) { + gw[k] = gw[k] * inv; + k = k + 1; + } + var i : usize = 0; + while (i < OUT_DIM) { + gb[i] = gb[i] * inv; + i = i + 1; + } + ternary_ste::ste_pass_through(&gw, &gs); + ternary_ste::clip_grad_norm(&gs, CLIP_NORM); + ternary_ste::sgd_shadow_step(&shadow, &gs, lr); + ternary_ste::sgd_shadow_step(&bias, &gb, lr); + requantize_all(); + return total * inv; + } + + // ═════════════════════════════════════════════════════════════════ + // 5. Probes + // ═════════════════════════════════════════════════════════════════ + + // d loss / d trit at idx for one example (token 3, class 7), by a central difference + fn fd_trit(idx: usize) -> f64 { + const eps = 0.00001; + nudge_index = idx; + nudge_delta = eps; + forward(3); + const up = cross_entropy(&logits, 7); + nudge_delta = 0.0 - eps; + forward(3); + const down = cross_entropy(&logits, 7); + nudge_delta = 0.0; + return (up - down) / (2.0 * eps); + } + + fn analytic_one() -> void { + init_model(); + zero_gradients(); + forward(3); + backward(7); + } + + fn trit_grad_error(idx: usize) -> f64 { + analytic_one(); + const analytic = gw[idx]; + return @abs(analytic - fd_trit(idx)); + } + + fn trit_grad_size(idx: usize) -> f64 { + analytic_one(); + return @abs(gw[idx]); + } + + fn bias_grad_error(idx: usize) -> f64 { + analytic_one(); + const analytic = gb[idx]; + const eps = 0.00001; + bias[idx] = bias[idx] + eps; + forward(3); + const up = cross_entropy(&logits, 7); + bias[idx] = bias[idx] - 2.0 * eps; + forward(3); + const down = cross_entropy(&logits, 7); + bias[idx] = bias[idx] + eps; + return @abs(analytic - (up - down) / (2.0 * eps)); + } + + fn input_grad_error(idx: usize) -> f64 { + analytic_one(); + const analytic = dx[idx]; + const eps = 0.00001; + x_index = idx; + x_delta = eps; + forward(3); + const up = cross_entropy(&logits, 7); + x_delta = 0.0 - eps; + forward(3); + const down = cross_entropy(&logits, 7); + x_delta = 0.0; + x_index = 0; + return @abs(analytic - (up - down) / (2.0 * eps)); + } + + fn input_grad_size(idx: usize) -> f64 { + analytic_one(); + return @abs(dx[idx]); + } + + // The softmax gradient sums to zero over the classes + fn dl_sum() -> f64 { + analytic_one(); + var s : f64 = 0.0; + var i : usize = 0; + while (i < OUT_DIM) { + s = s + dl[i]; + i = i + 1; + } + return s; + } + + // gw[j * OUT_DIM + i] = alpha * x[j] * dl[i], the whole gradient in one formula + fn gw_formula_error() -> f64 { + analytic_one(); + var worst : f64 = 0.0; + var j : usize = 0; + while (j < IN_DIM) { + var i : usize = 0; + while (i < OUT_DIM) { + const want = alpha * LOGIT_SCALE * act_x[j] * dl[i]; + const gap = @abs(gw[j * OUT_DIM + i] - want); + if (gap > worst) { + worst = gap; + } + i = i + 1; + } + j = j + 1; + } + return worst; + } + + fn init_loss() -> f64 { + init_model(); + return mean_loss(); + } + + fn init_hits() -> usize { + init_model(); + return correct_count(); + } + + // Loss after `steps` steps of learning rate lr + fn loss_after(steps: u32, lr: f64) -> f64 { + init_model(); + var n : u32 = 0; + while (n < steps) { + const before = train_step(lr); + n = n + 1; + } + return mean_loss(); + } + + fn hits_after(steps: u32, lr: f64) -> usize { + init_model(); + var n : u32 = 0; + while (n < steps) { + const before = train_step(lr); + n = n + 1; + } + return correct_count(); + } + + // How many of the 243 trits differ after `steps` steps from where they started + fn trits_flipped(steps: u32, lr: f64) -> usize { + init_model(); + var start : [243]Trit = undefined; + var k : usize = 0; + while (k < IN_DIM * OUT_DIM) { + start[k] = trits[k]; + k = k + 1; + } + var n : u32 = 0; + while (n < steps) { + const before = train_step(lr); + n = n + 1; + } + var flipped : usize = 0; + k = 0; + while (k < IN_DIM * OUT_DIM) { + if (trits[k] != start[k]) { + flipped = flipped + 1; + } + k = k + 1; + } + return flipped; + } + + // The clipped gradient norm after one step's worth of backward, for a tight bound + fn clipped_norm(bound: f64) -> f64 { + analytic_one(); + ternary_ste::ste_pass_through(&gw, &gs); + ternary_ste::clip_grad_norm(&gs, bound); + return ternary_ste::grad_norm(&gs); + } + + fn raw_norm() -> f64 { + analytic_one(); + return ternary_ste::grad_norm(&gw); + } + + fn nonzero_trits() -> usize { + init_model(); + return IN_DIM * OUT_DIM - ternary_ste::count_zero(&trits); + } + + // ═════════════════════════════════════════════════════════════════ + // TDD-Inside-Spec: Tests and Invariants + // ═════════════════════════════════════════════════════════════════ + + // Each backward entry against a central difference of the forward pass: a trit at the first + // row, a trit in the middle, the last trit. A zero trit is checked too: its gradient is the + // straight-through one, the slope of the loss in the trit's value. + test backward_trit_gradient_matches_finite_difference + given e0 = trit_grad_error(0) + when e1 = trit_grad_error(100) + then e0 < 0.000001 and e1 < 0.000001 and trit_grad_error(242) < 0.000001 and trit_grad_error(7) < 0.000001 + + // the check above is not about zeros: the gradients it compares are not tiny + test backward_trit_gradients_are_not_zero + given s0 = trit_grad_size(0) + then s0 > 0.00001 and trit_grad_size(100) > 0.00001 and trit_grad_size(7) > 0.00001 + + test backward_bias_gradient_matches_finite_difference + given e0 = bias_grad_error(0) + when e7 = bias_grad_error(7) + then e0 < 0.000001 and e7 < 0.000001 and bias_grad_error(26) < 0.000001 + + test backward_input_gradient_matches_finite_difference + given e0 = input_grad_error(0) + when e4 = input_grad_error(4) + then e0 < 0.000001 and e4 < 0.000001 and input_grad_error(8) < 0.000001 and input_grad_size(4) > 0.00001 + + test backward_logits_gradient_sums_to_zero + given s = dl_sum() + then @abs(s) < 0.000000000001 + + // the straight-through gradient is exactly alpha * x (x) dlogits + test backward_weight_gradient_is_alpha_x_outer_dlogits + given gap = gw_formula_error() + then gap < 0.000000000001 + + test clip_leaves_a_small_norm_alone + given n = clipped_norm(1000.0) + then @abs(n - raw_norm()) < 0.000000000001 + + test clip_pulls_a_large_norm_down_to_the_bound + given n = clipped_norm(0.01) + then @abs(n - 0.01) < 0.000000000001 and raw_norm() > 0.01 + + test init_quantization_uses_both_signs_and_zeros + given nz = nonzero_trits() + then nz > 20 and nz < 243 + + test init_loss_is_near_uniform + given l = init_loss() + then l > 2.0 and l < 5.0 + + test a_zero_learning_rate_changes_nothing + given l0 = loss_after(0, LEARNING_RATE) + when l1 = loss_after(3, 0.0) + then @abs(l0 - l1) < 0.000000000001 + + test one_step_lowers_the_loss + given l0 = loss_after(0, LEARNING_RATE) + when l1 = loss_after(1, LEARNING_RATE) + then l1 < l0 + + test training_lowers_the_loss_a_lot + given l0 = loss_after(0, LEARNING_RATE) + when l20 = loss_after(20, LEARNING_RATE) + then l20 < l0 * 0.5 + + test training_flips_trits + given f = trits_flipped(20, LEARNING_RATE) + then f > 0 + + test training_makes_predictions_right + given h0 = hits_after(0, LEARNING_RATE) + when h20 = hits_after(20, LEARNING_RATE) + then h20 > h0 and h20 >= 7 + + invariant resonance_ctx_is_three_cubed + assert CTX == 27 and CTX == 3 * 3 * 3 + + invariant resonance_head_dim_divides_ctx + assert CTX % HEAD_DIM == 0 and IN_DIM == NUM_HEADS * HEAD_DIM + + invariant batch_fits_the_context + assert N_EXAMPLES <= CTX and OUT_DIM == CTX +} diff --git a/specs/nn/hslm_data.t27 b/specs/nn/hslm_data.t27 new file mode 100644 index 0000000000..18e0539bc8 --- /dev/null +++ b/specs/nn/hslm_data.t27 @@ -0,0 +1,653 @@ +// SPDX-License-Identifier: Apache-2.0 +// t27/specs/nn/hslm_data.t27 +// The HSLM data loader and checkpoint format, in t27: next-token batches with a curriculum on +// the sequence length, and the byte layout of the MSLH checkpoint, checked against the one real +// checkpoint file that survives. +// phi^2 + 1/phi^2 = 3 | TRINITY +// +// Source of truth: gHashTag/trinity-training @ 3e9648c1817a97c5f9b8c727b0f8e2291928eead +// src/hslm/data.zig Dataset.nextBatch, Dataset.numSequences, Curriculum.updateDifficulty, +// Dataset.getCurrentSeqLen +// src/hslm/trainer.zig CHECKPOINT_MAGIC, CHECKPOINT_VERSION, saveCheckpoint, loadCheckpoint +// and the file gHashTag/trinity data/wave9/worker-2/hslm_step_10000.bin: 4971796 bytes, whose +// first 16 bytes are REAL_HEADER below. +// +// What is in it and what is not, so nobody reads it as more: +// - There is no file I/O and no heap in t27c yet (cells 24 and 25 of specs/igla/own_model.t27). +// So nothing here opens a file or allocates: the loader reads a token slice the caller +// holds and writes into batch slices the caller holds, and the checkpoint is read and written +// as a byte slice. The reading and writing of that slice to a file is the missing piece. +// - Not ported: Dataset.addText and the tokenizer (tokenizer.zig), loadTextFile and +// loadTextFileShard, shuffle (Zig's PRNG), sortByComplexity, splitTrainVal. +// - The checkpoint holds the f32 shadow weights and biases only (no trits, no alphas: loading +// requantizes). f32 is handled as its 32 bits, converted to and from f64 by decode_f32 and +// encode_f32 (normal numbers and zero; no NaN, infinity or denormals). +// - The layout is checked by arithmetic: the byte count that header + sections give equals +// the size of the real file. Its weights are not read here. +// - The curriculum arithmetic runs in f64; data.zig runs it in f32. They differ only when +// difficulty * (max - min) lands within one f32 ulp of a whole number. +// +// Tokens are u16 (a vocabulary of 729 fits), batches are row-major batch_size x seq_len. + +module HslmData { + use base::types; + + // ═════════════════════════════════════════════════════════════════ + // 1. Constants (trinity-training constants.zig and data.zig defaults) + // ═════════════════════════════════════════════════════════════════ + + const CONTEXT_LEN : usize = 27; // 3^3 + const VOCAB_SIZE : usize = 729; // 3^6 + const EMBED_DIM : usize = 243; + const FF_DIM : usize = 729; + const NUM_BLOCKS : usize = 3; + + const CHECKPOINT_MAGIC : u32 = 0x484C534D; // the bytes 4D 53 4C 48, "MSLH" on disk + const CHECKPOINT_VERSION : u32 = 1; + const HEADER_BYTES : usize = 16; // magic, version, step, loss (f32): 4 bytes each + const REAL_FILE_BYTES : usize = 4971796; // hslm_step_10000.bin + + // CurriculumSchedule + const SCHED_LINEAR : u8 = 0; + const SCHED_EXPONENTIAL : u8 = 1; + const SCHED_STEP : u8 = 2; + const SCHED_COSINE : u8 = 3; + const SCHED_NONE : u8 = 4; + + const DEFAULT_TOTAL_STEPS : u64 = 100000; + const DEFAULT_MIN_SEQ_LEN : usize = 8; + const DEFAULT_GROWTH_RATE : f64 = 5.0; + const DEFAULT_MILESTONES : usize = 10; + const PI : f64 = 3.141592653589793; + + // status codes of parse_header + const HEADER_OK : u32 = 0; + const HEADER_BAD_MAGIC : u32 = 1; + const HEADER_BAD_VERSION : u32 = 2; + const HEADER_TOO_SHORT : u32 = 3; + + // ═════════════════════════════════════════════════════════════════ + // 2. Curriculum + // ═════════════════════════════════════════════════════════════════ + + fn floor_pos(x: f64) -> f64 { + const n : usize = @intFromFloat(x); + return @as(f64, @floatFromInt(n)); + } + + // difficulty in [0, 1] at a step; progress = step / total_steps (1.0 when total is 0) + fn difficulty(schedule: u8, step: u64, total_steps: u64, growth_rate: f64, milestones: usize) -> f64 { + var progress : f64 = 1.0; + if (total_steps > 0) { + progress = @as(f64, @floatFromInt(step)) / @as(f64, @floatFromInt(total_steps)); + } + var d : f64 = 1.0; + if (schedule == SCHED_LINEAR) { + d = @min(1.0, progress); + } else if (schedule == SCHED_EXPONENTIAL) { + d = 1.0 - exp(0.0 - growth_rate * progress); + } else if (schedule == SCHED_STEP) { + const m = @as(f64, @floatFromInt(milestones)); + d = floor_pos(progress * m) / m; + } else if (schedule == SCHED_COSINE) { + d = (1.0 - cos(progress * PI)) / 2.0; + } + return @max(0.0, @min(1.0, d)); + } + + // min_seq_len + trunc(difficulty * (max_seq_len - min_seq_len)) + fn curriculum_seq_len(schedule: u8, step: u64, total_steps: u64, min_len: usize, max_len: usize) -> usize { + const d = difficulty(schedule, step, total_steps, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES); + const range = @as(f64, @floatFromInt(max_len - min_len)); + const grown : usize = @intFromFloat(d * range); + return min_len + grown; + } + + // Dataset.getCurrentSeqLen: with no schedule the batch's own seq_len stands + fn current_seq_len(schedule: u8, step: u64, total_steps: u64, min_len: usize, max_len: usize, fixed_len: usize) -> usize { + if (schedule == SCHED_NONE) { + return fixed_len; + } + return curriculum_seq_len(schedule, step, total_steps, min_len, max_len); + } + + // Curriculum.isComplete + fn curriculum_complete(schedule: u8, step: u64, total_steps: u64) -> bool { + const d = difficulty(schedule, step, total_steps, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES); + return d >= 0.99 or step >= total_steps; + } + + // ═════════════════════════════════════════════════════════════════ + // 3. Batches + // ═════════════════════════════════════════════════════════════════ + + // Dataset.numSequences: complete sequences at a length (zero when the data is shorter) + fn num_sequences(total_tokens: usize, cur_len: usize) -> usize { + if (total_tokens < cur_len + 1) { + return 0; + } + return total_tokens - cur_len; + } + + // Dataset.nextBatch. Fills inputs[b*seq_len + s] = data[cursor + s] and the target shifted + // by one, wraps the cursor to 0 when a sequence would run off the end, zero-pads each row + // from cur_len to seq_len, and advances by cur_len per row. Returns the new cursor. When the + // data is shorter than cur_len + 1 it writes nothing and returns the cursor unchanged. + fn next_batch(data: []const u16, cursor: usize, batch_size: usize, seq_len: usize, cur_len: usize, inputs: []u16, targets: []u16) -> usize { + if (data.len < cur_len + 1) { + return cursor; + } + var cur = cursor; + var b : usize = 0; + while (b < batch_size) { + if (cur + cur_len + 1 > data.len) { + cur = 0; + } + const row = b * seq_len; + var s : usize = 0; + while (s < cur_len) { + inputs[row + s] = data[cur + s]; + targets[row + s] = data[cur + s + 1]; + s = s + 1; + } + while (s < seq_len) { + inputs[row + s] = 0; + targets[row + s] = 0; + s = s + 1; + } + cur = cur + cur_len; + b = b + 1; + } + return cur; + } + + // ═════════════════════════════════════════════════════════════════ + // 4. Bytes: little-endian u32 and the f32 bit pattern + // ═════════════════════════════════════════════════════════════════ + + fn read_u32_le(buf: []const u8, off: usize) -> u32 { + const b0 : u32 = @intCast(buf[off]); + const b1 : u32 = @intCast(buf[off + 1]); + const b2 : u32 = @intCast(buf[off + 2]); + const b3 : u32 = @intCast(buf[off + 3]); + return b0 | (b1 << 8) | (b2 << 16) | (b3 << 24); + } + + fn write_u32_le(buf: []u8, off: usize, v: u32) -> void { + buf[off] = @intCast(v & 255); + buf[off + 1] = @intCast((v >> 8) & 255); + buf[off + 2] = @intCast((v >> 16) & 255); + buf[off + 3] = @intCast((v >> 24) & 255); + } + + // 2^k for 0 <= k <= 150 + fn pow2(k: u32) -> f64 { + var r : f64 = 1.0; + var i : u32 = 0; + while (i < k) { + r = r * 2.0; + i = i + 1; + } + return r; + } + + // The f64 value of an f32 bit pattern: normal numbers and zero + fn decode_f32(bits: u32) -> f64 { + const exponent = (bits >> 23) & 255; + const mantissa = bits & 8388607; + var magnitude : f64 = 0.0; + if (exponent != 0) { + const frac = 1.0 + @as(f64, @floatFromInt(mantissa)) / 8388608.0; + if (exponent >= 127) { + magnitude = frac * pow2(exponent - 127); + } else { + magnitude = frac / pow2(127 - exponent); + } + } + if ((bits >> 31) == 1) { + return 0.0 - magnitude; + } + return magnitude; + } + + // The f32 bit pattern nearest an f64 (round half up on the mantissa): normal numbers and zero + fn encode_f32(x: f64) -> u32 { + if (x == 0.0) { + return 0; + } + var sign : u32 = 0; + var m = x; + if (x < 0.0) { + sign = 1; + m = 0.0 - x; + } + var e : u32 = 127; + while (m >= 2.0) { + m = m / 2.0; + e = e + 1; + } + while (m < 1.0) { + m = m * 2.0; + e = e - 1; + } + const frac : u32 = @intFromFloat((m - 1.0) * 8388608.0 + 0.5); + // a mantissa that rounds up to 2.0 carries into the exponent + if (frac >= 8388608) { + return (sign << 31) | ((e + 1) << 23); + } + return (sign << 31) | (e << 23) | frac; + } + + // ═════════════════════════════════════════════════════════════════ + // 5. Checkpoint: header and layout + // ═════════════════════════════════════════════════════════════════ + + // saveCheckpoint's header: magic, version, step, loss bits + fn write_header(buf: []u8, step: u32, loss: f64) -> void { + write_u32_le(buf, 0, CHECKPOINT_MAGIC); + write_u32_le(buf, 4, CHECKPOINT_VERSION); + write_u32_le(buf, 8, step); + write_u32_le(buf, 12, encode_f32(loss)); + } + + // loadCheckpoint's checks, in its order: magic, then version + fn parse_header(buf: []const u8) -> u32 { + if (buf.len < HEADER_BYTES) { + return HEADER_TOO_SHORT; + } + if (read_u32_le(buf, 0) != CHECKPOINT_MAGIC) { + return HEADER_BAD_MAGIC; + } + if (read_u32_le(buf, 4) != CHECKPOINT_VERSION) { + return HEADER_BAD_VERSION; + } + return HEADER_OK; + } + + fn header_step(buf: []const u8) -> u32 { + return read_u32_le(buf, 8); + } + + fn header_loss(buf: []const u8) -> f64 { + return decode_f32(read_u32_le(buf, 12)); + } + + // floats in the output projection's shadow weights + fn output_floats(vocab: usize, embed: usize) -> usize { + return embed * vocab; + } + + // floats per block: shadow_up (embed x ff), shadow_down (ff x embed), bias_up, bias_down + fn block_floats(embed: usize, ff: usize) -> usize { + return embed * ff + ff * embed + ff + embed; + } + + fn param_floats(vocab: usize, embed: usize, ff: usize, blocks: usize) -> usize { + return output_floats(vocab, embed) + blocks * block_floats(embed, ff); + } + + fn file_bytes(vocab: usize, embed: usize, ff: usize, blocks: usize) -> usize { + return HEADER_BYTES + 4 * param_floats(vocab, embed, ff, blocks); + } + + // Byte offsets, in the order saveCheckpoint writes them + fn output_offset() -> usize { + return HEADER_BYTES; + } + + fn block_offset(vocab: usize, embed: usize, ff: usize, b: usize) -> usize { + return HEADER_BYTES + 4 * (output_floats(vocab, embed) + b * block_floats(embed, ff)); + } + + fn shadow_up_offset(vocab: usize, embed: usize, ff: usize, b: usize) -> usize { + return block_offset(vocab, embed, ff, b); + } + + fn shadow_down_offset(vocab: usize, embed: usize, ff: usize, b: usize) -> usize { + return shadow_up_offset(vocab, embed, ff, b) + 4 * embed * ff; + } + + fn bias_up_offset(vocab: usize, embed: usize, ff: usize, b: usize) -> usize { + return shadow_down_offset(vocab, embed, ff, b) + 4 * ff * embed; + } + + fn bias_down_offset(vocab: usize, embed: usize, ff: usize, b: usize) -> usize { + return bias_up_offset(vocab, embed, ff, b) + 4 * ff; + } + + // where the last section ends: must be the file size + fn last_section_end(vocab: usize, embed: usize, ff: usize, blocks: usize) -> usize { + return bias_down_offset(vocab, embed, ff, blocks - 1) + 4 * embed; + } + + fn real_file_bytes() -> usize { + return file_bytes(VOCAB_SIZE, EMBED_DIM, FF_DIM, NUM_BLOCKS); + } + + // ═════════════════════════════════════════════════════════════════ + // 6. Probes (fixed inputs, so each test reads one number) + // ═════════════════════════════════════════════════════════════════ + + // the first 16 bytes of data/wave9/worker-2/hslm_step_10000.bin + fn real_header_status() -> u32 { + const h : [16]u8 = [0x4D, 0x53, 0x4C, 0x48, 0x01, 0x00, 0x00, 0x00, 0x10, 0x27, 0x00, 0x00, 0x17, 0xAA, 0xB1, 0x40]; + return parse_header(&h); + } + + fn real_header_step() -> u32 { + const h : [16]u8 = [0x4D, 0x53, 0x4C, 0x48, 0x01, 0x00, 0x00, 0x00, 0x10, 0x27, 0x00, 0x00, 0x17, 0xAA, 0xB1, 0x40]; + return header_step(&h); + } + + fn real_header_loss() -> f64 { + const h : [16]u8 = [0x4D, 0x53, 0x4C, 0x48, 0x01, 0x00, 0x00, 0x00, 0x10, 0x27, 0x00, 0x00, 0x17, 0xAA, 0xB1, 0x40]; + return header_loss(&h); + } + + fn real_header_loss_bits() -> u32 { + const h : [16]u8 = [0x4D, 0x53, 0x4C, 0x48, 0x01, 0x00, 0x00, 0x00, 0x10, 0x27, 0x00, 0x00, 0x17, 0xAA, 0xB1, 0x40]; + return read_u32_le(&h, 12); + } + + // the same header written by write_header: byte k of the 16 + fn written_header_byte(step: u32, loss: f64, k: usize) -> u32 { + var buf : [16]u8 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]; + write_header(&buf, step, loss); + const v : u32 = @intCast(buf[k]); + return v; + } + + fn written_header_status(step: u32, loss: f64) -> u32 { + var buf : [16]u8 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]; + write_header(&buf, step, loss); + return parse_header(&buf); + } + + fn written_header_step(step: u32, loss: f64) -> u32 { + var buf : [16]u8 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]; + write_header(&buf, step, loss); + return header_step(&buf); + } + + fn written_header_loss(step: u32, loss: f64) -> f64 { + var buf : [16]u8 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]; + write_header(&buf, step, loss); + return header_loss(&buf); + } + + fn bad_magic_status() -> u32 { + const h : [16]u8 = [0x4D, 0x53, 0x4C, 0x49, 0x01, 0x00, 0x00, 0x00, 0x10, 0x27, 0x00, 0x00, 0x17, 0xAA, 0xB1, 0x40]; + return parse_header(&h); + } + + fn bad_version_status() -> u32 { + const h : [16]u8 = [0x4D, 0x53, 0x4C, 0x48, 0x02, 0x00, 0x00, 0x00, 0x10, 0x27, 0x00, 0x00, 0x17, 0xAA, 0xB1, 0x40]; + return parse_header(&h); + } + + fn short_status() -> u32 { + const h : [12]u8 = [0x4D, 0x53, 0x4C, 0x48, 0x01, 0x00, 0x00, 0x00, 0x10, 0x27, 0x00, 0x00]; + return parse_header(&h); + } + + // the order matters: a wrong magic is reported before a wrong version + fn both_wrong_status() -> u32 { + const h : [16]u8 = [0x00, 0x53, 0x4C, 0x48, 0x09, 0x00, 0x00, 0x00, 0x10, 0x27, 0x00, 0x00, 0x17, 0xAA, 0xB1, 0x40]; + return parse_header(&h); + } + + // a float section written into a buffer and read back + fn float_round_trip(x: f64) -> f64 { + var buf : [4]u8 = [0, 0, 0, 0]; + write_u32_le(&buf, 0, encode_f32(x)); + return decode_f32(read_u32_le(&buf, 0)); + } + + // tokens 0..11, a 12-token stream; batch 2 x seq 4 at full length 4 + fn batch_fixture(cursor: usize, cur_len: usize, which: usize) -> usize { + const data : [12]u16 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]; + var inputs : [8]u16 = [99, 99, 99, 99, 99, 99, 99, 99]; + var targets : [8]u16 = [99, 99, 99, 99, 99, 99, 99, 99]; + const next = next_batch(&data, cursor, 2, 4, cur_len, &inputs, &targets); + if (which == 100) { + return next; + } + if (which >= 50) { + const t : usize = @intCast(targets[which - 50]); + return t; + } + const i : usize = @intCast(inputs[which]); + return i; + } + + // a 3-token stream is shorter than cur_len + 1 = 4 + fn short_data_batch(which: usize) -> usize { + const data : [3]u16 = [5, 6, 7]; + var inputs : [4]u16 = [99, 99, 99, 99]; + var targets : [4]u16 = [99, 99, 99, 99]; + const next = next_batch(&data, 2, 1, 4, 3, &inputs, &targets); + if (which == 100) { + return next; + } + const i : usize = @intCast(inputs[which]); + return i; + } + + // ═════════════════════════════════════════════════════════════════ + // TDD-Inside-Spec: Tests and Invariants + // ═════════════════════════════════════════════════════════════════ + + // --- the real checkpoint ------------------------------------------------- + + test real_header_parses_ok + given s = real_header_status() + then s == HEADER_OK + + // 0x2710 = 10000 steps + test real_header_step_is_ten_thousand + given step = real_header_step() + then step == 10000 + + // 0x40B1AA17 = 5.5520129... as f32 + test real_header_loss_bits_and_value + given bits = real_header_loss_bits() + when loss = real_header_loss() + then bits == 0x40B1AA17 and @abs(loss - 5.5520129) < 0.000001 + + // 3 blocks of 243 / 729 and a 729-class head: 1242945 floats after the header + test layout_floats_match_the_model + given n = param_floats(VOCAB_SIZE, EMBED_DIM, FF_DIM, NUM_BLOCKS) + when per_block = block_floats(EMBED_DIM, FF_DIM) + then n == 1242945 and per_block == 355266 and output_floats(VOCAB_SIZE, EMBED_DIM) == 177147 + + // the layout adds up to the size of hslm_step_10000.bin, byte for byte + test layout_bytes_equal_the_real_file + given n = real_file_bytes() + then n == 4971796 and n == REAL_FILE_BYTES + + test layout_sections_tile_the_file + given first = output_offset() + when end0 = shadow_up_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, 0) + then first == 16 and end0 == 16 + 4 * 177147 and last_section_end(VOCAB_SIZE, EMBED_DIM, FF_DIM, NUM_BLOCKS) == REAL_FILE_BYTES + + test layout_each_block_starts_where_the_last_ended + given b1 = block_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, 1) + when b0end = bias_down_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, 0) + 4 * EMBED_DIM + then b1 == b0end and block_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, 2) == bias_down_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, 1) + 4 * EMBED_DIM + + // section order inside a block: up, down, bias_up, bias_down + test layout_block_section_order + given up = shadow_up_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, 0) + when down = shadow_down_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, 0) + then down - up == 4 * 243 * 729 and bias_up_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, 0) - down == 4 * 729 * 243 and bias_down_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, 0) - bias_up_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, 0) == 4 * 729 + + // a small model keeps the same rules: 2 classes, width 3, ff 9, 1 block + test layout_holds_for_a_small_model + given n = file_bytes(2, 3, 9, 1) + then n == 16 + 4 * (6 + 27 + 27 + 9 + 3) and last_section_end(2, 3, 9, 1) == n + + // --- header write and read ----------------------------------------------- + + test header_bytes_follow_the_format + given b0 = written_header_byte(10000, 5.5520129, 0) + when b3 = written_header_byte(10000, 5.5520129, 3) + then b0 == 0x4D and written_header_byte(10000, 5.5520129, 1) == 0x53 and written_header_byte(10000, 5.5520129, 2) == 0x4C and b3 == 0x48 and written_header_byte(10000, 5.5520129, 4) == 1 and written_header_byte(10000, 5.5520129, 8) == 0x10 and written_header_byte(10000, 5.5520129, 9) == 0x27 + + test header_written_then_parsed + given s = written_header_status(10000, 5.5520129) + when step = written_header_step(10000, 5.5520129) + then s == HEADER_OK and step == 10000 and @abs(written_header_loss(10000, 5.5520129) - 5.5520129) < 0.000001 + + // writing the real header's step and loss reproduces the real loss bits exactly + test header_loss_bits_reproduce_the_real_ones + given b12 = written_header_byte(10000, 5.552012920379639, 12) + when b15 = written_header_byte(10000, 5.552012920379639, 15) + then b12 == 0x17 and written_header_byte(10000, 5.552012920379639, 13) == 0xAA and written_header_byte(10000, 5.552012920379639, 14) == 0xB1 and b15 == 0x40 + + test header_rejects_a_wrong_magic + given s = bad_magic_status() + then s == HEADER_BAD_MAGIC + + test header_rejects_a_wrong_version + given s = bad_version_status() + then s == HEADER_BAD_VERSION + + test header_rejects_a_short_buffer + given s = short_status() + then s == HEADER_TOO_SHORT + + test header_reports_the_magic_first + given s = both_wrong_status() + then s == HEADER_BAD_MAGIC + + // --- f32 bit patterns ---------------------------------------------------- + + test f32_known_patterns + given one = encode_f32(1.0) + when small = encode_f32(0.15625) + then one == 0x3F800000 and small == 0x3E200000 and encode_f32(0.0 - 2.0) == 0xC0000000 and encode_f32(0.0) == 0 + + test f32_decode_known_patterns + given one = decode_f32(0x3F800000) + when small = decode_f32(0x3E200000) + then one == 1.0 and small == 0.15625 and decode_f32(0xC0000000) == 0.0 - 2.0 and decode_f32(0) == 0.0 + + // decode then encode returns the same 32 bits, for the real loss and for a negative and a tiny value + test f32_round_trip_on_bits + given real = encode_f32(decode_f32(0x40B1AA17)) + when neg = encode_f32(decode_f32(0xBF99999A)) + then real == 0x40B1AA17 and neg == 0xBF99999A and encode_f32(decode_f32(0x38D1B717)) == 0x38D1B717 + + test f32_round_trip_through_bytes + given a = float_round_trip(0.15625) + when b = float_round_trip(0.0 - 3.5) + then a == 0.15625 and b == 0.0 - 3.5 and @abs(float_round_trip(0.3) - 0.3) < 0.00000003 + + // a mantissa that rounds to 2.0 carries into the exponent: 1.99999999 -> 2.0 + test f32_rounding_carries_into_the_exponent + given bits = encode_f32(1.999999999) + then bits == 0x40000000 + + // --- curriculum ---------------------------------------------------------- + + test curriculum_none_is_full_difficulty + given d = difficulty(SCHED_NONE, 0, DEFAULT_TOTAL_STEPS, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES) + then d == 1.0 + + test curriculum_linear_follows_progress + given d = difficulty(SCHED_LINEAR, 25000, DEFAULT_TOTAL_STEPS, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES) + when past = difficulty(SCHED_LINEAR, 250000, DEFAULT_TOTAL_STEPS, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES) + then @abs(d - 0.25) < 0.000000000001 and past == 1.0 + + test curriculum_exponential + given d = difficulty(SCHED_EXPONENTIAL, 50000, DEFAULT_TOTAL_STEPS, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES) + then @abs(d - 0.917915001376101) < 0.000000000001 + + // milestones: floor(0.55 * 10) / 10 = 0.5 + test curriculum_step_holds_between_milestones + given d = difficulty(SCHED_STEP, 55000, DEFAULT_TOTAL_STEPS, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES) + when same = difficulty(SCHED_STEP, 50000, DEFAULT_TOTAL_STEPS, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES) + then @abs(d - 0.5) < 0.000000000001 and @abs(same - 0.5) < 0.000000000001 + + test curriculum_cosine + given half = difficulty(SCHED_COSINE, 50000, DEFAULT_TOTAL_STEPS, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES) + when quarter = difficulty(SCHED_COSINE, 25000, DEFAULT_TOTAL_STEPS, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES) + then @abs(half - 0.5) < 0.000000000001 and @abs(quarter - 0.14644660940672627) < 0.000000000001 + + test curriculum_zero_total_steps_is_full_progress + given d = difficulty(SCHED_LINEAR, 0, 0, DEFAULT_GROWTH_RATE, DEFAULT_MILESTONES) + then d == 1.0 + + // 8 + trunc(d * 19): linear 0.5 -> 17, cosine 0.25 -> 8 + trunc(2.78) = 10, exponential 0.918 -> 8 + 17 = 25 + test curriculum_sequence_length + given lin = curriculum_seq_len(SCHED_LINEAR, 50000, DEFAULT_TOTAL_STEPS, DEFAULT_MIN_SEQ_LEN, CONTEXT_LEN) + when cosq = curriculum_seq_len(SCHED_COSINE, 25000, DEFAULT_TOTAL_STEPS, DEFAULT_MIN_SEQ_LEN, CONTEXT_LEN) + then lin == 17 and cosq == 10 and curriculum_seq_len(SCHED_EXPONENTIAL, 50000, DEFAULT_TOTAL_STEPS, DEFAULT_MIN_SEQ_LEN, CONTEXT_LEN) == 25 + + test curriculum_starts_at_the_minimum_and_ends_at_the_maximum + given start = curriculum_seq_len(SCHED_LINEAR, 0, DEFAULT_TOTAL_STEPS, DEFAULT_MIN_SEQ_LEN, CONTEXT_LEN) + when finish = curriculum_seq_len(SCHED_LINEAR, 100000, DEFAULT_TOTAL_STEPS, DEFAULT_MIN_SEQ_LEN, CONTEXT_LEN) + then start == 8 and finish == 27 + + test current_seq_len_without_a_schedule_is_the_fixed_one + given n = current_seq_len(SCHED_NONE, 50000, DEFAULT_TOTAL_STEPS, DEFAULT_MIN_SEQ_LEN, CONTEXT_LEN, 12) + when m = current_seq_len(SCHED_LINEAR, 50000, DEFAULT_TOTAL_STEPS, DEFAULT_MIN_SEQ_LEN, CONTEXT_LEN, 12) + then n == 12 and m == 17 + + test curriculum_completes_at_the_end + given early = curriculum_complete(SCHED_LINEAR, 1000, DEFAULT_TOTAL_STEPS) + when late = curriculum_complete(SCHED_LINEAR, 99000, DEFAULT_TOTAL_STEPS) + then early == false and late == true and curriculum_complete(SCHED_LINEAR, 100000, DEFAULT_TOTAL_STEPS) == true and curriculum_complete(SCHED_STEP, 98000, DEFAULT_TOTAL_STEPS) == false + + // --- batches ------------------------------------------------------------- + + test num_sequences_counts_start_positions + given a = num_sequences(12, 4) + then a == 8 and num_sequences(4, 4) == 0 and num_sequences(5, 4) == 1 and num_sequences(0, 4) == 0 + + // row 0 = tokens 0..3, row 1 = tokens 4..7: inputs[0] = 0, inputs[4] = 4 + test batch_inputs_are_consecutive_tokens + given i0 = batch_fixture(0, 4, 0) + when i3 = batch_fixture(0, 4, 3) + then i0 == 0 and i3 == 3 and batch_fixture(0, 4, 4) == 4 and batch_fixture(0, 4, 7) == 7 + + // the target is the input shifted by one + test batch_targets_are_shifted_by_one + given t0 = batch_fixture(0, 4, 50) + when t3 = batch_fixture(0, 4, 53) + then t0 == 1 and t3 == 4 and batch_fixture(0, 4, 54) == 5 and batch_fixture(0, 4, 57) == 8 + + test batch_advances_the_cursor_by_one_row_per_example + given next = batch_fixture(0, 4, 100) + then next == 8 + + // cursor 4: row 0 = 4..7, row 1 would need 8..12 but 8 + 5 > 12, so the cursor wraps to 0 + test batch_wraps_when_a_sequence_would_run_off_the_end + given i4 = batch_fixture(4, 4, 4) + when next = batch_fixture(4, 4, 100) + then i4 == 0 and batch_fixture(4, 4, 0) == 4 and next == 4 + + // current length 2 in a row of 4: two tokens, then two PAD (0) entries + test batch_pads_a_short_curriculum_row + given i1 = batch_fixture(0, 2, 1) + when i2 = batch_fixture(0, 2, 2) + then i1 == 1 and i2 == 0 and batch_fixture(0, 2, 3) == 0 and batch_fixture(0, 2, 4) == 2 and batch_fixture(0, 2, 53) == 0 + + test batch_with_a_short_curriculum_advances_by_the_current_length + given next = batch_fixture(0, 2, 100) + then next == 4 + + // shorter than cur_len + 1: nothing is written, the cursor stays + test batch_leaves_everything_alone_when_data_is_too_short + given next = short_data_batch(100) + then next == 2 and short_data_batch(0) == 99 and short_data_batch(3) == 99 + + invariant checkpoint_magic_is_hslm_in_little_endian + assert CHECKPOINT_MAGIC == 0x484C534D and CHECKPOINT_VERSION == 1 and HEADER_BYTES == 16 + + invariant resonance_ctx_is_three_cubed + assert CONTEXT_LEN == 27 and CONTEXT_LEN == 3 * 3 * 3 and VOCAB_SIZE == 729 and EMBED_DIM == 243 + + invariant schedules_are_distinct + assert SCHED_LINEAR != SCHED_EXPONENTIAL and SCHED_STEP != SCHED_COSINE and SCHED_COSINE != SCHED_NONE +} diff --git a/specs/nn/ternary_ste.t27 b/specs/nn/ternary_ste.t27 new file mode 100644 index 0000000000..6b846c6211 --- /dev/null +++ b/specs/nn/ternary_ste.t27 @@ -0,0 +1,548 @@ +// SPDX-License-Identifier: Apache-2.0 +// t27/specs/nn/ternary_ste.t27 +// Ternary straight-through estimator (STE) for HSLM training: shadow float weights are +// quantized to {-1, 0, +1} for the forward pass, and the gradient of the loss with respect +// to a ternary weight is handed to the shadow weight as if the quantizer were the identity. +// phi^2 + 1/phi^2 = 3 | TRINITY +// +// Source of truth: gHashTag/trinity-training @ 3e9648c1817a97c5f9b8c727b0f8e2291928eead +// src/hslm/ste.zig quantizeAbsMean, quantizeVanilla, quantizeTwn, quantizeProgressive, +// quantizeForMode, applyAlpha, computeMetrics (the quantizers) +// src/hslm/autograd.zig clipGradNorm (the gradient clip the trainer applies to every +// shadow-weight gradient before the AdamW step) +// src/hslm/trainer.zig the shadow weights are the trained parameters; after each +// optimizer step the model requantizes ternary weights from them +// +// What this spec adds to ste.zig, and what it leaves out, so nobody reads it as more: +// - ste.zig has no backward function: its gradient is the identity, taken by the autograd +// tape. ste_pass_through is that identity, written out. ste_pass is the clipped variant +// (the gradient is zero where |shadow| > clip), the usual hardtanh STE; the trainer does +// not use it, so it is a choice of this spec, not a copy. +// - Floats: ste.zig computes in f32, this spec in f64. They differ only for a weight that +// sits within one f32 ulp of a threshold. +// - quantize_progressive follows ste.zig exactly, including its odd blend: during the +// transition the threshold does not tighten gradually, it switches from the abs-mean rule +// to TWN once progress passes one half. +// - The optimizer is plain SGD here (sgd_shadow_step); the trainer uses AdamW (specs/ml/ +// optimizer/adam.t27 holds Adam). +// +// Trits are base::types Trit values (neg, zero, pos); trit_value turns one into -1.0, 0.0, 1.0. + +module TernarySte { + use base::types; + + // ═════════════════════════════════════════════════════════════════ + // 1. Constants (ste.zig defaults) + // ═════════════════════════════════════════════════════════════════ + + // Li et al. 2016: delta = 0.7 * mean(|w|) + const TWN_DELTA_FACTOR : f64 = 0.7; + // quantizeAbsMean: a scaled weight above 0.5 in magnitude rounds to a trit + const ABS_MEAN_ROUND : f64 = 0.5; + // quantizeAbsMean: below this mean the scale falls back to 1.0 + const ABS_MEAN_FLOOR : f64 = 0.000001; + const DEFAULT_VANILLA_THRESHOLD : f64 = 0.5; + const DEFAULT_WARMUP_STEPS : u32 = 10000; + const DEFAULT_TRANSITION_STEPS : u32 = 10000; + + // SteMode + const MODE_NONE : u8 = 0; + const MODE_VANILLA : u8 = 1; + const MODE_TWN : u8 = 2; + const MODE_PROGRESSIVE : u8 = 3; + + // ═════════════════════════════════════════════════════════════════ + // 2. Quantizers: shadow floats -> trits, each returns alpha + // ═════════════════════════════════════════════════════════════════ + + fn trit_value(t: Trit) -> f64 { + if (t == Trit.pos) { + return 1.0; + } else if (t == Trit.neg) { + return -1.0; + } + return 0.0; + } + + fn mean_abs(w: []const f64) -> f64 { + var sum : f64 = 0.0; + var i : usize = 0; + while (i < w.len) { + sum = sum + @abs(w[i]); + i = i + 1; + } + return sum / @as(f64, @floatFromInt(w.len)); + } + + // Standard rule: scale = mean(|w|) (1.0 when that is below the floor), trit = round of + // w / scale to {-1, 0, 1} with the cut at one half. Returns the scale. + fn quantize_abs_mean(w: []const f64, out: []Trit) -> f64 { + var scale : f64 = mean_abs(w); + if (scale <= ABS_MEAN_FLOOR) { + scale = 1.0; + } + var i : usize = 0; + while (i < w.len) { + const scaled = w[i] / scale; + if (scaled > ABS_MEAN_ROUND) { + out[i] = Trit.pos; + } else if (scaled < -ABS_MEAN_ROUND) { + out[i] = Trit.neg; + } else { + out[i] = Trit.zero; + } + i = i + 1; + } + return scale; + } + + // Vanilla STE: a fixed threshold, no scale. Returns 1.0. + fn quantize_vanilla(w: []const f64, out: []Trit, threshold: f64) -> f64 { + var i : usize = 0; + while (i < w.len) { + if (w[i] > threshold) { + out[i] = Trit.pos; + } else if (w[i] < -threshold) { + out[i] = Trit.neg; + } else { + out[i] = Trit.zero; + } + i = i + 1; + } + return 1.0; + } + + // The TWN threshold, delta = 0.7 * mean(|w|) + fn twn_delta(w: []const f64) -> f64 { + return TWN_DELTA_FACTOR * mean_abs(w); + } + + // TWN: trit = sign(w) where |w| > delta, else 0; alpha = mean(|w|) over the nonzero + // trits, or 1.0 when there is none. Forward: y = alpha * (x @ W_trits). + fn quantize_twn(w: []const f64, out: []Trit) -> f64 { + const delta = twn_delta(w); + var alpha_sum : f64 = 0.0; + var alpha_count : usize = 0; + var i : usize = 0; + while (i < w.len) { + if (w[i] > delta) { + out[i] = Trit.pos; + alpha_sum = alpha_sum + @abs(w[i]); + alpha_count = alpha_count + 1; + } else if (w[i] < -delta) { + out[i] = Trit.neg; + alpha_sum = alpha_sum + @abs(w[i]); + alpha_count = alpha_count + 1; + } else { + out[i] = Trit.zero; + } + i = i + 1; + } + if (alpha_count > 0) { + return alpha_sum / @as(f64, @floatFromInt(alpha_count)); + } + return 1.0; + } + + // Warmup: abs-mean. Transition: abs-mean until progress passes 0.5, then TWN. After: TWN. + fn quantize_progressive(w: []const f64, out: []Trit, step: u32, warmup: u32, transition: u32) -> f64 { + if (step < warmup) { + return quantize_abs_mean(w, out); + } else if (step < warmup + transition) { + const progress = @as(f64, @floatFromInt(step - warmup)) / @as(f64, @floatFromInt(transition)); + if (progress > 0.5) { + return quantize_twn(w, out); + } + return quantize_abs_mean(w, out); + } + return quantize_twn(w, out); + } + + fn quantize_for_mode(w: []const f64, out: []Trit, mode: u8, threshold: f64, step: u32, warmup: u32, transition: u32) -> f64 { + if (mode == MODE_VANILLA) { + return quantize_vanilla(w, out, threshold); + } else if (mode == MODE_TWN) { + return quantize_twn(w, out); + } else if (mode == MODE_PROGRESSIVE) { + return quantize_progressive(w, out, step, warmup, transition); + } + return quantize_abs_mean(w, out); + } + + // What the forward pass sees: alpha * trit + fn dequantize(trits: []const Trit, alpha: f64, out: []f64) -> void { + var i : usize = 0; + while (i < trits.len) { + out[i] = alpha * trit_value(trits[i]); + i = i + 1; + } + } + + // After the ternary matvec: scale the output (a no-op at alpha exactly 1.0, as in ste.zig) + fn apply_alpha(output: []f64, alpha: f64) -> void { + if (alpha == 1.0) { + return; + } + var i : usize = 0; + while (i < output.len) { + output[i] = output[i] * alpha; + i = i + 1; + } + } + + fn count_zero(trits: []const Trit) -> usize { + var zeros : usize = 0; + var i : usize = 0; + while (i < trits.len) { + if (trits[i] == Trit.zero) { + zeros = zeros + 1; + } + i = i + 1; + } + return zeros; + } + + // computeMetrics: percent of trits that are zero + fn sparsity_pct(trits: []const Trit) -> f64 { + return @as(f64, @floatFromInt(count_zero(trits))) / @as(f64, @floatFromInt(trits.len)) * 100.0; + } + + // ═════════════════════════════════════════════════════════════════ + // 3. Straight-through backward and the shadow-weight update + // ═════════════════════════════════════════════════════════════════ + + // The identity STE: the gradient of the ternary weight is the gradient of the shadow weight + fn ste_pass_through(g: []const f64, out: []f64) -> void { + var i : usize = 0; + while (i < g.len) { + out[i] = g[i]; + i = i + 1; + } + } + + // The clipped STE: the gradient passes where |shadow| <= clip, and is zero beyond it + fn ste_pass(shadow: []const f64, g: []const f64, clip: f64, out: []f64) -> void { + var i : usize = 0; + while (i < g.len) { + if (@abs(shadow[i]) <= clip) { + out[i] = g[i]; + } else { + out[i] = 0.0; + } + i = i + 1; + } + } + + fn grad_norm(g: []const f64) -> f64 { + var sq : f64 = 0.0; + var i : usize = 0; + while (i < g.len) { + sq = sq + g[i] * g[i]; + i = i + 1; + } + return sqrt(sq); + } + + // autograd.clipGradNorm: when the L2 norm exceeds max_norm, scale the gradient down to it + fn clip_grad_norm(g: []f64, max_norm: f64) -> void { + const norm = grad_norm(g); + if (norm > max_norm) { + const scale = max_norm / norm; + var i : usize = 0; + while (i < g.len) { + g[i] = g[i] * scale; + i = i + 1; + } + } + } + + // shadow = shadow - lr * g. The shadow weights move in float; the trits follow on requantize. + fn sgd_shadow_step(shadow: []f64, g: []const f64, lr: f64) -> void { + var i : usize = 0; + while (i < shadow.len) { + shadow[i] = shadow[i] - lr * g[i]; + i = i + 1; + } + } + + // The trainer's requantize: TWN trits and alpha from the shadow weights + fn requantize(shadow: []const f64, out: []Trit) -> f64 { + return quantize_twn(shadow, out); + } + + // ═════════════════════════════════════════════════════════════════ + // 4. Probes (fixed inputs, so each test reads one number) + // ═════════════════════════════════════════════════════════════════ + + // ste.zig tests: [0.7, -0.8, 0.3, -0.2, 0.0, 1.5, -1.5] at threshold 0.5 + fn vanilla_trit(i: usize) -> f64 { + const w : [7]f64 = [0.7, -0.8, 0.3, -0.2, 0.0, 1.5, -1.5]; + var out : [7]Trit = [Trit.zero, Trit.zero, Trit.zero, Trit.zero, Trit.zero, Trit.zero, Trit.zero]; + const alpha = quantize_vanilla(&w, &out, DEFAULT_VANILLA_THRESHOLD); + return trit_value(out[i]) + alpha - 1.0; + } + + fn vanilla_alpha() -> f64 { + const w : [3]f64 = [0.7, -0.8, 0.3]; + var out : [3]Trit = [Trit.zero, Trit.zero, Trit.zero]; + return quantize_vanilla(&w, &out, DEFAULT_VANILLA_THRESHOLD); + } + + // mean(|w|) = 3.75 / 6 = 0.625, delta = 0.4375, nonzero |w| = 0.7 0.8 0.9 1.2 + fn twn_fixture(i: usize) -> f64 { + const w : [6]f64 = [0.7, -0.8, 0.1, -0.05, 0.9, -1.2]; + var out : [6]Trit = [Trit.zero, Trit.zero, Trit.zero, Trit.zero, Trit.zero, Trit.zero]; + const alpha = quantize_twn(&w, &out); + if (i == 6) { + return alpha; + } else if (i == 7) { + return sparsity_pct(&out); + } + return trit_value(out[i]); + } + + fn twn_delta_fixture() -> f64 { + const w : [6]f64 = [0.7, -0.8, 0.1, -0.05, 0.9, -1.2]; + return twn_delta(&w); + } + + // all zero: delta 0, no weight above it, alpha falls back to 1.0 + fn twn_zero_fixture(i: usize) -> f64 { + const w : [3]f64 = [0.0, 0.0, 0.0]; + var out : [3]Trit = [Trit.pos, Trit.pos, Trit.pos]; + const alpha = quantize_twn(&w, &out); + if (i == 3) { + return alpha; + } + return trit_value(out[i]); + } + + // mean(|w|) = 0.5 = scale; w / scale = 1.8, -0.2, 0.8, -1.2 + fn abs_mean_fixture(i: usize) -> f64 { + const w : [4]f64 = [0.9, -0.1, 0.4, -0.6]; + var out : [4]Trit = [Trit.zero, Trit.zero, Trit.zero, Trit.zero]; + const alpha = quantize_abs_mean(&w, &out); + if (i == 4) { + return alpha; + } + return trit_value(out[i]); + } + + // mean 0.475: abs-mean gives alpha 0.475 and trits 1 -1 1 0 (0.3 / 0.475 = 0.63); + // TWN has delta 0.3325, so 0.3 drops to 0, and alpha = (0.7 + 0.8) / 2 = 0.75 + fn progressive_fixture(step: u32, i: usize) -> f64 { + const w : [4]f64 = [0.7, -0.8, 0.3, -0.1]; + var out : [4]Trit = [Trit.zero, Trit.zero, Trit.zero, Trit.zero]; + const alpha = quantize_progressive(&w, &out, step, 100, 100); + if (i == 4) { + return alpha; + } + return trit_value(out[i]); + } + + fn mode_fixture(mode: u8) -> f64 { + const w : [4]f64 = [0.7, -0.8, 0.3, -0.1]; + var out : [4]Trit = [Trit.zero, Trit.zero, Trit.zero, Trit.zero]; + return quantize_for_mode(&w, &out, mode, 0.5, 0, 100, 100); + } + + fn dequantize_fixture(i: usize) -> f64 { + const trits : [3]Trit = [Trit.pos, Trit.zero, Trit.neg]; + var out : [3]f64 = [9.0, 9.0, 9.0]; + dequantize(&trits, 0.75, &out); + return out[i]; + } + + fn apply_alpha_fixture(alpha: f64, i: usize) -> f64 { + var data : [4]f64 = [1.0, -2.0, 3.0, 0.0]; + apply_alpha(&data, alpha); + return data[i]; + } + + // |shadow| = 0.5 1.5 1.0 1.0 against clip 1.0: the weight at 1.0 still passes + fn ste_pass_fixture(clip: f64, i: usize) -> f64 { + const shadow : [4]f64 = [0.5, -1.5, 1.0, -1.0]; + const g : [4]f64 = [0.1, 0.2, -0.3, 0.4]; + var out : [4]f64 = [9.0, 9.0, 9.0, 9.0]; + ste_pass(&shadow, &g, clip, &out); + return out[i]; + } + + fn ste_identity_fixture(i: usize) -> f64 { + const g : [4]f64 = [0.1, 0.2, -0.3, 0.4]; + var out : [4]f64 = [9.0, 9.0, 9.0, 9.0]; + ste_pass_through(&g, &out); + return out[i]; + } + + fn clip_fixture(max_norm: f64, i: usize) -> f64 { + var g : [2]f64 = [3.0, 4.0]; + clip_grad_norm(&g, max_norm); + return g[i]; + } + + fn norm_fixture() -> f64 { + const g : [2]f64 = [3.0, 4.0]; + return grad_norm(&g); + } + + // shadow 0.3 0.9 -0.9 0.0: delta 0.3675, trits 0 1 -1 0. A step of lr 1 against + // g0 = -0.4 moves shadow[0] to 0.7: mean 0.625, delta 0.4375, trit 1, alpha 2.5 / 3. + fn flip_fixture(steps: u32, i: usize) -> f64 { + var shadow : [4]f64 = [0.3, 0.9, -0.9, 0.0]; + const g : [4]f64 = [-0.4, 0.0, 0.0, 0.0]; + var out : [4]Trit = [Trit.zero, Trit.zero, Trit.zero, Trit.zero]; + var n : u32 = 0; + while (n < steps) { + sgd_shadow_step(&shadow, &g, 1.0); + n = n + 1; + } + const alpha = requantize(&shadow, &out); + if (i == 4) { + return alpha; + } else if (i == 5) { + return shadow[0]; + } + return trit_value(out[i]); + } + + // ═════════════════════════════════════════════════════════════════ + // TDD-Inside-Spec: Tests and Invariants for the ternary STE + // ═════════════════════════════════════════════════════════════════ + + // The ste.zig vanilla test, value by value: 0.7 and 1.5 are above 0.5, -0.8 and -1.5 below -0.5 + test ste_vanilla_trits_follow_the_threshold + given t0 = vanilla_trit(0) + when t1 = vanilla_trit(1) + then t0 == 1.0 and t1 == -1.0 and vanilla_trit(2) == 0.0 and vanilla_trit(3) == 0.0 and vanilla_trit(4) == 0.0 and vanilla_trit(5) == 1.0 and vanilla_trit(6) == -1.0 + + test ste_vanilla_alpha_is_one + given a = vanilla_alpha() + then a == 1.0 + + test ste_twn_threshold_is_seven_tenths_of_the_mean + given d = twn_delta_fixture() + then @abs(d - 0.4375) < 0.000000000001 + + test ste_twn_trits + given t0 = twn_fixture(0) + when t1 = twn_fixture(1) + then t0 == 1.0 and t1 == -1.0 and twn_fixture(2) == 0.0 and twn_fixture(3) == 0.0 and twn_fixture(4) == 1.0 and twn_fixture(5) == -1.0 + + // alpha = (0.7 + 0.8 + 0.9 + 1.2) / 4 = 0.9, over the four nonzero trits only + test ste_twn_alpha_is_the_mean_of_the_nonzero_magnitudes + given a = twn_fixture(6) + then @abs(a - 0.9) < 0.000000000001 + + // two of six trits are zero + test ste_twn_sparsity + given s = twn_fixture(7) + then @abs(s - 33.33333333333333) < 0.000001 + + test ste_twn_all_zero_weights_give_zero_trits_and_alpha_one + given a = twn_zero_fixture(3) + then a == 1.0 and twn_zero_fixture(0) == 0.0 and twn_zero_fixture(1) == 0.0 and twn_zero_fixture(2) == 0.0 + + test ste_abs_mean_trits_and_scale + given a = abs_mean_fixture(4) + when t0 = abs_mean_fixture(0) + then @abs(a - 0.5) < 0.000000000001 and t0 == 1.0 and abs_mean_fixture(1) == 0.0 and abs_mean_fixture(2) == 1.0 and abs_mean_fixture(3) == -1.0 + + test ste_progressive_warmup_is_abs_mean + given a = progressive_fixture(50, 4) + when t2 = progressive_fixture(50, 2) + then @abs(a - 0.475) < 0.000000000001 and t2 == 1.0 + + // progress 0.5 is not above 0.5: still abs-mean + test ste_progressive_transition_first_half_is_abs_mean + given a = progressive_fixture(150, 4) + then @abs(a - 0.475) < 0.000000000001 and progressive_fixture(150, 2) == 1.0 + + // progress 0.6: TWN, and the 0.3 weight falls into the dead zone + test ste_progressive_transition_second_half_is_twn + given a = progressive_fixture(160, 4) + then @abs(a - 0.75) < 0.000000000001 and progressive_fixture(160, 2) == 0.0 + + test ste_progressive_after_transition_is_twn + given a = progressive_fixture(250, 4) + then @abs(a - 0.75) < 0.000000000001 and progressive_fixture(250, 3) == 0.0 + + test ste_mode_dispatch + given none = mode_fixture(MODE_NONE) + when vanilla = mode_fixture(MODE_VANILLA) + then @abs(none - 0.475) < 0.000000000001 and vanilla == 1.0 and @abs(mode_fixture(MODE_TWN) - 0.75) < 0.000000000001 and @abs(mode_fixture(MODE_PROGRESSIVE) - 0.475) < 0.000000000001 + + test ste_dequantize_is_alpha_times_trit + given d0 = dequantize_fixture(0) + when d2 = dequantize_fixture(2) + then d0 == 0.75 and dequantize_fixture(1) == 0.0 and d2 == -0.75 + + test ste_apply_alpha_scales + given a0 = apply_alpha_fixture(0.75, 0) + when a1 = apply_alpha_fixture(0.75, 1) + then a0 == 0.75 and a1 == -1.5 and apply_alpha_fixture(0.75, 2) == 2.25 and apply_alpha_fixture(0.75, 3) == 0.0 + + test ste_apply_alpha_one_changes_nothing + given a1 = apply_alpha_fixture(1.0, 1) + then a1 == -2.0 and apply_alpha_fixture(1.0, 2) == 3.0 + + test ste_identity_passes_every_gradient + given g0 = ste_identity_fixture(0) + when g1 = ste_identity_fixture(1) + then g0 == 0.1 and g1 == 0.2 and ste_identity_fixture(2) == -0.3 and ste_identity_fixture(3) == 0.4 + + // inside the window the gradient is the input gradient; outside it is zero + test ste_clipped_passes_inside_and_blocks_outside + given g0 = ste_pass_fixture(1.0, 0) + when g1 = ste_pass_fixture(1.0, 1) + then g0 == 0.1 and g1 == 0.0 and ste_pass_fixture(1.0, 2) == -0.3 and ste_pass_fixture(1.0, 3) == 0.4 + + // a window wider than every shadow weight is the identity STE + test ste_clipped_with_a_wide_window_is_the_identity + given g1 = ste_pass_fixture(2.0, 1) + then g1 == 0.2 and ste_pass_fixture(2.0, 0) == ste_identity_fixture(0) + + // a window narrower than |0.5| blocks all four + test ste_clipped_with_a_narrow_window_blocks_everything + given g0 = ste_pass_fixture(0.25, 0) + then g0 == 0.0 and ste_pass_fixture(0.25, 1) == 0.0 and ste_pass_fixture(0.25, 2) == 0.0 and ste_pass_fixture(0.25, 3) == 0.0 + + test ste_grad_norm_is_euclidean + given n = norm_fixture() + then n == 5.0 + + // norm 5 against max 1: scaled by 0.2, so [0.6, 0.8] + test ste_clip_grad_norm_scales_down_to_the_bound + given g0 = clip_fixture(1.0, 0) + when g1 = clip_fixture(1.0, 1) + then @abs(g0 - 0.6) < 0.000000000001 and @abs(g1 - 0.8) < 0.000000000001 + + test ste_clip_grad_norm_leaves_a_small_gradient_alone + given g0 = clip_fixture(10.0, 0) + then g0 == 3.0 and clip_fixture(10.0, 1) == 4.0 + + // no step: shadow 0.3 is inside delta 0.3675, so its trit is 0, alpha = (0.9 + 0.9) / 2 + test ste_requantize_before_the_step + given t0 = flip_fixture(0, 0) + when a = flip_fixture(0, 4) + then t0 == 0.0 and flip_fixture(0, 1) == 1.0 and flip_fixture(0, 2) == -1.0 and @abs(a - 0.9) < 0.000000000001 + + // one SGD step moves the shadow weight to 0.7 and requantize flips its trit to +1 + test ste_one_step_flips_a_trit + given s0 = flip_fixture(1, 5) + when t0 = flip_fixture(1, 0) + then @abs(s0 - 0.7) < 0.000000000001 and t0 == 1.0 and flip_fixture(1, 1) == 1.0 and flip_fixture(1, 2) == -1.0 + + test ste_alpha_follows_the_shadow_weights + given a = flip_fixture(1, 4) + then @abs(a - 0.8333333333333334) < 0.000000000001 + + invariant ste_modes_are_distinct + assert MODE_NONE != MODE_VANILLA and MODE_VANILLA != MODE_TWN and MODE_TWN != MODE_PROGRESSIVE and MODE_NONE != MODE_TWN + + invariant ste_twn_factor_is_seven_tenths + assert TWN_DELTA_FACTOR > 0.69999 and TWN_DELTA_FACTOR < 0.70001 + + invariant ste_defaults_match_ste_zig + assert DEFAULT_WARMUP_STEPS == 10000 and DEFAULT_TRANSITION_STEPS == 10000 +}