From 1a6e210325550e9cb5fdb08e37b9536c372fb2c8 Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Sun, 18 Oct 2026 01:04:52 +0700 Subject: [PATCH] spec(hslm): nn/hslm is the model trinity-training trained (Closes #10062) specs/nn/hslm.t27 now describes src/hslm in trinity-training @3e9648c1: 3 blocks, context 27 (the PPL 2.96 run), a ReLU FFN 243-729-243 with biases and TWN alphas, phi position encoding and phi-RoPE, a token embedding and the vocab-729 head with cross-entropy on the last position. The header cites the trained source files and names what is not modelled (the per-block VSA path, f32, the trainer's STE update). The backward pass reaches the embedding; every gradient (embedding, Q/K/V/O, W_up/W_down, biases, both RMSNorm gammas, W_out, b_out) is checked against a central difference or an exact trit flip. Lab (t27c @efe3c94): 46/46 pass, 0 vacuous, 10 invariants proved. A mutant that drops alpha_q fails 10 of the 46. own_model: forward 60 -> 85, backward 40 -> 70; HSLM lens 51 -> 57, overall 42 -> 44; rows and header updated. own_model_board mirrors it. Lab: own_model 8/8, own_model_board 6/6, 0 vacuous. Seals HSLM and nn_HSLM re-saved on the lab with t27c seal --save. Co-Authored-By: Claude Opus 5.5 --- .trinity/seals/HSLM.json | 16 +- .trinity/seals/nn_HSLM.json | 16 +- specs/igla/own_model.t27 | 28 +- specs/igla/own_model_board.t27 | 22 +- specs/nn/hslm.t27 | 1477 +++++++++++++++++++++++--------- 5 files changed, 1109 insertions(+), 450 deletions(-) diff --git a/.trinity/seals/HSLM.json b/.trinity/seals/HSLM.json index dcca23f704..45b7345a3b 100644 --- a/.trinity/seals/HSLM.json +++ b/.trinity/seals/HSLM.json @@ -1,19 +1,19 @@ { - "gen_hash_c": "sha256:fd94f6962089194da7478e83cd9a015b83a28707b6cb35a9c8e629a96f502fa8", - "gen_hash_rust": "sha256:7f60cf6c5c5c2bf0ee07982585e981bb3e84a5444eab90063406f36940d9d76d", - "gen_hash_verilog": "sha256:25e87d2ab265533cab0de48e751110380236e26686ba78a711d412eb625cf1a8", - "gen_hash_zig": "sha256:769110ea52ada18974a5c561ea7c364bf4c7169c29b55619fb7daf16a2a91451", + "gen_hash_c": "sha256:5e6c04320ab12b7f1451bdfd4162ca52a3eb54c1badfb47d84d6fa8411724d19", + "gen_hash_rust": "sha256:435b3e5b895f99ea146078d091715ee5a71c82aa13a204090bd4b53939983b26", + "gen_hash_verilog": "sha256:cafcce94d7360dcf804d9ef3dc6fe5de0b4de7335c7fcc05a747ee4db0be9800", + "gen_hash_zig": "sha256:0b168aaa160d077bfbcb66a19487dce330289422c64da51cec4154fd5ff7799a", "module": "HSLM", "ring": 12, - "sealed_at": "2026-10-17T17:34:44Z", - "spec_hash": "sha256:1ffe306efd7cc6b092829cbcfb7645cbadde70fed20b3c3a7e3cfedd3abd27ce", + "sealed_at": "2026-10-17T18:02:24Z", + "spec_hash": "sha256:df24bc0cb83391c7131fc940940bae4804f95fde9e854ae03065ec8caff85a08", "spec_path": "specs/nn/hslm.t27", "tests": { "failed": 0, "failing": [], "forced": false, - "passed": 26, - "total": 26, + "passed": 46, + "total": 46, "vacuous": 0 } } \ No newline at end of file diff --git a/.trinity/seals/nn_HSLM.json b/.trinity/seals/nn_HSLM.json index af89890f1e..a0e9e7c4ef 100644 --- a/.trinity/seals/nn_HSLM.json +++ b/.trinity/seals/nn_HSLM.json @@ -7,23 +7,23 @@ } ], "config": "gen=c,rust,verilog,zig;tests=zig", - "gen_hash_c": "sha256:fd94f6962089194da7478e83cd9a015b83a28707b6cb35a9c8e629a96f502fa8", - "gen_hash_rust": "sha256:7f60cf6c5c5c2bf0ee07982585e981bb3e84a5444eab90063406f36940d9d76d", - "gen_hash_verilog": "sha256:25e87d2ab265533cab0de48e751110380236e26686ba78a711d412eb625cf1a8", - "gen_hash_zig": "sha256:769110ea52ada18974a5c561ea7c364bf4c7169c29b55619fb7daf16a2a91451", + "gen_hash_c": "sha256:5e6c04320ab12b7f1451bdfd4162ca52a3eb54c1badfb47d84d6fa8411724d19", + "gen_hash_rust": "sha256:435b3e5b895f99ea146078d091715ee5a71c82aa13a204090bd4b53939983b26", + "gen_hash_verilog": "sha256:cafcce94d7360dcf804d9ef3dc6fe5de0b4de7335c7fcc05a747ee4db0be9800", + "gen_hash_zig": "sha256:0b168aaa160d077bfbcb66a19487dce330289422c64da51cec4154fd5ff7799a", "module": "HSLM", "ring": 12, "seal_schema": 2, - "sealed_at": "2026-10-17T17:34:44Z", + "sealed_at": "2026-10-17T18:02:24Z", "sealed_by": "t27c-bootstrap@0.5.2", - "spec_hash": "sha256:1ffe306efd7cc6b092829cbcfb7645cbadde70fed20b3c3a7e3cfedd3abd27ce", + "spec_hash": "sha256:df24bc0cb83391c7131fc940940bae4804f95fde9e854ae03065ec8caff85a08", "spec_path": "specs/nn/hslm.t27", "tests": { "failed": 0, "failing": [], "forced": false, - "passed": 26, - "total": 26, + "passed": 46, + "total": 46, "vacuous": 0 }, "toolchain": { diff --git a/specs/igla/own_model.t27 b/specs/igla/own_model.t27 index 42c1096680..5a4fb11fe1 100644 --- a/specs/igla/own_model.t27 +++ b/specs/igla/own_model.t27 @@ -22,9 +22,15 @@ ; - re-test on 2026-10-18, same t27c, after the PR that rewrote both HSLM specs (#7778): ; nn/hslm 26/26 and hslm/forward_pass 14/14 pass, 0 vacuous, so 6 of the 17 compile. The ; spec's backward pass is now real, each gradient checked against a finite difference. -; The spec is not the trained model: 6 layers at ctx 81, FFN 243-972-243 with GELU and no +; That spec was not the trained model: 6 layers at ctx 81, FFN 243-972-243 with GELU and no ; embedding or vocab head, where trinity-training ran 3 blocks, its best run at ctx 27, and an ; FFN 243-729-243 with ReLU, biases and a TWN scale. +; - re-test on 2026-10-18, same t27c, after the PR that rewrote nn/hslm to the trained model +; (#10062): nn/hslm 46/46 pass, 0 vacuous. It now has the trained shape (3 blocks, ctx 27, +; phi position encoding and phi-RoPE, a ReLU FFN 243-729-243 with biases and TWN alphas, the +; 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). ; - 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 @@ -39,9 +45,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, but that spec is not the trained -; model (see MEASURED); ternary STE, data loader, checkpoint I/O, heap and file I/O and -; `t27c --backend ml` do not exist yet. +; 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. ; Not claimed: no model quality beyond the numbers in SOURCES; no GIF or MP4, stills only. ; phi^2 + 1/phi^2 = 3 | TRINITY @@ -178,7 +184,7 @@ pub const STATUS : [27]u8 = [4, 4, 4, 3, 3, 3, 3, 0, 0, 4, 4, 4, 4, 4, 4, 3, 0, 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, 60, 40, 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, 70, 0, 0, 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; @@ -204,9 +210,9 @@ pub const ROWS : [20]str = [ "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 42 percent overall", - "Match the spec to the trained HSLM", - "Ternary STE, embedding and head", + "7 gaps, about 44 percent overall", + "HSLM spec: add the VSA path", + "Ternary STE with shadow weights", "Loader, checkpoint I/O, --backend ml", "" ]; @@ -313,10 +319,10 @@ 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) == 51; + assert mean_pct(9, 9) == 57; assert mean_pct(18, 9) == 30; - // the evidence slide's "about 42 percent overall" - assert mean_pct(0, CELL_COUNT) == 42; + // the evidence slide's "about 44 percent overall" + assert mean_pct(0, CELL_COUNT) == 44; } test the_camera_is_the_demo_camera { diff --git a/specs/igla/own_model_board.t27 b/specs/igla/own_model_board.t27 index abfb6f3639..6a5d0cf671 100644 --- a/specs/igla/own_model_board.t27 +++ b/specs/igla/own_model_board.t27 @@ -48,20 +48,20 @@ 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; 42%"; +pub const MSG : str = "1 published, 12 verified, 7 drafts, 7 gaps; 44%"; /// What MSG says once its numbers are taken out. pub const MSG_WORDS : str = " published, verified, drafts, gaps; %"; -pub const OVERALL_PCT : u8 = 42; -pub const PROGRESS : f64 = 0.42; +pub const OVERALL_PCT : u8 = 44; +pub const PROGRESS : f64 = 0.44; 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, 51, 30]; +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, 60, 40, 0, 0, 70, 15, 60, 40, 35, 0, 0, 0, 50]; +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]; /// The size of the compact metrics object, in bytes, computed from the keys and values above. pub const METRICS_BYTES : i32 = 622; @@ -273,19 +273,19 @@ 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; 42%", STATUS, overall) == false; - assert msg_matches("1 published, 12 verified, 7 drafts, 7 holes; 42%", STATUS, overall) == false; - assert msg_matches("1 published, 12 verified, 7 drafts, 7 gaps; 42% of 100", STATUS, overall) == false; - assert msg_matches(MSG, STATUS, 43) == false; + 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(MSG, STATUS, 45) == false; assert nth_number(MSG, 1) == 12; - assert nth_number(MSG, 4) == 42; + assert nth_number(MSG, 4) == 44; assert nth_number(MSG, 5) == -1; } test progress_is_the_overall_percent_over_a_hundred { assert progress_matches(PROGRESS, mean_pct(0, CELL_COUNT)); // a bar one point off, or past the end of the board's 0..1, is refused - assert progress_matches(0.43, mean_pct(0, CELL_COUNT)) == false; + assert progress_matches(0.45, mean_pct(0, CELL_COUNT)) == false; assert progress_matches(1.39, 139) == false; assert progress_matches(-0.01, -1) == false; assert JOB.len > 0; diff --git a/specs/nn/hslm.t27 b/specs/nn/hslm.t27 index f1795d2d54..35e5039825 100644 --- a/specs/nn/hslm.t27 +++ b/specs/nn/hslm.t27 @@ -1,31 +1,49 @@ // SPDX-License-Identifier: Apache-2.0 // t27/specs/nn/hslm.t27 // HSLM (Hierarchical Sacred Learning Model) Specification -// Ternary transformer: RMSNorm, causal multi-head attention with a KV cache, -// ternary feed-forward network, and the backward pass for one position. +// The ternary language model trinity-training trained: a token embedding with +// a phi position encoding, three Trinity blocks (sacred attention, then a +// ternary ReLU FFN), and a vocab-729 output head; and the backward pass for +// the last position, from the logits back to the embedding. // φ² + 1/φ² = 3 | TRINITY // -// This spec used to inline nn::attention, which reads constants::PHI_CUBED_INV; -// math::constants does not define it, so the generated Zig never compiled -// ("use of undeclared identifier 'constants'"). The layer is written out here -// instead, with the sacred attention scale 1 / HEAD_DIM^(phi^-3) as a number. +// Source of truth: gHashTag/trinity-training @ 3e9648c1817a97c5f9b8c727b0f8e2291928eead +// src/hslm/constants.zig VOCAB_SIZE 729, EMBED_DIM 243, HIDDEN_DIM 729, +// DEFAULT_BLOCKS 3, NUM_HEADS 3, HEAD_DIM 81 +// src/hslm/embedding.zig x = table[token] + phi position encoding: +// freq = phi^(-2(i/2)/243) * 3/pi, sin even, cos odd +// src/hslm/sacred_attention.zig RMSNorm with gamma (eps 1e-6), Q/K/V = alpha W n, +// phi-RoPE on Q and K (pairs i, i+40 per head, angle +// pos * phi^(-2i/81), dim 80 unrotated), causal softmax +// of q.k / 81^(phi^-3), out = x + alpha_o (ctx @ Wo) +// src/hslm/trinity_block.zig TernaryDense: h = relu(alpha_up (r @ W_up) + b_up), +// out = alpha_down (h @ W_down) + b_down + r; no norm +// src/hslm/model.zig head: n = x / sqrt(mean(x^2) + 1e-6), no gamma; +// logits = alpha_out (n @ W_out) / 243^(phi^-3) + b_out +// src/hslm/ste.zig applyAlpha: the TWN scale, one per matrix +// src/hslm/trainer.zig cross-entropy on the last position's logits +// Context 27: gHashTag/trinity @ ee512359254bcd7438bb6a79c8a4a7f8b554b620, +// docs/experiments/FOUND_EXPERIMENTS_SUMMARY.md: ctx 18 -> PPL 5.58, ctx 27 -> +// PPL 2.96 (run R5, the best), ctx 54 -> PPL 6.05; ctx must divide HEAD_DIM +// (EXP-012). constants.zig still says CONTEXT_LEN 81; the run used 27. // -// The four backward functions used to fill gradients with zeros. They now -// compute them, and the tests check every gradient against a finite -// difference of the loss: a central difference with a small step for inputs -// and RMSNorm gains, and the -1 / +1 trit flip for ternary weights (exact for -// the last layer's W2, where the loss is linear in the weight). +// What still differs from the trained model, named so nobody reads this spec +// as more than it is: +// - The VSA path is not here: per block the trained model also runs VSA +// attention over 1024-trit hypervectors, and when its consciousness gate +// fires (similarity > phi^-1) adds 0.1 * projectVsaToEmbed(reasoned) to +// the block output. +// - Floats: the trained model computes in f32, this spec in f64. +// - Training: the trainer's STE backward leaves the alphas out of the input +// gradients and stops at the embedding (the table is never trained); this +// spec computes the exact gradient, embedding included. Alphas and the +// ternary requantize step are inputs here, not learned. // -// Measured against the trained model (gHashTag/trinity-training, -// src/hslm/constants.zig and trinity_block.zig), this spec does not describe -// it yet: -// - NUM_LAYERS 6 here; the trained model has DEFAULT_BLOCKS 3. -// - CONTEXT_LEN 81 here; the best run (PPL 2.96) trained at context 27. -// - The FFN here is 243 -> 972 -> 243 with GELU and no bias; the trained -// block is 243 -> 729 -> 243 with ReLU, a TWN alpha scale per matrix, and -// a bias after each projection. -// - No token embedding and no vocab-729 output head here. -// - Attention matches: 3 heads x 81, pre-RMSNorm, sacred scale, residual. +// The backward pass is checked against central differences of the loss: on +// inputs, embedding rows, gains, biases and, for ternary weights, a real-valued +// nudge of one weight (the straight-through gradient is d loss / d weight as if +// it were real). Trit flips (-1 / +1) stay for W_out, where the loss is linear +// in the weight and the flip is exact, and for the last block's W_down. module HSLM { use base::types; @@ -34,20 +52,25 @@ module HSLM { // 1. HSLM Configuration // ═════════════════════════════════════════════════════════════════ - const NUM_LAYERS : usize = 6; // transformer layers - const NUM_HEADS : usize = 3; // attention heads per layer + const NUM_BLOCKS : usize = 3; // DEFAULT_BLOCKS: Trinity blocks + const NUM_HEADS : usize = 3; // attention heads per block const HEAD_DIM : usize = 81; // 3^4 per head const EMBED_DIM : usize = 243; // NUM_HEADS * HEAD_DIM - const FF_DIM : usize = 972; // 4 * EMBED_DIM - const CONTEXT_LEN : usize = 81; // KV cache positions - const VSA_DIM : usize = 1024; // VSA hypervector dimension + const FF_DIM : usize = 729; // HIDDEN_DIM, 3^6 + const VOCAB_SIZE : usize = 729; // 3^6 tokens + const CONTEXT_LEN : usize = 27; // 3^3, the PPL 2.96 run + const ROPE_PAIRS : usize = 40; // HEAD_DIM / 2; the last dim is unrotated + const VSA_DIM : usize = 1024; // VSA hypervector dimension (path not specified here) const RMS_EPS : f64 = 0.000001; // 1 / 81^(phi^-3), phi^-3 = 0.2360679775 const SACRED_ATTN_SCALE : f64 = 0.3543788557382518; - // sqrt(2 / pi) and the cubic term of the tanh GELU approximation - const SQRT_2_OVER_PI : f64 = 0.7978845608028654; - const GELU_CUBIC : f64 = 0.044715; + // 1 / 243^(phi^-3) + const SACRED_LOGIT_SCALE : f64 = 0.2734228160400333; + // ln(phi), for phi^-t = exp(-t ln phi) + const LN_PHI : f64 = 0.48121182505960347; + // 3 / pi, the position encoding's frequency scale + const TRINITY_SCALE : f64 = 0.954929658551372; // Layer phases const PHASE_NORM : u8 = 0; @@ -55,84 +78,102 @@ module HSLM { const PHASE_FFN : u8 = 2; const PHASE_RESIDUAL : u8 = 3; - // Activation functions - const ACT_RELU : u8 = 0; - const ACT_GELU : u8 = 1; - const ACT_SWISH : u8 = 2; - const ACT_TERNARY : u8 = 3; - // Training phases const PHASE_FORWARD : u8 = 0; const PHASE_BACKWARD : u8 = 1; const PHASE_UPDATE : u8 = 2; - // Weight matrices, for the test probes + // Weight matrices and biases, for the test probes const MAT_Q : u8 = 0; const MAT_K : u8 = 1; const MAT_V : u8 = 2; const MAT_O : u8 = 3; - const MAT_W1 : u8 = 4; - const MAT_W2 : u8 = 5; + const MAT_UP : u8 = 4; + const MAT_DOWN : u8 = 5; + const MAT_OUT : u8 = 6; + const MAT_NONE : u8 = 255; + const BIAS_UP : u8 = 0; + const BIAS_DOWN : u8 = 1; + const BIAS_OUT : u8 = 2; // ═════════════════════════════════════════════════════════════════ // 2. Weights and state // ═════════════════════════════════════════════════════════════════ - // Ternary weights, row-major in_dim x out_dim: w[j * out_dim + i] - struct LayerWeights { + // Ternary weights, row-major in_dim x out_dim: w[j * out_dim + i]. + // Each matrix has one TWN alpha: y = alpha * (x @ W). + struct BlockWeights { w_q : [EMBED_DIM * EMBED_DIM]Trit, w_k : [EMBED_DIM * EMBED_DIM]Trit, w_v : [EMBED_DIM * EMBED_DIM]Trit, w_o : [EMBED_DIM * EMBED_DIM]Trit, - w1 : [EMBED_DIM * FF_DIM]Trit, - w2 : [FF_DIM * EMBED_DIM]Trit, + w_up : [EMBED_DIM * FF_DIM]Trit, + w_down : [FF_DIM * EMBED_DIM]Trit, norm1_gamma : [EMBED_DIM]f64, - norm2_gamma : [EMBED_DIM]f64, + b_up : [FF_DIM]f64, + b_down : [EMBED_DIM]f64, + alpha_q : f64, + alpha_k : f64, + alpha_v : f64, + alpha_o : f64, + alpha_up : f64, + alpha_down : f64, } struct HSLMWeights { - layers : [NUM_LAYERS]LayerWeights, + embed : [VOCAB_SIZE * EMBED_DIM]f64, + blocks : [NUM_BLOCKS]BlockWeights, + w_out : [EMBED_DIM * VOCAB_SIZE]Trit, + b_out : [VOCAB_SIZE]f64, + alpha_out : f64, } var hslm_mode : u8 = PHASE_FORWARD; - // The model (about 4.3 MB of trits) lives in module storage, not on a stack + // The model (about 2 M trits and 180 k floats) lives in module storage var model : HSLMWeights = undefined; - // KV cache per layer: [layer][position][EMBED_DIM] - var kv_k : [NUM_LAYERS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; - var kv_v : [NUM_LAYERS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; - - // Activations of the last position each layer processed, kept for backward - var act_pos : [NUM_LAYERS]usize = undefined; - var act_x : [NUM_LAYERS * EMBED_DIM]f64 = undefined; // layer input - var act_n1 : [NUM_LAYERS * EMBED_DIM]f64 = undefined; // RMSNorm 1 output - var act_rms1 : [NUM_LAYERS]f64 = undefined; - var act_q : [NUM_LAYERS * EMBED_DIM]f64 = undefined; - var act_attn : [NUM_LAYERS * NUM_HEADS * CONTEXT_LEN]f64 = undefined; // softmax - var act_ctx : [NUM_LAYERS * EMBED_DIM]f64 = undefined; // attention context - var act_r : [NUM_LAYERS * EMBED_DIM]f64 = undefined; // input + attention - var act_n2 : [NUM_LAYERS * EMBED_DIM]f64 = undefined; // RMSNorm 2 output - var act_rms2 : [NUM_LAYERS]f64 = undefined; - var act_pre : [NUM_LAYERS * FF_DIM]f64 = undefined; // n2 @ W1 - var act_h : [NUM_LAYERS * FF_DIM]f64 = undefined; // gelu(pre) - var act_out : [NUM_LAYERS * EMBED_DIM]f64 = undefined; // layer output - - // Weight gradients (straight-through: d loss / d weight as if it were real), - // accumulated until zero_weight_gradients - var gw_q : [NUM_LAYERS * EMBED_DIM * EMBED_DIM]f64 = undefined; - var gw_k : [NUM_LAYERS * EMBED_DIM * EMBED_DIM]f64 = undefined; - var gw_v : [NUM_LAYERS * EMBED_DIM * EMBED_DIM]f64 = undefined; - var gw_o : [NUM_LAYERS * EMBED_DIM * EMBED_DIM]f64 = undefined; - var gw_1 : [NUM_LAYERS * EMBED_DIM * FF_DIM]f64 = undefined; - var gw_2 : [NUM_LAYERS * FF_DIM * EMBED_DIM]f64 = undefined; - var g_gamma1 : [NUM_LAYERS * EMBED_DIM]f64 = undefined; - var g_gamma2 : [NUM_LAYERS * EMBED_DIM]f64 = undefined; - - // Backward scratch for the layer being processed - var bw_out : [EMBED_DIM]f64 = undefined; // d loss / d layer output - var bw_h : [FF_DIM]f64 = undefined; // d loss / d gelu output - var bw_pre : [FF_DIM]f64 = undefined; // d loss / d W1 output + // KV cache per block: [block][position][EMBED_DIM]; keys after RoPE + var kv_k : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; + var kv_v : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; + + // Activations of the last position each block processed, kept for backward + var act_token : usize = 0; + var act_pos : [NUM_BLOCKS]usize = undefined; + var act_x : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; // block input + var act_n1 : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; // RMSNorm output + var act_rms1 : [NUM_BLOCKS]f64 = undefined; + var act_q : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; // query after RoPE + var act_attn : [NUM_BLOCKS * NUM_HEADS * CONTEXT_LEN]f64 = undefined; // softmax + var act_ctx : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; // attention context + var act_r : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; // input + attention + var act_pre : [NUM_BLOCKS * FF_DIM]f64 = undefined; // alpha_up (r @ W_up) + b_up + var act_h : [NUM_BLOCKS * FF_DIM]f64 = undefined; // relu(pre) + var act_out : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; // block output + var act_hn : [EMBED_DIM]f64 = undefined; // head RMSNorm output + var act_hrms : f64 = 1.0; + + // Gradients (straight-through for trits: d loss / d weight as if it were + // real), accumulated until zero_weight_gradients + var gw_q : [NUM_BLOCKS * EMBED_DIM * EMBED_DIM]f64 = undefined; + var gw_k : [NUM_BLOCKS * EMBED_DIM * EMBED_DIM]f64 = undefined; + var gw_v : [NUM_BLOCKS * EMBED_DIM * EMBED_DIM]f64 = undefined; + var gw_o : [NUM_BLOCKS * EMBED_DIM * EMBED_DIM]f64 = undefined; + var gw_up : [NUM_BLOCKS * EMBED_DIM * FF_DIM]f64 = undefined; + var gw_down : [NUM_BLOCKS * FF_DIM * EMBED_DIM]f64 = undefined; + var gw_out : [EMBED_DIM * VOCAB_SIZE]f64 = undefined; + var g_gamma1 : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; + var g_b_up : [NUM_BLOCKS * FF_DIM]f64 = undefined; + var g_b_down : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; + var g_b_out : [VOCAB_SIZE]f64 = undefined; + var g_embed : [VOCAB_SIZE * EMBED_DIM]f64 = undefined; + + // Backward scratch for the block being processed + var bw_out : [EMBED_DIM]f64 = undefined; // d loss / d block output + var bw_h : [FF_DIM]f64 = undefined; // d loss / d relu output + var bw_pre : [FF_DIM]f64 = undefined; // d loss / d pre-activation + var bw_sf : [FF_DIM]f64 = undefined; // an alpha-scaled gradient, FF_DIM wide + var bw_se : [EMBED_DIM]f64 = undefined; // an alpha-scaled gradient, EMBED_DIM wide var bw_n : [EMBED_DIM]f64 = undefined; // d loss / d RMSNorm output var bw_tmp : [EMBED_DIM]f64 = undefined; var bw_r : [EMBED_DIM]f64 = undefined; // d loss / d (input + attention) @@ -141,7 +182,16 @@ module HSLM { var bw_k : [EMBED_DIM]f64 = undefined; // current position's key var bw_v : [EMBED_DIM]f64 = undefined; // current position's value var bw_da : [CONTEXT_LEN]f64 = undefined; - var bw_x : [EMBED_DIM]f64 = undefined; // d loss / d layer input + var bw_x : [EMBED_DIM]f64 = undefined; // d loss / d block input + var bw_dz : [VOCAB_SIZE]f64 = undefined; // d loss / d (n @ W_out) + + // A real-valued nudge of one ternary weight, for the gradient checks: + // the forward pass computes as if weight `nudge_index` of matrix + // `nudge_which` in block `nudge_block` were w + nudge_delta. Off by default. + var nudge_which : u8 = MAT_NONE; + var nudge_block : usize = 0; + var nudge_index : usize = 0; + var nudge_delta : f64 = 0.0; // ═════════════════════════════════════════════════════════════════ // 3. Kernels @@ -155,6 +205,14 @@ module HSLM { } } + fn scale_into(x: []const f64, alpha: f64, out: []f64) -> void { + var i : usize = 0; + while (i < x.len) { + out[i] = alpha * x[i]; + i = i + 1; + } + } + fn rms_of(x: []const f64) -> f64 { var sum_squares : f64 = 0.0; var i : usize = 0; @@ -204,6 +262,22 @@ module HSLM { } } + // The head's RMSNorm has no gamma: y_i = x_i / r + fn rms_norm_plain_backward(x: []const f64, rms: f64, dy: []const f64, dx: []f64) -> void { + const n = @as(f64, @floatFromInt(x.len)); + var s : f64 = 0.0; + var i : usize = 0; + while (i < x.len) { + s = s + dy[i] * x[i]; + i = i + 1; + } + i = 0; + while (i < x.len) { + dx[i] = dy[i] / rms - x[i] * s / (n * rms * rms * rms); + 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; @@ -257,49 +331,161 @@ module HSLM { } } - // GELU, tanh approximation: 0.5 x (1 + tanh(sqrt(2/pi) (x + 0.044715 x^3))) - fn gelu(x: f64) -> f64 { - const inner = SQRT_2_OVER_PI * (x + GELU_CUBIC * x * x * x); - return 0.5 * x * (1.0 + tanh(inner)); + // Applies the gradient-check nudge, if it targets this matrix: the output + // moves as if the weight were w + nudge_delta. Runs before the alpha scale. + fn nudge(which: u8, block_idx: usize, input: []const f64, output: []f64, out_dim: usize) -> void { + if (which == nudge_which and block_idx == nudge_block) { + const j = nudge_index / out_dim; + const i = nudge_index % out_dim; + output[i] = output[i] + nudge_delta * input[j]; + } + } + + fn relu(x: f64) -> f64 { + if (x > 0.0) { + return x; + } + return 0.0; + } + + // d relu / dx: 1 above zero, 0 at and below it (the trainer zeroes the + // gradient where the cached output is 0) + fn relu_derivative(x: f64) -> f64 { + if (x > 0.0) { + return 1.0; + } + return 0.0; } - // d gelu / dx = 0.5 (1 + t) + 0.5 x (1 - t^2) sqrt(2/pi) (1 + 3 * 0.044715 x^2) - fn gelu_derivative(x: f64) -> f64 { - const inner = SQRT_2_OVER_PI * (x + GELU_CUBIC * x * x * x); - const t = tanh(inner); - return 0.5 * (1.0 + t) + 0.5 * x * (1.0 - t * t) * SQRT_2_OVER_PI * (1.0 + 3.0 * GELU_CUBIC * x * x); + // embedding.zig: t = 2 (i / 2) / 243, freq = phi^-t * 3/pi, angle = position * freq; + // sin for even i, cos for odd i + fn pos_encoding(position: usize, i: usize) -> f64 { + const t = @as(f64, @floatFromInt(2 * (i / 2))) / @as(f64, @floatFromInt(EMBED_DIM)); + const angle = @as(f64, @floatFromInt(position)) * exp(0.0 - t * LN_PHI) * TRINITY_SCALE; + if (i % 2 == 0) { + return sin(angle); + } + return cos(angle); } - fn gelu_activation(x: []f64) -> void { + // x = table[token] + pos_encoding(position) + fn embed_into(token: usize, position: usize, out: []f64) -> void { var i : usize = 0; - while (i < x.len) { - x[i] = gelu(x[i]); + while (i < EMBED_DIM) { + out[i] = model.embed[token * EMBED_DIM + i] + pos_encoding(position, i); i = i + 1; } } + // phi-RoPE angle of pair i: position * phi^(-2 i / HEAD_DIM) + fn rope_angle(position: usize, i: usize) -> f64 { + const t = 2.0 * @as(f64, @floatFromInt(i)) / @as(f64, @floatFromInt(HEAD_DIM)); + return @as(f64, @floatFromInt(position)) * exp(0.0 - t * LN_PHI); + } + + // Rotates pairs (h + i, h + i + 40) of every head by the position's angles; + // inverse rotates the other way, which is the transpose (backward) + fn rope(x: []f64, position: usize, inverse: bool) -> void { + var i : usize = 0; + while (i < ROPE_PAIRS) { + const angle = rope_angle(position, i); + const c = cos(angle); + var s = sin(angle); + if (inverse) { + s = 0.0 - s; + } + var head : usize = 0; + while (head < NUM_HEADS) { + const i0 = head * HEAD_DIM + i; + const i1 = i0 + ROPE_PAIRS; + const x0 = x[i0]; + const x1 = x[i1]; + x[i0] = x0 * c - x1 * s; + x[i1] = x0 * s + x1 * c; + head = head + 1; + } + i = i + 1; + } + } + + // loss = log(sum_j exp(l_j)) - l_target + fn cross_entropy(logits: []const f64, target: usize) -> f64 { + var max_logit : f64 = logits[0]; + var j : usize = 1; + while (j < logits.len) { + if (logits[j] > max_logit) { + max_logit = logits[j]; + } + j = j + 1; + } + var sum : f64 = 0.0; + j = 0; + while (j < logits.len) { + sum = sum + exp(logits[j] - max_logit); + j = j + 1; + } + return log(sum) + max_logit - logits[target]; + } + + // d loss / d l_j = softmax_j - [j == target] + fn cross_entropy_backward(logits: []const f64, target: usize, dl: []f64) -> void { + var max_logit : f64 = logits[0]; + var j : usize = 1; + while (j < logits.len) { + if (logits[j] > max_logit) { + max_logit = logits[j]; + } + j = j + 1; + } + var sum : f64 = 0.0; + j = 0; + while (j < logits.len) { + dl[j] = exp(logits[j] - max_logit); + sum = sum + dl[j]; + j = j + 1; + } + j = 0; + while (j < logits.len) { + dl[j] = dl[j] / sum; + j = j + 1; + } + dl[target] = dl[target] - 1.0; + } + // ═════════════════════════════════════════════════════════════════ // 4. Forward pass (one position, causal, KV cache) // ═════════════════════════════════════════════════════════════════ - // q, k, v = n1 @ Wq, Wk, Wv; k and v go into the cache at `position`; - // per head, softmax(q . k_j * scale) over j <= position weights the v_j; - // out = context @ Wo - fn attention_forward(layer_idx: usize, position: usize, out: []f64) -> void { - const w = &model.layers[layer_idx]; - const eo = layer_idx * EMBED_DIM; - const ko = (layer_idx * CONTEXT_LEN + position) * EMBED_DIM; - ternary_matmul(act_n1[eo..eo + EMBED_DIM], &w.w_q, act_q[eo..eo + EMBED_DIM], EMBED_DIM, EMBED_DIM); - ternary_matmul(act_n1[eo..eo + EMBED_DIM], &w.w_k, kv_k[ko..ko + EMBED_DIM], EMBED_DIM, EMBED_DIM); - ternary_matmul(act_n1[eo..eo + EMBED_DIM], &w.w_v, kv_v[ko..ko + EMBED_DIM], EMBED_DIM, EMBED_DIM); + // q, k, v = alpha (n1 @ Wq, Wk, Wv); RoPE on q and k; k and v go into the + // cache at `position`; per head, softmax(q . k_j * scale) over j <= position + // weights the v_j; out = alpha_o (context @ Wo) + fn attention_forward(block_idx: usize, position: usize, out: []f64) -> void { + const w = &model.blocks[block_idx]; + const eo = block_idx * EMBED_DIM; + const ko = (block_idx * CONTEXT_LEN + position) * EMBED_DIM; + const n1 = act_n1[eo..eo + EMBED_DIM]; + const q = act_q[eo..eo + EMBED_DIM]; + const k = kv_k[ko..ko + EMBED_DIM]; + const v = kv_v[ko..ko + EMBED_DIM]; + ternary_matmul(n1, &w.w_q, q, EMBED_DIM, EMBED_DIM); + nudge(MAT_Q, block_idx, n1, q, EMBED_DIM); + scale_into(q, w.alpha_q, q); + ternary_matmul(n1, &w.w_k, k, EMBED_DIM, EMBED_DIM); + nudge(MAT_K, block_idx, n1, k, EMBED_DIM); + scale_into(k, w.alpha_k, k); + ternary_matmul(n1, &w.w_v, v, EMBED_DIM, EMBED_DIM); + nudge(MAT_V, block_idx, n1, v, EMBED_DIM); + scale_into(v, w.alpha_v, v); + rope(q, position, false); + rope(k, position, false); var head : usize = 0; while (head < NUM_HEADS) { const ho = head * HEAD_DIM; - const ao = (layer_idx * NUM_HEADS + head) * CONTEXT_LEN; + const ao = (block_idx * NUM_HEADS + head) * CONTEXT_LEN; var j : usize = 0; while (j <= position) { - const kj = (layer_idx * CONTEXT_LEN + j) * EMBED_DIM + ho; + const kj = (block_idx * CONTEXT_LEN + j) * EMBED_DIM + ho; var dot : f64 = 0.0; var d : usize = 0; while (d < HEAD_DIM) { @@ -334,7 +520,7 @@ module HSLM { var acc : f64 = 0.0; var m : usize = 0; while (m <= position) { - acc = acc + act_attn[ao + m] * kv_v[(layer_idx * CONTEXT_LEN + m) * EMBED_DIM + ho + c]; + acc = acc + act_attn[ao + m] * kv_v[(block_idx * CONTEXT_LEN + m) * EMBED_DIM + ho + c]; m = m + 1; } act_ctx[eo + ho + c] = acc; @@ -343,41 +529,53 @@ module HSLM { head = head + 1; } ternary_matmul(act_ctx[eo..eo + EMBED_DIM], &w.w_o, out, EMBED_DIM, EMBED_DIM); + nudge(MAT_O, block_idx, act_ctx[eo..eo + EMBED_DIM], out, EMBED_DIM); + scale_into(out, w.alpha_o, out); } - // f = gelu(n2 @ W1) @ W2 - fn ffn_forward(layer_idx: usize, out: []f64) -> void { - const w = &model.layers[layer_idx]; - const eo = layer_idx * EMBED_DIM; - const fo = layer_idx * FF_DIM; - ternary_matmul(act_n2[eo..eo + EMBED_DIM], &w.w1, act_pre[fo..fo + FF_DIM], EMBED_DIM, FF_DIM); + // TernaryDense: h = relu(alpha_up (r @ W_up) + b_up); f = alpha_down (h @ W_down) + b_down + fn ffn_forward(block_idx: usize, out: []f64) -> void { + const w = &model.blocks[block_idx]; + const eo = block_idx * EMBED_DIM; + const fo = block_idx * FF_DIM; + const r = act_r[eo..eo + EMBED_DIM]; + const pre = act_pre[fo..fo + FF_DIM]; + ternary_matmul(r, &w.w_up, pre, EMBED_DIM, FF_DIM); + nudge(MAT_UP, block_idx, r, pre, FF_DIM); var i : usize = 0; while (i < FF_DIM) { - act_h[fo + i] = gelu(act_pre[fo + i]); + act_pre[fo + i] = w.alpha_up * act_pre[fo + i] + w.b_up[i]; + act_h[fo + i] = relu(act_pre[fo + i]); + i = i + 1; + } + ternary_matmul(act_h[fo..fo + FF_DIM], &w.w_down, out, FF_DIM, EMBED_DIM); + nudge(MAT_DOWN, block_idx, act_h[fo..fo + FF_DIM], out, EMBED_DIM); + i = 0; + while (i < EMBED_DIM) { + out[i] = w.alpha_down * out[i] + w.b_down[i]; i = i + 1; } - ternary_matmul(act_h[fo..fo + FF_DIM], &w.w2, out, FF_DIM, EMBED_DIM); } - // r = x + attention(rms_norm(x)); out = r + ffn(rms_norm(r)) - fn hslm_layer_forward(layer_idx: usize, position: usize, input: []const f64, output: []f64) -> void { - const w = &model.layers[layer_idx]; - const eo = layer_idx * EMBED_DIM; - act_pos[layer_idx] = position; + // r = x + attention(rms_norm(x)); out = r + ffn(r). The trained block has + // no norm in front of the FFN. + fn hslm_block_forward(block_idx: usize, position: usize, input: []const f64, output: []f64) -> void { + const w = &model.blocks[block_idx]; + const eo = block_idx * EMBED_DIM; + act_pos[block_idx] = position; var i : usize = 0; while (i < EMBED_DIM) { act_x[eo + i] = input[i]; i = i + 1; } - act_rms1[layer_idx] = rms_norm_into(input, &w.norm1_gamma, act_n1[eo..eo + EMBED_DIM]); - attention_forward(layer_idx, position, output); + act_rms1[block_idx] = rms_norm_into(input, &w.norm1_gamma, act_n1[eo..eo + EMBED_DIM]); + attention_forward(block_idx, position, output); i = 0; while (i < EMBED_DIM) { act_r[eo + i] = input[i] + output[i]; i = i + 1; } - act_rms2[layer_idx] = rms_norm_into(act_r[eo..eo + EMBED_DIM], &w.norm2_gamma, act_n2[eo..eo + EMBED_DIM]); - ffn_forward(layer_idx, output); + ffn_forward(block_idx, output); i = 0; while (i < EMBED_DIM) { output[i] = act_r[eo + i] + output[i]; @@ -385,140 +583,172 @@ module HSLM { } } - // All layers at one position; earlier positions must already be in the cache - fn hslm_forward(input: []const f64, position: usize, output: []f64) -> void { - hslm_layer_forward(0, position, input, act_out[0..EMBED_DIM]); - var layer_idx : usize = 1; - while (layer_idx < NUM_LAYERS) { - const prev = (layer_idx - 1) * EMBED_DIM; - const cur = layer_idx * EMBED_DIM; - hslm_layer_forward(layer_idx, position, act_out[prev..prev + EMBED_DIM], act_out[cur..cur + EMBED_DIM]); - layer_idx = layer_idx + 1; - } - const last = (NUM_LAYERS - 1) * EMBED_DIM; + // n = x / rms(x); logits = alpha_out (n @ W_out) * SACRED_LOGIT_SCALE + b_out + fn head_forward(x: []const f64, logits: []f64) -> void { + act_hrms = rms_of(x); var i : usize = 0; while (i < EMBED_DIM) { - output[i] = act_out[last + i]; + act_hn[i] = x[i] / act_hrms; i = i + 1; } + ternary_matmul(&act_hn, &model.w_out, logits, EMBED_DIM, VOCAB_SIZE); + nudge(MAT_OUT, 0, &act_hn, logits, VOCAB_SIZE); + var j : usize = 0; + while (j < VOCAB_SIZE) { + logits[j] = model.alpha_out * logits[j] * SACRED_LOGIT_SCALE + model.b_out[j]; + j = j + 1; + } + } + + // Token at one position to logits; earlier positions must already be in + // the cache + fn hslm_forward(token: usize, position: usize, logits: []f64) -> void { + act_token = token; + embed_into(token, position, &bw_tmp); + hslm_block_forward(0, position, &bw_tmp, act_out[0..EMBED_DIM]); + var block_idx : usize = 1; + while (block_idx < NUM_BLOCKS) { + const prev = (block_idx - 1) * EMBED_DIM; + const cur = block_idx * EMBED_DIM; + hslm_block_forward(block_idx, position, act_out[prev..prev + EMBED_DIM], act_out[cur..cur + EMBED_DIM]); + block_idx = block_idx + 1; + } + const last = (NUM_BLOCKS - 1) * EMBED_DIM; + head_forward(act_out[last..last + EMBED_DIM], logits); } // ═════════════════════════════════════════════════════════════════ - // 5. Backward pass (the last position each layer processed) + // 5. Backward pass (the last position each block processed) // ═════════════════════════════════════════════════════════════════ // - // The gradient is that of the current position's output. Keys and values + // The gradient is that of the current position's logits. Keys and values // of earlier positions are cached constants: no gradient flows into them. // The `grads` argument of the backward functions carries the ternary weights - // the gradient flows back through; the f64 gradients land in gw_*, g_gamma*. + // the gradient flows back through; the f64 gradients land in gw_*, g_*. fn zero_weight_gradients() -> void { fill(&gw_q, 0.0); fill(&gw_k, 0.0); fill(&gw_v, 0.0); fill(&gw_o, 0.0); - fill(&gw_1, 0.0); - fill(&gw_2, 0.0); + fill(&gw_up, 0.0); + fill(&gw_down, 0.0); + fill(&gw_out, 0.0); fill(&g_gamma1, 0.0); - fill(&g_gamma2, 0.0); + fill(&g_b_up, 0.0); + fill(&g_b_down, 0.0); + fill(&g_b_out, 0.0); + fill(&g_embed, 0.0); } - // grad_input = d loss / d input, given grad_output = d loss / d output - fn hslm_backward(grad_output: []const f64, grad_input: []f64) -> void { - var i : usize = 0; - while (i < EMBED_DIM) { - bw_out[i] = grad_output[i]; - i = i + 1; - } - var layer_idx : usize = NUM_LAYERS; - while (layer_idx > 0) { - layer_idx = layer_idx - 1; - hslm_layer_backward(&model, layer_idx); - i = 0; - while (i < EMBED_DIM) { - bw_out[i] = bw_x[i]; - i = i + 1; + // grad_input = d loss / d embedded input, given grad_logits = d loss / d logits; + // the input gradient also lands in the token's embedding row + fn hslm_backward(grad_logits: []const f64, grad_input: []f64) -> void { + head_backward(&model, grad_logits); + var block_idx : usize = NUM_BLOCKS; + while (block_idx > 0) { + block_idx = block_idx - 1; + hslm_block_backward(&model, block_idx); + var c : usize = 0; + while (c < EMBED_DIM) { + bw_out[c] = bw_x[c]; + c = c + 1; } } - i = 0; + var i : usize = 0; while (i < EMBED_DIM) { grad_input[i] = bw_x[i]; + g_embed[act_token * EMBED_DIM + i] = g_embed[act_token * EMBED_DIM + i] + bw_x[i]; i = i + 1; } } - // bw_out -> bw_x through FFN, RMSNorm 2, attention, RMSNorm 1 and both residuals - fn hslm_layer_backward(grads: *HSLMWeights, layer_idx: usize) -> void { - ffn_backward(grads, layer_idx); - attention_backward(grads, layer_idx); + // logits -> bw_out (d loss / d last block output): bias, W_out, plain RMSNorm + fn head_backward(grads: *HSLMWeights, grad_logits: []const f64) -> void { + var j : usize = 0; + while (j < VOCAB_SIZE) { + g_b_out[j] = g_b_out[j] + grad_logits[j]; + bw_dz[j] = grads.alpha_out * SACRED_LOGIT_SCALE * grad_logits[j]; + j = j + 1; + } + outer_add(&act_hn, &bw_dz, &gw_out, EMBED_DIM, VOCAB_SIZE); + fill(&bw_n, 0.0); + ternary_matmul_t_add(&bw_dz, &grads.w_out, &bw_n, EMBED_DIM, VOCAB_SIZE); + const last = (NUM_BLOCKS - 1) * EMBED_DIM; + rms_norm_plain_backward(act_out[last..last + EMBED_DIM], act_hrms, &bw_n, &bw_out); + } + + // bw_out -> bw_x through the FFN, attention, RMSNorm and both residuals + fn hslm_block_backward(grads: *HSLMWeights, block_idx: usize) -> void { + ffn_backward(grads, block_idx); + attention_backward(grads, block_idx); } - // bw_out -> bw_r = bw_out + d loss / d r through the FFN and RMSNorm 2 - fn ffn_backward(grads: *HSLMWeights, layer_idx: usize) -> void { - const w = &grads.layers[layer_idx]; - const eo = layer_idx * EMBED_DIM; - ffn_backward_w2(grads, layer_idx); - gelu_backward(layer_idx); - ffn_backward_w1(grads, layer_idx); - rms_norm_backward(act_r[eo..eo + EMBED_DIM], &w.norm2_gamma, act_rms2[layer_idx], &bw_n, &bw_tmp, g_gamma2[eo..eo + EMBED_DIM]); + // bw_out -> bw_r = bw_out + d loss / d r through the FFN + fn ffn_backward(grads: *HSLMWeights, block_idx: usize) -> void { + const w = &grads.blocks[block_idx]; + const eo = block_idx * EMBED_DIM; + const fo = block_idx * FF_DIM; + const go = block_idx * FF_DIM * EMBED_DIM; + // f = alpha_down (h @ W_down) + b_down var i : usize = 0; while (i < EMBED_DIM) { - bw_r[i] = bw_out[i] + bw_tmp[i]; + g_b_down[eo + i] = g_b_down[eo + i] + bw_out[i]; i = i + 1; } - } - - // f = h @ W2: dW2 += h (x) bw_out, bw_h = W2^T bw_out - fn ffn_backward_w2(grads: *HSLMWeights, layer_idx: usize) -> void { - const w = &grads.layers[layer_idx]; - const fo = layer_idx * FF_DIM; - const go = layer_idx * FF_DIM * EMBED_DIM; - outer_add(act_h[fo..fo + FF_DIM], &bw_out, gw_2[go..go + FF_DIM * EMBED_DIM], FF_DIM, EMBED_DIM); + scale_into(&bw_out, w.alpha_down, &bw_se); + outer_add(act_h[fo..fo + FF_DIM], &bw_se, gw_down[go..go + FF_DIM * EMBED_DIM], FF_DIM, EMBED_DIM); fill(&bw_h, 0.0); - ternary_matmul_t_add(&bw_out, &w.w2, &bw_h, FF_DIM, EMBED_DIM); + ternary_matmul_t_add(&bw_se, &w.w_down, &bw_h, FF_DIM, EMBED_DIM); + relu_backward(block_idx); + // pre = alpha_up (r @ W_up) + b_up + i = 0; + while (i < FF_DIM) { + g_b_up[fo + i] = g_b_up[fo + i] + bw_pre[i]; + i = i + 1; + } + scale_into(&bw_pre, w.alpha_up, &bw_sf); + outer_add(act_r[eo..eo + EMBED_DIM], &bw_sf, gw_up[go..go + EMBED_DIM * FF_DIM], EMBED_DIM, FF_DIM); + i = 0; + while (i < EMBED_DIM) { + bw_r[i] = bw_out[i]; + i = i + 1; + } + ternary_matmul_t_add(&bw_sf, &w.w_up, &bw_r, EMBED_DIM, FF_DIM); } - // h = gelu(pre): bw_pre = gelu'(pre) * bw_h - fn gelu_backward(layer_idx: usize) -> void { - const fo = layer_idx * FF_DIM; + // h = relu(pre): bw_pre = relu'(pre) * bw_h + fn relu_backward(block_idx: usize) -> void { + const fo = block_idx * FF_DIM; var i : usize = 0; while (i < FF_DIM) { - bw_pre[i] = gelu_derivative(act_pre[fo + i]) * bw_h[i]; + bw_pre[i] = relu_derivative(act_pre[fo + i]) * bw_h[i]; i = i + 1; } } - // pre = n2 @ W1: dW1 += n2 (x) bw_pre, bw_n = W1^T bw_pre - fn ffn_backward_w1(grads: *HSLMWeights, layer_idx: usize) -> void { - const w = &grads.layers[layer_idx]; - const eo = layer_idx * EMBED_DIM; - const go = layer_idx * EMBED_DIM * FF_DIM; - outer_add(act_n2[eo..eo + EMBED_DIM], &bw_pre, gw_1[go..go + EMBED_DIM * FF_DIM], EMBED_DIM, FF_DIM); - fill(&bw_n, 0.0); - ternary_matmul_t_add(&bw_pre, &w.w1, &bw_n, EMBED_DIM, FF_DIM); - } - - // bw_r -> bw_x: Wo, softmax, the current position's q, k, v, Wq/Wk/Wv, - // RMSNorm 1, and the residual around the block - fn attention_backward(grads: *HSLMWeights, layer_idx: usize) -> void { - const w = &grads.layers[layer_idx]; - const eo = layer_idx * EMBED_DIM; - const go = layer_idx * EMBED_DIM * EMBED_DIM; - const position = act_pos[layer_idx]; + // bw_r -> bw_x: Wo, softmax, the current position's q, k, v, RoPE, + // Wq/Wk/Wv, RMSNorm, and the residual around the block + fn attention_backward(grads: *HSLMWeights, block_idx: usize) -> void { + const w = &grads.blocks[block_idx]; + const eo = block_idx * EMBED_DIM; + const go = block_idx * EMBED_DIM * EMBED_DIM; + const position = act_pos[block_idx]; - outer_add(act_ctx[eo..eo + EMBED_DIM], &bw_r, gw_o[go..go + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM); + scale_into(&bw_r, w.alpha_o, &bw_se); + outer_add(act_ctx[eo..eo + EMBED_DIM], &bw_se, gw_o[go..go + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM); fill(&bw_ctx, 0.0); - ternary_matmul_t_add(&bw_r, &w.w_o, &bw_ctx, EMBED_DIM, EMBED_DIM); + ternary_matmul_t_add(&bw_se, &w.w_o, &bw_ctx, EMBED_DIM, EMBED_DIM); var head : usize = 0; while (head < NUM_HEADS) { const ho = head * HEAD_DIM; - const ao = (layer_idx * NUM_HEADS + head) * CONTEXT_LEN; + const ao = (block_idx * NUM_HEADS + head) * CONTEXT_LEN; // da_j = dctx . v_j, then the softmax: ds_j = a_j (da_j - sum_m a_m da_m) var mean_da : f64 = 0.0; var j : usize = 0; while (j <= position) { - const vj = (layer_idx * CONTEXT_LEN + j) * EMBED_DIM + ho; + const vj = (block_idx * CONTEXT_LEN + j) * EMBED_DIM + ho; var dot : f64 = 0.0; var d : usize = 0; while (d < HEAD_DIM) { @@ -540,7 +770,7 @@ module HSLM { var acc : f64 = 0.0; var m : usize = 0; while (m <= position) { - acc = acc + bw_da[m] * kv_k[(layer_idx * CONTEXT_LEN + m) * EMBED_DIM + ho + c]; + acc = acc + bw_da[m] * kv_k[(block_idx * CONTEXT_LEN + m) * EMBED_DIM + ho + c]; m = m + 1; } bw_q[ho + c] = SACRED_ATTN_SCALE * acc; @@ -551,15 +781,24 @@ module HSLM { head = head + 1; } - outer_add(act_n1[eo..eo + EMBED_DIM], &bw_q, gw_q[go..go + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM); - outer_add(act_n1[eo..eo + EMBED_DIM], &bw_k, gw_k[go..go + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM); - outer_add(act_n1[eo..eo + EMBED_DIM], &bw_v, gw_v[go..go + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM); + // Back through RoPE (a rotation: its transpose is the inverse), then + // through the alphas: d (x @ W) = alpha * d y + rope(&bw_q, position, true); + rope(&bw_k, position, true); + scale_into(&bw_q, w.alpha_q, &bw_q); + scale_into(&bw_k, w.alpha_k, &bw_k); + scale_into(&bw_v, w.alpha_v, &bw_v); + + const n1 = act_n1[eo..eo + EMBED_DIM]; + outer_add(n1, &bw_q, gw_q[go..go + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM); + outer_add(n1, &bw_k, gw_k[go..go + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM); + outer_add(n1, &bw_v, gw_v[go..go + EMBED_DIM * EMBED_DIM], EMBED_DIM, EMBED_DIM); fill(&bw_n, 0.0); ternary_matmul_t_add(&bw_q, &w.w_q, &bw_n, EMBED_DIM, EMBED_DIM); ternary_matmul_t_add(&bw_k, &w.w_k, &bw_n, EMBED_DIM, EMBED_DIM); ternary_matmul_t_add(&bw_v, &w.w_v, &bw_n, EMBED_DIM, EMBED_DIM); - rms_norm_backward(act_x[eo..eo + EMBED_DIM], &w.norm1_gamma, act_rms1[layer_idx], &bw_n, &bw_tmp, g_gamma1[eo..eo + EMBED_DIM]); + rms_norm_backward(act_x[eo..eo + EMBED_DIM], &w.norm1_gamma, act_rms1[block_idx], &bw_n, &bw_tmp, g_gamma1[eo..eo + EMBED_DIM]); var i : usize = 0; while (i < EMBED_DIM) { bw_x[i] = bw_r[i] + bw_tmp[i]; @@ -586,13 +825,21 @@ module HSLM { // 7. Test probes // ═════════════════════════════════════════════════════════════════ - var probe_x : [EMBED_DIM]f64 = undefined; - var probe_out : [EMBED_DIM]f64 = undefined; - var probe_up : [EMBED_DIM]f64 = undefined; + // Position 0 holds TOKEN_FIRST; position 1, the probed one, TOKEN_PROBE. + // A target of VOCAB_SIZE means the linear loss sum_j logits_j * upstream(j); + // a smaller one, cross-entropy against that token. + const TOKEN_FIRST : usize = 5; + const TOKEN_PROBE : usize = 42; + const LINEAR_LOSS : usize = 729; + const PROBE_TARGET : usize = 3; + + var probe_vec : [EMBED_DIM]f64 = undefined; + var probe_logits : [VOCAB_SIZE]f64 = undefined; + var probe_dl : [VOCAB_SIZE]f64 = undefined; var probe_gin : [EMBED_DIM]f64 = undefined; - // About 2 in 9 trits nonzero, a different pattern per matrix and layer. - // A multiplicative hash mod a prime: 243 and 972 are multiples of 9, so a + // About 2 in 9 trits nonzero, a different pattern per matrix and block. + // A multiplicative hash mod a prime: 243 and 729 are multiples of 9, so a // linear pattern mod 9 repeats down every column (measured: whole columns // of one trit, and weight gradients that were exactly zero). fn trit_pattern(seed: usize) -> Trit { @@ -614,180 +861,455 @@ module HSLM { } } - // Small gains keep the attention scores and GELU inputs near 1, where the - // softmax and GELU are far from saturated - fn init_model(gain: f64) -> void { - var l : usize = 0; - while (l < NUM_LAYERS) { - const w = &model.layers[l]; - const salt = l * 1000003; + // A small signed pattern in [-scale/2, scale/2) + fn small(seed: usize, scale: f64) -> f64 { + return scale * (@as(f64, @floatFromInt((seed * 7) % 13)) / 13.0 - 0.5); + } + + // Alphas below 1 (a TWN alpha is a mean |shadow weight|) keep the + // attention scores and the ReLU inputs near 1; none is 1, so a dropped + // alpha shows in the gradient checks + fn init_model() -> void { + var b : usize = 0; + while (b < NUM_BLOCKS) { + const w = &model.blocks[b]; + const salt = b * 1000003; init_matrix(&w.w_q, salt + 11); init_matrix(&w.w_k, salt + 23); init_matrix(&w.w_v, salt + 37); init_matrix(&w.w_o, salt + 41); - init_matrix(&w.w1, salt + 53); - init_matrix(&w.w2, salt + 67); + init_matrix(&w.w_up, salt + 53); + init_matrix(&w.w_down, salt + 67); var i : usize = 0; while (i < EMBED_DIM) { - w.norm1_gamma[i] = gain * (1.0 + 0.01 * @as(f64, @floatFromInt(i % 7))); - w.norm2_gamma[i] = gain * (1.0 + 0.01 * @as(f64, @floatFromInt(i % 5))); + w.norm1_gamma[i] = 0.5 * (1.0 + 0.01 * @as(f64, @floatFromInt(i % 7))); + w.b_down[i] = small(i + b * 3 + 1, 0.04); i = i + 1; } - l = l + 1; + i = 0; + while (i < FF_DIM) { + w.b_up[i] = small(i + b * 5 + 2, 0.1); + i = i + 1; + } + w.alpha_q = 0.1; + w.alpha_k = 0.11; + w.alpha_v = 0.09; + w.alpha_o = 0.8; + w.alpha_up = 0.1; + w.alpha_down = 0.15; + b = b + 1; + } + init_matrix(&model.w_out, 3 * 1000003 + 79); + var e : usize = 0; + while (e < VOCAB_SIZE * EMBED_DIM) { + model.embed[e] = small(e, 0.12); + e = e + 1; } + e = 0; + while (e < VOCAB_SIZE) { + model.b_out[e] = small(e + 4, 0.02); + e = e + 1; + } + model.alpha_out = 0.5; + nudge_which = MAT_NONE; + nudge_delta = 0.0; } - fn set_trit(which: u8, layer_idx: usize, index: usize, t: Trit) -> void { - const w = &model.layers[layer_idx]; - if (which == MAT_Q) { - w.w_q[index] = t; - } else if (which == MAT_K) { - w.w_k[index] = t; - } else if (which == MAT_V) { - w.w_v[index] = t; - } else if (which == MAT_O) { - w.w_o[index] = t; - } else if (which == MAT_W1) { - w.w1[index] = t; + // The linear loss's upstream gradient + fn upstream(j: usize) -> f64 { + return 0.01 * @as(f64, @floatFromInt((j * 5) % 11)) - 0.05; + } + + fn probe_run(position: usize, token: usize) -> void { + hslm_forward(token, position, &probe_logits); + } + + fn probe_loss(target: usize) -> f64 { + probe_run(1, TOKEN_PROBE); + if (target < VOCAB_SIZE) { + return cross_entropy(&probe_logits, target); + } + var loss : f64 = 0.0; + var j : usize = 0; + while (j < VOCAB_SIZE) { + loss = loss + probe_logits[j] * upstream(j); + j = j + 1; + } + return loss; + } + + // The probe model with position 0 in the cache; position 1 is probed + fn probe_setup() -> void { + init_model(); + probe_run(0, TOKEN_FIRST); + } + + // Forward position 1 at the unperturbed point, then backward from the loss + fn probe_backward(target: usize) -> void { + probe_run(1, TOKEN_PROBE); + zero_weight_gradients(); + if (target < VOCAB_SIZE) { + cross_entropy_backward(&probe_logits, target, &probe_dl); } else { - w.w2[index] = t; + var j : usize = 0; + while (j < VOCAB_SIZE) { + probe_dl[j] = upstream(j); + j = j + 1; + } } + hslm_backward(&probe_dl, &probe_gin); } - fn weight_gradient(which: u8, layer_idx: usize, index: usize) -> f64 { + // Embedding entry k of the probed token, moved by +-1e-5 + fn embed_slope(k: usize, target: usize) -> f64 { + probe_setup(); + const e0 = model.embed[TOKEN_PROBE * EMBED_DIM + k]; + model.embed[TOKEN_PROBE * EMBED_DIM + k] = e0 + 0.00001; + const lp = probe_loss(target); + model.embed[TOKEN_PROBE * EMBED_DIM + k] = e0 - 0.00001; + const lm = probe_loss(target); + model.embed[TOKEN_PROBE * EMBED_DIM + k] = e0; + return (lp - lm) / 0.00002; + } + + fn embed_gradient(k: usize, target: usize) -> f64 { + probe_setup(); + probe_backward(target); + return g_embed[TOKEN_PROBE * EMBED_DIM + k]; + } + + // The input gradient hslm_backward returns is the embedding row's + fn input_gradient(k: usize) -> f64 { + probe_setup(); + probe_backward(LINEAR_LOSS); + return probe_gin[k]; + } + + // One ternary weight nudged by +-1e-5 as a real number, position 1 only + fn nudge_slope(which: u8, block_idx: usize, index: usize) -> f64 { + probe_setup(); + nudge_which = which; + nudge_block = block_idx; + nudge_index = index; + nudge_delta = 0.00001; + const lp = probe_loss(LINEAR_LOSS); + nudge_delta = -0.00001; + const lm = probe_loss(LINEAR_LOSS); + nudge_which = MAT_NONE; + nudge_delta = 0.0; + return (lp - lm) / 0.00002; + } + + fn weight_gradient(which: u8, block_idx: usize, index: usize) -> f64 { if (which == MAT_Q) { - return gw_q[layer_idx * EMBED_DIM * EMBED_DIM + index]; + return gw_q[block_idx * EMBED_DIM * EMBED_DIM + index]; } if (which == MAT_K) { - return gw_k[layer_idx * EMBED_DIM * EMBED_DIM + index]; + return gw_k[block_idx * EMBED_DIM * EMBED_DIM + index]; } if (which == MAT_V) { - return gw_v[layer_idx * EMBED_DIM * EMBED_DIM + index]; + return gw_v[block_idx * EMBED_DIM * EMBED_DIM + index]; } if (which == MAT_O) { - return gw_o[layer_idx * EMBED_DIM * EMBED_DIM + index]; + return gw_o[block_idx * EMBED_DIM * EMBED_DIM + index]; + } + if (which == MAT_UP) { + return gw_up[block_idx * EMBED_DIM * FF_DIM + index]; + } + if (which == MAT_DOWN) { + return gw_down[block_idx * FF_DIM * EMBED_DIM + index]; } - if (which == MAT_W1) { - return gw_1[layer_idx * EMBED_DIM * FF_DIM + index]; + return gw_out[index]; + } + + fn weight_grad_at(which: u8, block_idx: usize, index: usize) -> f64 { + probe_setup(); + probe_backward(LINEAR_LOSS); + return weight_gradient(which, block_idx, index); + } + + fn set_trit(which: u8, block_idx: usize, index: usize, t: Trit) -> void { + if (which == MAT_OUT) { + model.w_out[index] = t; + } else { + model.blocks[block_idx].w_down[index] = t; + } + } + + // Flip one trit (W_out or W_down) to +1 and -1, position 0 cached with it + // at 0: (loss(+1) - loss(-1)) / 2 + fn flip_slope(which: u8, block_idx: usize, index: usize) -> f64 { + init_model(); + set_trit(which, block_idx, index, Trit.zero); + probe_run(0, TOKEN_FIRST); + set_trit(which, block_idx, index, Trit.pos); + const lp = probe_loss(LINEAR_LOSS); + set_trit(which, block_idx, index, Trit.neg); + const lm = probe_loss(LINEAR_LOSS); + return (lp - lm) / 2.0; + } + + fn flip_grad_at(which: u8, block_idx: usize, index: usize) -> f64 { + init_model(); + set_trit(which, block_idx, index, Trit.zero); + probe_run(0, TOKEN_FIRST); + probe_backward(LINEAR_LOSS); + return weight_gradient(which, block_idx, index); + } + + fn gamma_slope(block_idx: usize, index: usize) -> f64 { + probe_setup(); + const w = &model.blocks[block_idx]; + const g0 = w.norm1_gamma[index]; + w.norm1_gamma[index] = g0 + 0.00001; + const lp = probe_loss(LINEAR_LOSS); + w.norm1_gamma[index] = g0 - 0.00001; + const lm = probe_loss(LINEAR_LOSS); + w.norm1_gamma[index] = g0; + return (lp - lm) / 0.00002; + } + + fn gamma_gradient(block_idx: usize, index: usize) -> f64 { + probe_setup(); + probe_backward(LINEAR_LOSS); + return g_gamma1[block_idx * EMBED_DIM + index]; + } + + fn set_bias(kind: u8, block_idx: usize, index: usize, value: f64) -> void { + if (kind == BIAS_UP) { + model.blocks[block_idx].b_up[index] = value; + } else if (kind == BIAS_DOWN) { + model.blocks[block_idx].b_down[index] = value; + } else { + model.b_out[index] = value; } - return gw_2[layer_idx * FF_DIM * EMBED_DIM + index]; } - // The upstream gradient: loss = sum_i out_i * upstream(i) - fn upstream(i: usize) -> f64 { - return 0.01 * @as(f64, @floatFromInt((i * 5) % 11)) - 0.05; + fn get_bias(kind: u8, block_idx: usize, index: usize) -> f64 { + if (kind == BIAS_UP) { + return model.blocks[block_idx].b_up[index]; + } + if (kind == BIAS_DOWN) { + return model.blocks[block_idx].b_down[index]; + } + return model.b_out[index]; + } + + fn bias_slope(kind: u8, block_idx: usize, index: usize) -> f64 { + probe_setup(); + const b0 = get_bias(kind, block_idx, index); + set_bias(kind, block_idx, index, b0 + 0.00001); + const lp = probe_loss(LINEAR_LOSS); + set_bias(kind, block_idx, index, b0 - 0.00001); + const lm = probe_loss(LINEAR_LOSS); + set_bias(kind, block_idx, index, b0); + return (lp - lm) / 0.00002; + } + + fn bias_gradient(kind: u8, block_idx: usize, index: usize) -> f64 { + probe_setup(); + probe_backward(LINEAR_LOSS); + if (kind == BIAS_UP) { + return g_b_up[block_idx * FF_DIM + index]; + } + if (kind == BIAS_DOWN) { + return g_b_down[block_idx * EMBED_DIM + index]; + } + return g_b_out[index]; } - // Run one position with input element k moved by delta - fn probe_run(position: usize, k: usize, delta: f64) -> void { + // The first hidden unit of a block whose pre-activation at the probed + // position is above `margin` (active) or below -margin (dead); FF_DIM if none + fn first_unit(block_idx: usize, active: bool, margin: f64) -> usize { + probe_setup(); + probe_run(1, TOKEN_PROBE); + var j : usize = 0; + while (j < FF_DIM) { + const pre = act_pre[block_idx * FF_DIM + j]; + if (active and pre > margin) { + return j; + } + if (active == false and pre < 0.0 - margin) { + return j; + } + j = j + 1; + } + return FF_DIM; + } + + // W_up entry (row, an active unit): its gradient is not trivially zero + fn active_up_index(block_idx: usize, row: usize) -> usize { + return row * FF_DIM + first_unit(block_idx, true, 0.1); + } + + // relu_backward on block 2 with pre = x and an upstream gradient of 3 + fn relu_backward_at(x: f64) -> f64 { + act_pre[2 * FF_DIM + 7] = x; + bw_h[7] = 3.0; + relu_backward(2); + return bw_pre[7]; + } + + // The ReLU kills a dead unit's output: count units of block 0 with h == 0 + // where pre <= 0, and with h == pre where pre > 0 + fn relu_mismatches() -> usize { + probe_setup(); + probe_run(1, TOKEN_PROBE); + var bad : usize = 0; + var j : usize = 0; + while (j < FF_DIM) { + const pre = act_pre[j]; + const h = act_h[j]; + if (pre > 0.0 and h != pre) { + bad = bad + 1; + } + if (pre <= 0.0 and h != 0.0) { + bad = bad + 1; + } + j = j + 1; + } + return bad; + } + + fn dead_units(block_idx: usize) -> usize { + probe_setup(); + probe_run(1, TOKEN_PROBE); + var dead : usize = 0; + var j : usize = 0; + while (j < FF_DIM) { + if (act_h[block_idx * FF_DIM + j] == 0.0) { + dead = dead + 1; + } + j = j + 1; + } + return dead; + } + + // embed_into minus the table row: the position encoding alone + fn embed_probe(token: usize, position: usize, i: usize) -> f64 { + init_model(); + embed_into(token, position, &probe_vec); + return probe_vec[i] - model.embed[token * EMBED_DIM + i]; + } + + // |rope(x)|^2 - |x|^2 at a position, and rope then its inverse, for a + // fixed x + fn rope_norm_change(position: usize) -> f64 { + var before : f64 = 0.0; var i : usize = 0; while (i < EMBED_DIM) { - probe_x[i] = 0.1 * @as(f64, @floatFromInt((i * 3 + position * 7) % 13)) - 0.6; + probe_vec[i] = small(i * 3 + 1, 2.0); + before = before + probe_vec[i] * probe_vec[i]; + i = i + 1; + } + rope(&probe_vec, position, false); + var after : f64 = 0.0; + i = 0; + while (i < EMBED_DIM) { + after = after + probe_vec[i] * probe_vec[i]; i = i + 1; } - probe_x[k] = probe_x[k] + delta; - hslm_forward(&probe_x, position, &probe_out); + return after - before; } - // The same, returning the loss - fn probe_loss(position: usize, k: usize, delta: f64) -> f64 { - probe_run(position, k, delta); - var loss : f64 = 0.0; + fn rope_round_trip_gap(position: usize) -> f64 { var i : usize = 0; while (i < EMBED_DIM) { - loss = loss + probe_out[i] * upstream(i); + probe_vec[i] = small(i * 3 + 1, 2.0); i = i + 1; } - return loss; + rope(&probe_vec, position, false); + rope(&probe_vec, position, true); + var gap : f64 = 0.0; + i = 0; + while (i < EMBED_DIM) { + gap = gap + @abs(probe_vec[i] - small(i * 3 + 1, 2.0)); + i = i + 1; + } + return gap; } - // Model at gain 0.05 with position 0 in the cache; position 1 is probed - fn probe_setup() -> void { - init_model(0.05); - probe_run(0, 0, 0.0); + // Element i of rope(e_k) for a unit vector e_k + fn rope_unit(position: usize, k: usize, i: usize) -> f64 { + fill(&probe_vec, 0.0); + probe_vec[k] = 1.0; + rope(&probe_vec, position, false); + return probe_vec[i]; } - // Forward position 1 at the unperturbed point, then backward from upstream - fn probe_backward() -> void { - probe_run(1, 0, 0.0); - zero_weight_gradients(); + // Head logit j with W_out set to zero: alpha_out * 0 * scale + b_out[j] + fn head_bias_only(j: usize) -> f64 { + init_model(); var i : usize = 0; - while (i < EMBED_DIM) { - probe_up[i] = upstream(i); + while (i < EMBED_DIM * VOCAB_SIZE) { + model.w_out[i] = Trit.zero; i = i + 1; } - hslm_backward(&probe_up, &probe_gin); + probe_run(0, TOKEN_FIRST); + return probe_logits[j] - model.b_out[j]; } - fn input_slope(k: usize) -> f64 { - probe_setup(); - return (probe_loss(1, k, 0.0001) - probe_loss(1, k, -0.0001)) / 0.0002; + // Logit j minus b_out[j], divided by alpha_out * SACRED_LOGIT_SCALE: the + // raw n @ W_out, recomputed by hand for column j + fn head_scale_gap(j: usize) -> f64 { + init_model(); + probe_run(0, TOKEN_FIRST); + var z : f64 = 0.0; + var i : usize = 0; + while (i < EMBED_DIM) { + const w = model.w_out[i * VOCAB_SIZE + j]; + if (w == Trit.pos) { + z = z + act_hn[i]; + } else if (w == Trit.neg) { + z = z - act_hn[i]; + } + i = i + 1; + } + return probe_logits[j] - (0.5 * z * SACRED_LOGIT_SCALE + model.b_out[j]); } - fn input_gradient(k: usize) -> f64 { - probe_setup(); - probe_backward(); - return probe_gin[k]; + // RMS of the head's normalized vector: 1 up to eps + fn head_norm_rms() -> f64 { + init_model(); + probe_run(0, TOKEN_FIRST); + return rms_of(&act_hn); } - // Flip one trit to +1 and -1 (position 0 cached with it at 0): - // (loss(+1) - loss(-1)) / 2 - fn weight_slope(which: u8, layer_idx: usize, index: usize) -> f64 { - init_model(0.05); - set_trit(which, layer_idx, index, Trit.zero); - probe_run(0, 0, 0.0); - set_trit(which, layer_idx, index, Trit.pos); - const lp = probe_loss(1, 0, 0.0); - set_trit(which, layer_idx, index, Trit.neg); - const lm = probe_loss(1, 0, 0.0); - return (lp - lm) / 2.0; + fn ce_uniform() -> f64 { + fill(&probe_logits, 0.3); + return cross_entropy(&probe_logits, 17); } - fn weight_grad_at(which: u8, layer_idx: usize, index: usize) -> f64 { - init_model(0.05); - set_trit(which, layer_idx, index, Trit.zero); - probe_run(0, 0, 0.0); - probe_backward(); - return weight_gradient(which, layer_idx, index); + fn ce_pattern(j: usize, delta: f64) -> f64 { + var i : usize = 0; + while (i < VOCAB_SIZE) { + probe_logits[i] = 0.1 * @as(f64, @floatFromInt((i * 7) % 17)) - 0.8; + i = i + 1; + } + probe_logits[j] = probe_logits[j] + delta; + return cross_entropy(&probe_logits, PROBE_TARGET); } - fn gamma_slope(second: bool, layer_idx: usize, index: usize) -> f64 { - probe_setup(); - const w = &model.layers[layer_idx]; - var g0 : f64 = w.norm1_gamma[index]; - if (second) { - g0 = w.norm2_gamma[index]; - } - if (second) { - w.norm2_gamma[index] = g0 + 0.00001; - } else { - w.norm1_gamma[index] = g0 + 0.00001; - } - const lp = probe_loss(1, 0, 0.0); - if (second) { - w.norm2_gamma[index] = g0 - 0.00001; - } else { - w.norm1_gamma[index] = g0 - 0.00001; - } - const lm = probe_loss(1, 0, 0.0); - return (lp - lm) / 0.00002; + fn ce_slope(j: usize) -> f64 { + return (ce_pattern(j, 0.00001) - ce_pattern(j, -0.00001)) / 0.00002; } - fn gamma_gradient(second: bool, layer_idx: usize, index: usize) -> f64 { - probe_setup(); - probe_backward(); - if (second) { - return g_gamma2[layer_idx * EMBED_DIM + index]; - } - return g_gamma1[layer_idx * EMBED_DIM + index]; + fn ce_gradient(j: usize) -> f64 { + const unused = ce_pattern(j, 0.0); + cross_entropy_backward(&probe_logits, PROBE_TARGET, &probe_dl); + return probe_dl[j] + 0.0 * unused; } - // gelu_backward on layer 2 with pre = x and an upstream gradient of 3 - fn gelu_backward_at(x: f64) -> f64 { - act_pre[2 * FF_DIM + 7] = x; - bw_h[7] = 3.0; - gelu_backward(2); - return bw_pre[7]; + fn ce_gradient_sum() -> f64 { + const unused = ce_pattern(0, 0.0); + cross_entropy_backward(&probe_logits, PROBE_TARGET, &probe_dl); + var sum : f64 = 0.0; + var j : usize = 0; + while (j < VOCAB_SIZE) { + sum = sum + probe_dl[j]; + j = j + 1; + } + return sum + 0.0 * unused; } fn rms_norm_probe(index: usize) -> f64 { @@ -828,40 +1350,40 @@ module HSLM { return output[index]; } - // Sum of the softmax weights of layer 0, head 1, after position 1 + // Sum of the softmax weights of block 0, head 1, after position 1 fn attention_weight_sum() -> f64 { probe_setup(); - probe_run(1, 0, 0.0); + probe_run(1, TOKEN_PROBE); const ao = (0 * NUM_HEADS + 1) * CONTEXT_LEN; return act_attn[ao] + act_attn[ao + 1]; } fn attention_weight(j: usize) -> f64 { probe_setup(); - probe_run(1, 0, 0.0); + probe_run(1, TOKEN_PROBE); return act_attn[(0 * NUM_HEADS + 1) * CONTEXT_LEN + j]; } // At position 0 a head can only attend to itself: its context is v_0 fn self_attention_gap() -> f64 { - init_model(0.05); - probe_run(0, 0, 0.0); + init_model(); + probe_run(0, TOKEN_FIRST); var gap : f64 = 0.0; var i : usize = 0; while (i < EMBED_DIM) { - gap = gap + @abs(act_ctx[3 * EMBED_DIM + i] - kv_v[(3 * CONTEXT_LEN + 0) * EMBED_DIM + i]); + gap = gap + @abs(act_ctx[2 * EMBED_DIM + i] - kv_v[(2 * CONTEXT_LEN + 0) * EMBED_DIM + i]); i = i + 1; } return gap; } - // The output of position 1 depends on what position 0 left in the cache + // The output of position 1 depends on which token position 0 left in the cache fn context_effect() -> f64 { - init_model(0.05); - probe_run(0, 0, 0.0); - const with_first = probe_loss(1, 0, 0.0); - probe_run(0, 0, 0.5); - const with_other_first = probe_loss(1, 0, 0.0); + init_model(); + probe_run(0, TOKEN_FIRST); + const with_first = probe_loss(LINEAR_LOSS); + probe_run(0, 600); + const with_other_first = probe_loss(LINEAR_LOSS); return with_first - with_other_first; } @@ -909,31 +1431,96 @@ module HSLM { when c1 = matmul_rect_probe(1) then c0 == -2.0 and c1 == -2.0 - test hslm_gelu_values - given at_one = gelu(1.0) - when at_minus_one = gelu(-1.0) - then @abs(at_one - 0.8411919906) < 0.0000001 and @abs(at_minus_one + 0.1588080094) < 0.0000001 and gelu(0.0) == 0.0 - - // #7778: the derivative against a central difference at -2.0, 0.0 and 1.5 - test hslm_gelu_derivative_matches_finite_difference - given at_minus_two = (gelu(-2.0 + 0.00001) - gelu(-2.0 - 0.00001)) / 0.00002 - when at_zero = (gelu(0.00001) - gelu(-0.00001)) / 0.00002 - then @abs(gelu_derivative(-2.0) - at_minus_two) < 0.000001 and @abs(gelu_derivative(0.0) - at_zero) < 0.000001 and @abs(gelu_derivative(1.5) - (gelu(1.5 + 0.00001) - gelu(1.5 - 0.00001)) / 0.00002) < 0.000001 - - test hslm_gelu_derivative_at_zero_is_one_half - given d = gelu_derivative(0.0) - then d == 0.5 - - test hslm_gelu_backward_scales_upstream_by_slope - given y = gelu_backward_at(1.5) - when slope = 3.0 * (gelu(1.5 + 0.00001) - gelu(1.5 - 0.00001)) / 0.00002 - then @abs(y - slope) < 0.000001 - - // gelu_backward multiplies the upstream gradient by gelu'(pre) - test hslm_gelu_backward_uses_gelu_derivative - given at_minus_two = gelu_backward_at(-2.0) - when at_zero = gelu_backward_at(0.0) - then @abs(at_minus_two - 3.0 * gelu_derivative(-2.0)) < 0.000000001 and at_zero == 1.5 + // ── Embedding: table row + phi position encoding ── + + // Position 0: sin(0) on even dims, cos(0) on odd dims + test hslm_position_encoding_at_zero + given even = embed_probe(TOKEN_PROBE, 0, 100) + when odd = embed_probe(TOKEN_PROBE, 0, 101) + then even == 0.0 and @abs(odd - 1.0) < 0.000000000000001 + + // Position 1, dims 0 and 1: angle 3/pi; position 5, dim 100: + // sin(5 phi^(-100/243) 3/pi); position 26, dims 241 and 242 share t = 240/243 + test hslm_position_encoding_values + given p1 = embed_probe(TOKEN_PROBE, 1, 0) + when p26 = embed_probe(TOKEN_PROBE, 26, 242) + then @abs(p1 - 0.8162731085894214) < 0.000000000001 and @abs(embed_probe(7, 1, 1) - 0.5776661771246112) < 0.000000000001 and @abs(embed_probe(7, 5, 100) + 0.6999082937334531) < 0.000000000001 and @abs(p26 - 0.32677817087282596) < 0.000000000001 and @abs(embed_probe(7, 26, 241) + 0.9632684200661288) < 0.000000000001 + + // The position encoding does not depend on the token + test hslm_position_encoding_is_token_independent + given a = embed_probe(0, 9, 33) + when b = embed_probe(728, 9, 33) + then @abs(a - b) < 0.000000000000001 and @abs(a) > 0.01 + + // ── phi-RoPE ── + + test hslm_rope_is_identity_at_position_zero + given same = rope_unit(0, 10, 10) + when other = rope_unit(0, 10, 50) + then same == 1.0 and other == 0.0 + + // Pair (10, 50) at position 3: angle 3 phi^(-20/81) + test hslm_rope_rotates_pair_i_with_i_plus_40 + given c = rope_unit(3, 10, 10) + when s = rope_unit(3, 10, 50) + then @abs(c + 0.8880621352130498) < 0.000000000001 and @abs(s - 0.4597234429533031) < 0.000000000001 and rope_unit(3, 80, 80) == 1.0 and rope_unit(3, 161, 161) == 1.0 + + test hslm_rope_preserves_norm_and_inverts + given change = rope_norm_change(19) + when gap = rope_round_trip_gap(19) + then @abs(change) < 0.000000001 and gap < 0.000000001 + + // ── TernaryDense: ReLU, biases ── + + test hslm_relu_values + given pos = relu(1.5) + when neg = relu(-2.0) + then pos == 1.5 and neg == 0.0 and relu(0.0) == 0.0 + + test hslm_relu_backward_masks_dead_units + given live = relu_backward_at(1.5) + when dead = relu_backward_at(-0.5) + then live == 3.0 and dead == 0.0 and relu_backward_at(0.0) == 0.0 + + // The forward pass applies the ReLU to every unit, and the probe model + // has both live and dead units in every block + test hslm_ffn_relu_applied + given bad = relu_mismatches() + when dead = dead_units(0) + then bad == 0 and dead > 100 and dead < 629 and dead_units(2) > 100 and dead_units(2) < 629 + + // ── Head: plain RMSNorm, W_out, alpha, 1 / 243^(phi^-3), bias ── + + test hslm_head_rmsnorm_has_unit_rms + given rms = head_norm_rms() + then @abs(rms - 1.0) < 0.000001 + + test hslm_head_logit_is_scaled_matmul_plus_bias + given gap = head_scale_gap(100) + when gap_last = head_scale_gap(728) + then @abs(gap) < 0.000000000001 and @abs(gap_last) < 0.000000000001 + + test hslm_head_with_zero_weights_is_the_bias + given rest = head_bias_only(250) + then rest == 0.0 + + // ── Cross-entropy ── + + // Equal logits: loss = ln 729 + test hslm_cross_entropy_uniform_is_ln_vocab + given loss = ce_uniform() + then @abs(loss - 6.591673732008658) < 0.000000000001 + + test hslm_cross_entropy_gradient_matches_finite_difference + given at_target = ce_gradient(PROBE_TARGET) + when other = ce_gradient(200) + then @abs(at_target - ce_slope(PROBE_TARGET)) < 0.00000001 and @abs(other - ce_slope(200)) < 0.00000001 and at_target < -0.9 and other > 0.0 + + test hslm_cross_entropy_gradient_sums_to_zero + given total = ce_gradient_sum() + then @abs(total) < 0.000000000001 + + // ── Attention ── test hslm_attention_softmax_sums_to_one given total = attention_weight_sum() @@ -948,65 +1535,118 @@ module HSLM { given effect = context_effect() then @abs(effect) > 0.000001 - // Through all six layers: d loss / d input against a central difference - test hslm_backward_input_gradient_matches_finite_difference - given slope = input_slope(5) - when grad = input_gradient(5) + // ── Backward pass against central differences ── + + // Through the head and all three blocks: d loss / d embedding entry + test hslm_backward_embedding_gradient_matches_finite_difference + given slope = embed_slope(5, LINEAR_LOSS) + when grad = embed_gradient(5, LINEAR_LOSS) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 + + test hslm_backward_embedding_gradient_second_element + given slope = embed_slope(200, LINEAR_LOSS) + when grad = embed_gradient(200, LINEAR_LOSS) then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 - test hslm_backward_input_gradient_second_element - given slope = input_slope(200) - when grad = input_gradient(200) + // The trainer's loss: cross-entropy against the next token + test hslm_backward_cross_entropy_matches_finite_difference + given slope = embed_slope(77, PROBE_TARGET) + when grad = embed_gradient(77, PROBE_TARGET) then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 - // The loss is linear in the last layer's W2, so the trit flip is exact. - // Measured on the lab (2026-10-18): input and gain gradients agree with the - // central difference to 2e-11, W2 to 1e-16, the other trit flips to 7e-4 - // relative (the flip is a step of 1, so curvature shows). - test hslm_ffn_backward_w2_matches_trit_flip - given slope = weight_slope(MAT_W2, NUM_LAYERS - 1, 100 * EMBED_DIM + 7) - when grad = weight_grad_at(MAT_W2, NUM_LAYERS - 1, 100 * EMBED_DIM + 7) + // The input gradient hslm_backward returns is the one the embedding row gets + test hslm_backward_input_gradient_is_embedding_row + given gin = input_gradient(5) + when grad = embed_gradient(5, LINEAR_LOSS) + then gin == grad and @abs(gin) > 0.000001 + + // The loss is linear in W_out, so the trit flip is exact + test hslm_head_w_out_matches_trit_flip + given slope = flip_slope(MAT_OUT, 0, 100 * VOCAB_SIZE + 7) + when grad = flip_grad_at(MAT_OUT, 0, 100 * VOCAB_SIZE + 7) then @abs(grad - slope) < 0.000000000001 and @abs(slope) > 0.000001 - test hslm_ffn_backward_w1_matches_trit_flip - given slope = weight_slope(MAT_W1, NUM_LAYERS - 1, 9 * FF_DIM + 40) - when grad = weight_grad_at(MAT_W1, NUM_LAYERS - 1, 9 * FF_DIM + 40) - then @abs(grad - slope) < 0.002 * @abs(slope) and @abs(slope) > 0.0000005 + test hslm_head_w_out_matches_nudge + given slope = nudge_slope(MAT_OUT, 0, 31 * VOCAB_SIZE + 500) + when grad = weight_grad_at(MAT_OUT, 0, 31 * VOCAB_SIZE + 500) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 - test hslm_attention_backward_wo_matches_trit_flip - given slope = weight_slope(MAT_O, NUM_LAYERS - 1, 12 * EMBED_DIM + 30) - when grad = weight_grad_at(MAT_O, NUM_LAYERS - 1, 12 * EMBED_DIM + 30) - then @abs(grad - slope) < 0.002 * @abs(slope) and @abs(slope) > 0.0000005 + // d loss / d b_out[j] is the upstream gradient itself + test hslm_head_bias_gradient_is_upstream + given slope = bias_slope(BIAS_OUT, 0, 300) + when grad = bias_gradient(BIAS_OUT, 0, 300) + then grad == upstream(300) and @abs(grad - slope) < 0.000000001 and @abs(slope) > 0.000001 - test hslm_attention_backward_wq_matches_trit_flip - given slope = weight_slope(MAT_Q, NUM_LAYERS - 1, 20 * EMBED_DIM + 90) - when grad = weight_grad_at(MAT_Q, NUM_LAYERS - 1, 20 * EMBED_DIM + 90) + test hslm_ffn_backward_w_down_matches_trit_flip + given slope = flip_slope(MAT_DOWN, NUM_BLOCKS - 1, first_unit(NUM_BLOCKS - 1, true, 0.1) * EMBED_DIM + 7) + when grad = flip_grad_at(MAT_DOWN, NUM_BLOCKS - 1, first_unit(NUM_BLOCKS - 1, true, 0.1) * EMBED_DIM + 7) then @abs(grad - slope) < 0.002 * @abs(slope) and @abs(slope) > 0.0000005 - test hslm_attention_backward_wk_matches_trit_flip - given slope = weight_slope(MAT_K, NUM_LAYERS - 1, 20 * EMBED_DIM + 90) - when grad = weight_grad_at(MAT_K, NUM_LAYERS - 1, 20 * EMBED_DIM + 90) - then @abs(grad - slope) < 0.002 * @abs(slope) and @abs(slope) > 0.0000005 + test hslm_ffn_backward_w_down_matches_nudge + given slope = nudge_slope(MAT_DOWN, NUM_BLOCKS - 1, first_unit(NUM_BLOCKS - 1, true, 0.1) * EMBED_DIM + 30) + when grad = weight_grad_at(MAT_DOWN, NUM_BLOCKS - 1, first_unit(NUM_BLOCKS - 1, true, 0.1) * EMBED_DIM + 30) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 - test hslm_attention_backward_wv_matches_trit_flip - given slope = weight_slope(MAT_V, NUM_LAYERS - 1, 20 * EMBED_DIM + 170) - when grad = weight_grad_at(MAT_V, NUM_LAYERS - 1, 20 * EMBED_DIM + 170) - then @abs(grad - slope) < 0.002 * @abs(slope) and @abs(slope) > 0.0000005 + test hslm_ffn_backward_w_up_matches_nudge + given slope = nudge_slope(MAT_UP, NUM_BLOCKS - 1, active_up_index(NUM_BLOCKS - 1, 20)) + when grad = weight_grad_at(MAT_UP, NUM_BLOCKS - 1, active_up_index(NUM_BLOCKS - 1, 20)) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 - // A first-layer weight: the gradient crosses five more layers - test hslm_backward_first_layer_w1_matches_trit_flip - given slope = weight_slope(MAT_W1, 0, 3 * FF_DIM + 500) - when grad = weight_grad_at(MAT_W1, 0, 3 * FF_DIM + 500) - then @abs(grad - slope) < 0.002 * @abs(slope) and @abs(slope) > 0.0000005 + // A first-block weight: the gradient crosses two more blocks and the head + test hslm_backward_first_block_w_up_matches_nudge + given slope = nudge_slope(MAT_UP, 0, active_up_index(0, 3)) + when grad = weight_grad_at(MAT_UP, 0, active_up_index(0, 3)) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 + + test hslm_attention_backward_wo_matches_nudge + given slope = nudge_slope(MAT_O, NUM_BLOCKS - 1, 12 * EMBED_DIM + 30) + when grad = weight_grad_at(MAT_O, NUM_BLOCKS - 1, 12 * EMBED_DIM + 30) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 + + test hslm_attention_backward_wq_matches_nudge + given slope = nudge_slope(MAT_Q, NUM_BLOCKS - 1, 20 * EMBED_DIM + 90) + when grad = weight_grad_at(MAT_Q, NUM_BLOCKS - 1, 20 * EMBED_DIM + 90) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.0000001 + + test hslm_attention_backward_wk_matches_nudge + given slope = nudge_slope(MAT_K, NUM_BLOCKS - 1, 20 * EMBED_DIM + 90) + when grad = weight_grad_at(MAT_K, NUM_BLOCKS - 1, 20 * EMBED_DIM + 90) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.0000001 + + test hslm_attention_backward_wv_matches_nudge + given slope = nudge_slope(MAT_V, NUM_BLOCKS - 1, 20 * EMBED_DIM + 170) + when grad = weight_grad_at(MAT_V, NUM_BLOCKS - 1, 20 * EMBED_DIM + 170) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 + + test hslm_attention_backward_first_block_wq_matches_nudge + given slope = nudge_slope(MAT_Q, 0, 40 * EMBED_DIM + 15) + when grad = weight_grad_at(MAT_Q, 0, 40 * EMBED_DIM + 15) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.0000001 + + test hslm_backward_b_up_matches_finite_difference + given slope = bias_slope(BIAS_UP, 0, first_unit(0, true, 0.1)) + when grad = bias_gradient(BIAS_UP, 0, first_unit(0, true, 0.1)) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 + + // A dead unit passes no gradient, and moving its bias changes nothing + test hslm_backward_b_up_dead_unit_is_zero + given slope = bias_slope(BIAS_UP, 1, first_unit(1, false, 0.1)) + when grad = bias_gradient(BIAS_UP, 1, first_unit(1, false, 0.1)) + then slope == 0.0 and grad == 0.0 and first_unit(1, false, 0.1) < FF_DIM + + test hslm_backward_b_down_matches_finite_difference + given slope = bias_slope(BIAS_DOWN, 1, 111) + when grad = bias_gradient(BIAS_DOWN, 1, 111) + then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 test hslm_backward_norm1_gamma_matches_finite_difference - given slope = gamma_slope(false, 2, 17) - when grad = gamma_gradient(false, 2, 17) + given slope = gamma_slope(0, 17) + when grad = gamma_gradient(0, 17) then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 - test hslm_backward_norm2_gamma_matches_finite_difference - given slope = gamma_slope(true, 4, 101) - when grad = gamma_gradient(true, 4, 101) + test hslm_backward_last_block_gamma_matches_finite_difference + given slope = gamma_slope(NUM_BLOCKS - 1, 101) + when grad = gamma_gradient(NUM_BLOCKS - 1, 101) then @abs(grad - slope) < 0.000000001 * (1.0 + @abs(slope)) and @abs(slope) > 0.000001 invariant hslm_embedding_dimensional_consistency @@ -1015,15 +1655,28 @@ module HSLM { invariant hslm_head_dim_is_three_pow_four assert HEAD_DIM == 3 * 3 * 3 * 3 - invariant hslm_ffn_expansion_ratio - assert FF_DIM == 4 * EMBED_DIM + invariant hslm_three_trinity_blocks + assert NUM_BLOCKS == 3 - invariant hslm_context_fits_cache - assert CONTEXT_LEN == HEAD_DIM + invariant hslm_ffn_is_three_pow_six + assert FF_DIM == 3 * EMBED_DIM and VOCAB_SIZE == FF_DIM + + // EXP-012: the context divides the head dimension; EXP-014: it is a power of 3 + invariant hslm_context_divides_head_dim + assert CONTEXT_LEN == 3 * 3 * 3 and HEAD_DIM % CONTEXT_LEN == 0 + + invariant hslm_rope_pairs_leave_one_dim + assert 2 * ROPE_PAIRS + 1 == HEAD_DIM invariant hslm_phase_constants_are_unique assert PHASE_FORWARD != PHASE_BACKWARD and PHASE_BACKWARD != PHASE_UPDATE and PHASE_FORWARD != PHASE_UPDATE invariant hslm_attention_scale_is_sacred assert SACRED_ATTN_SCALE > 0.3543788 and SACRED_ATTN_SCALE < 0.3543789 + + invariant hslm_logit_scale_is_sacred + assert SACRED_LOGIT_SCALE > 0.2734228 and SACRED_LOGIT_SCALE < 0.2734229 + + invariant hslm_linear_loss_marker_is_vocab + assert LINEAR_LOSS == VOCAB_SIZE and PROBE_TARGET < VOCAB_SIZE and TOKEN_PROBE < VOCAB_SIZE }