From 6c200f83b13c50bbc8759d15546634ac4a9904d8 Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Sun, 18 Oct 2026 17:32:42 +0700 Subject: [PATCH 1/5] spec(ml): train_hslm -- main() trains the whole 3-block HSLM with STE and LAMB, writes MSLH checkpoints (Refs #10372) Co-Authored-By: Claude Opus 5.5 --- specs/ml/train_hslm.t27 | 2030 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 2030 insertions(+) create mode 100644 specs/ml/train_hslm.t27 diff --git a/specs/ml/train_hslm.t27 b/specs/ml/train_hslm.t27 new file mode 100644 index 0000000000..d77b8bb328 --- /dev/null +++ b/specs/ml/train_hslm.t27 @@ -0,0 +1,2030 @@ +// SPDX-License-Identifier: Apache-2.0 +; specs/ml/train_hslm.t27 -- main() trains the whole 3-block HSLM of specs/nn/hslm.t27 on real +; tokens: straight-through shadow weights, LAMB, a warmup-cosine schedule, global-norm clipping, +; held-out evaluation, and MSLH checkpoints that are read back and compared byte for byte. +; Issue: gHashTag/t27#10372. +; +; WHAT IS REUSED, NOT COPIED: the forward pass (hslm_forward, hslm_block_forward, the sacred +; attention, the FFN, the VSA stream and its gate), the backward pass (hslm_backward: the exact +; gradient of the last position's cross-entropy, embedding included), cross_entropy, the TWN +; quantizer (nn/ternary_ste quantize_twn) and the checkpoint header and layout (nn/hslm_data +; write_header, parse_header, the section offsets). They are reached through `use`. +; +; THE MODEL: ctx 27, vocab 729, embed 243, ff 729, 3 blocks, 3 heads of 81. Every ternary matrix +; (Q, K, V, O, up, down per block, and the output head: 19 matrices, 1771470 trits) has an f64 +; shadow; after every update the shadow is requantized with TWN (delta = 0.7 mean |w|, alpha = mean +; |w| over the kept entries), one alpha per matrix. The float parameters (embedding table 729 x 243, +; RMSNorm gains, FFN biases, output bias) are trained directly. NPARAM = 2130138 f64 in one flat +; vector, 30 segments. +; +; THE STEP: 66 windows of 27 tokens, starts drawn by a xorshift64 PRNG (pinned seed); per window, +; 27 forward positions, cross-entropy of the last position against the 28th token, hslm_backward. +; Gradients are averaged over the 66 windows, clipped to global L2 norm 1.0, and handed to LAMB +; (beta1 0.9, beta2 0.999, eps 1e-6, decoupled weight decay 0.01 on the matrices and the embedding, +; none on gains and biases, trust ratio ||w|| / ||r|| per segment, clamped to 10, 1 when either norm +; is 0 -- trinity-training autograd.zig lambStepSlice). Learning rate: 1e-3, linear warmup over +; 2000 updates, then cosine down to 1e-6 at TOTAL updates. +; +; THE FILES (all little endian): +; /tmp/hslm/train.tok, valid.tok "HTOK", u32 version 1, u32 N, N u16 ids below 729 (input) +; /tmp/hslm/train.cfg optional, key=value lines of integers: steps, total, warmup, +; eval_every, ckpt_every, seed, speed, resume, synth, eval_windows +; /tmp/hslm/ckpt_.bin MSLH: nn/hslm_data's header (magic 0x484C534D, version 1, step, +; f32 loss = the last minibatch loss) and its layout, f32: output +; shadow (243 x 729, row-major [in][out]), then per block shadow_up +; (243 x 729), shadow_down (729 x 243), bias_up 729, bias_down 243. +; 4971796 bytes, the size of hslm_step_10000.bin. +; /tmp/hslm/ckpt_.state the full resumable state, since MSLH holds no attention matrices, +; embedding, gains or output bias: a 64-byte header ("T27S", version +; 1, step, nparam, rng u64, beta1^t f64, beta2^t f64, loss f64, +; nseg) then the parameters, LAMB m and LAMB v, f64 each. +; Both files are read back after they are written and compared with what was meant to be on disk; +; any difference makes main() fail (assert, a nonzero exit). resume= in train.cfg restarts +; from ckpt_.state and continues bit-identically (same windows, same updates). +; +; EVERY EVAL_EVERY STEPS one line goes to stdout: step, the step's minibatch loss and its exp, the +; mean minibatch loss since the last line, the loss on EVAL_WINDOWS fixed held-out windows (evenly +; spaced starts in valid.tok, last position) and its exp, tokens/s (66 x 27 tokens per step over the +; training time only), s/step, the learning rate and the gradient norm before clipping. +; +; SPEED, HONESTLY: speed=0 runs the reference kernels of nn/hslm.t27 as they are. speed=1 replaces +; four kernels by local ones with the same results: ternary_matmul as a branch-free axpy over the +; rows of W (bit-identical: same summation order per output, x * 0.0 adds +0.0 to a sum that is never +; -0.0), trit_int and outer_add branch-free, and ternary_matmul_t_add with four partial sums (the one +; kernel that is not bit-identical; it only runs in backward). speed=2 (the default) also skips work +; whose result nobody reads: the head at positions 0..25 and, in the last block at those positions, +; everything but the key and value that the last position attends to. The tests check that speed 1 +; and speed 2 produce the same logits as speed 0 to the bit. The model is not reduced. +; +; WHAT DIFFERS FROM THE ZIG TRAINER (trinity-training @ 3e9648c, run R5): f64, not f32; LAMB, not +; AdamW; one global clip, not per tensor; linear warmup + cosine to 1e-6, not the phi schedule; the +; embedding table and RMSNorm gains are trained (the trainer froze the table); the gradient is the +; last position's, with earlier keys and values as constants (nn/hslm.t27 section 5); alphas are +; recomputed by TWN, not learned; the VSA table is nn/hslm.t27's hash, not Zig's PRNG; the MSLH body +; floats are IEEE f32 casts (round to nearest even), the header loss is hslm_data's encode_f32. +; phi^2 + 1/phi^2 = 3 | TRINITY + +module train_hslm; + +use base::types; +use nn::hslm; +use nn::ternary_ste; +use nn::hslm_data; + +; THE SHAPE. nn/hslm and nn/hslm_data both declare these five; a name two imports declare is not +; resolved, so this module declares them once, with the same type and value as both. +const CONTEXT_LEN : usize = 27; +const VOCAB_SIZE : usize = 729; +const EMBED_DIM : usize = 243; +const FF_DIM : usize = 729; +const NUM_BLOCKS : usize = 3; +const HEADER_BYTES : usize = 16; + +; THE STATE OF nn/hslm.t27. use_resolve splices fns, structs, enums and consts but not module-level +; `var`; the spliced forward and backward bind to these, which repeat nn/hslm.t27 section 2 (names, +; types, initial values) and nothing else. A variable added there and missing here stops the build +; with "use of undeclared identifier"; it cannot drift silently. +var model : HSLMWeights = undefined; +var kv_k : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var kv_v : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var act_token : usize = 0; +var act_pos : [NUM_BLOCKS]usize = undefined; +var act_x : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_n1 : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_rms1 : [NUM_BLOCKS]f64 = undefined; +var act_q : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_attn : [NUM_BLOCKS * NUM_HEADS * CONTEXT_LEN]f64 = undefined; +var act_ctx : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_r : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_pre : [NUM_BLOCKS * FF_DIM]f64 = undefined; +var act_h : [NUM_BLOCKS * FF_DIM]f64 = undefined; +var act_out : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_hn : [EMBED_DIM]f64 = undefined; +var act_hrms : f64 = 1.0; +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; +var bw_out : [EMBED_DIM]f64 = undefined; +var bw_h : [FF_DIM]f64 = undefined; +var bw_pre : [FF_DIM]f64 = undefined; +var bw_sf : [FF_DIM]f64 = undefined; +var bw_se : [EMBED_DIM]f64 = undefined; +var bw_n : [EMBED_DIM]f64 = undefined; +var bw_tmp : [EMBED_DIM]f64 = undefined; +var bw_r : [EMBED_DIM]f64 = undefined; +var bw_ctx : [EMBED_DIM]f64 = undefined; +var bw_q : [EMBED_DIM]f64 = undefined; +var bw_k : [EMBED_DIM]f64 = undefined; +var bw_v : [EMBED_DIM]f64 = undefined; +var bw_da : [CONTEXT_LEN]f64 = undefined; +var bw_x : [EMBED_DIM]f64 = undefined; +var bw_dz : [VOCAB_SIZE]f64 = undefined; +var nudge_which : u8 = MAT_NONE; +var nudge_block : usize = 0; +var nudge_index : usize = 0; +var nudge_delta : f64 = 0.0; +var vsa_on : bool = true; +var vsa_in : [NUM_BLOCKS * CONTEXT_LEN * VSA_DIM]Trit = undefined; +var vsa_out : [NUM_BLOCKS * VSA_DIM]Trit = undefined; +var vsa_ctx : [VSA_DIM]Trit = undefined; +var vsa_reasoned : [VSA_DIM]Trit = undefined; +var vsa_temp : [VSA_DIM]Trit = undefined; +var vsa_temp2 : [VSA_DIM]Trit = undefined; +var vsa_sim : [CONTEXT_LEN]f64 = undefined; +var vsa_accum : [VSA_DIM]i64 = undefined; +var vsa_max_sim : [NUM_BLOCKS]f64 = undefined; +var vsa_fired : [NUM_BLOCKS]bool = undefined; +var gate_total : [NUM_BLOCKS]usize = undefined; +var gate_conscious : [NUM_BLOCKS]usize = undefined; +var gate_ema : [NUM_BLOCKS]f64 = undefined; + +; THE TRAINING CONSTANTS (the issue's numbers). Steps, schedule length, warmup, intervals, seed, +; speed level and resume step can be overridden in /tmp/hslm/train.cfg; these are the defaults. +const BATCH : usize = 66; +const LR_PEAK : f64 = 0.001; +const LR_FLOOR : f64 = 0.000001; +const WARMUP_STEPS : u64 = 2000; +const DEFAULT_STEPS : u64 = 32000; +const DEFAULT_EVAL_EVERY : u64 = 500; +const DEFAULT_CKPT_EVERY : u64 = 2000; +const DEFAULT_SEED : u64 = 1729; +const DEFAULT_EVAL_WINDOWS : u64 = 243; +const CLIP_NORM : f64 = 1.0; +const WEIGHT_DECAY : f64 = 0.01; +const LAMB_BETA1 : f64 = 0.9; +const LAMB_BETA2 : f64 = 0.999; +const LAMB_EPS : f64 = 0.000001; +const TRUST_CLAMP : f64 = 10.0; +const TR_PI : f64 = 3.141592653589793; +const INIT_SALT : u64 = 7046029254386353131; + +; Parameter segments: per block Q K V O up down gamma b_up b_down, then out, b_out, embed +const K_Q : u8 = 0; +const K_K : u8 = 1; +const K_V : u8 = 2; +const K_O : u8 = 3; +const K_UP : u8 = 4; +const K_DOWN : u8 = 5; +const K_GAMMA : u8 = 6; +const K_BUP : u8 = 7; +const K_BDOWN : u8 = 8; +const K_OUT : u8 = 9; +const K_BOUT : u8 = 10; +const K_EMBED : u8 = 11; +const KINDS_PER_BLOCK : usize = 9; +const NSEG : usize = 30; +const NPARAM : usize = 2130138; + +; Token buffers, sized for the real corpus (train 8557893, valid 7689354 tokens) +const TRAIN_CAP : usize = 16777216; +const VALID_CAP : usize = 8388608; +const SYN_TRAIN : usize = 262144; +const SYN_VALID : usize = 32768; +const IO_CHUNK : usize = 1048576; +const MSLH_BYTES : usize = 4971796; +const STATE_HEADER : usize = 64; +const STATE_MAGIC : u32 = 0x53373254; +const STATE_VERSION : u32 = 1; +const HTOK_MAGIC : u32 = 0x4B4F5448; +const HTOK_VERSION : u32 = 1; + +; Exit codes of tr_main +const RC_OK : i32 = 0; +const RC_TRAIN_TOK : i32 = 10; +const RC_VALID_TOK : i32 = 11; +const RC_SYNTH_WRITE : i32 = 12; +const RC_CKPT_WRITE : i32 = 20; +const RC_CKPT_VERIFY : i32 = 21; +const RC_STATE_WRITE : i32 = 22; +const RC_STATE_VERIFY : i32 = 23; +const RC_RESUME : i32 = 30; +const RC_NAN : i32 = 40; + +; Trainer state +var p_val : [NPARAM]f64 = undefined; +var g_val : [NPARAM]f64 = undefined; +var m_val : [NPARAM]f64 = undefined; +var v_val : [NPARAM]f64 = undefined; +var seg_off : [NSEG]usize = undefined; +var seg_len : [NSEG]usize = undefined; +var seg_kind : [NSEG]u8 = undefined; +var seg_block : [NSEG]usize = undefined; +var beta1_pow : f64 = 1.0; +var beta2_pow : f64 = 1.0; +var rng_state : u64 = 1; +var speed_level : u8 = 2; +var last_lr : f64 = 0.0; +var last_gnorm : f64 = 0.0; +var last_loss : f64 = 0.0; +var io_fail : u32 = 0; + +var train_tok : [TRAIN_CAP]u16 = undefined; +var valid_tok : [VALID_CAP]u16 = undefined; +var ntrain : usize = 0; +var nvalid : usize = 0; +var data_synthetic : bool = false; +var fixed_starts : [8]usize = undefined; +var win_logits : [VOCAB_SIZE]f64 = undefined; +var win_dl : [VOCAB_SIZE]f64 = undefined; +var win_gin : [EMBED_DIM]f64 = undefined; +var ref_logits : [VOCAB_SIZE]f64 = undefined; +var ref_gin : [EMBED_DIM]f64 = undefined; +var kv_n1 : [EMBED_DIM]f64 = undefined; + +var io_chunk : [IO_CHUNK]u8 = undefined; +var io_back : [IO_CHUNK]u8 = undefined; +var io_hdr : [STATE_HEADER]u8 = undefined; +var io_hdr_back : [STATE_HEADER]u8 = undefined; +var ck_buf : [MSLH_BYTES]u8 = undefined; +var ck_back : [MSLH_BYTES]u8 = undefined; +var path_buf : [96]u8 = undefined; +var up_buf : [64]u8 = undefined; +var cfg_buf : [4096]u8 = undefined; +var ln_buf : [512]u8 = undefined; +var ln_n : usize = 0; +var clip_buf : [1000]f64 = undefined; +var clip_ref : [1000]f64 = undefined; +var st_save : [131072]f64 = undefined; +var tiny_before : f64 = 0.0; + +var cfg_steps : u64 = DEFAULT_STEPS; +var cfg_total : u64 = DEFAULT_STEPS; +var cfg_warmup : u64 = WARMUP_STEPS; +var cfg_eval_every : u64 = DEFAULT_EVAL_EVERY; +var cfg_ckpt_every : u64 = DEFAULT_CKPT_EVERY; +var cfg_seed : u64 = DEFAULT_SEED; +var cfg_speed : u64 = 2; +var cfg_resume : u64 = 0; +var cfg_synth : u64 = 0; +var cfg_eval_windows : u64 = DEFAULT_EVAL_WINDOWS; + +; --------------------------------------------------------------------------------------------- +; 1. Kernels. The four names below are nn/hslm.t27's; a local declaration wins over the spliced +; one, also inside the spliced forward and backward. speed 0 runs nn/hslm.t27's code unchanged. +; --------------------------------------------------------------------------------------------- + +/// output[i] = sum_j input[j] * w[j * out_dim + i]. Speed 0: nn/hslm.t27's branchy dot per output. +/// Otherwise: zero the output, then add x_j * trit row by row (an axpy the compiler vectorizes); +/// per output the terms are added in the same order j = 0, 1, ..., so the result is the same bits. +fn ternary_matmul(input: []const f64, weights: []const Trit, output: []f64, in_dim: usize, out_dim: usize) -> void { + if (speed_level == 0) { + tr_matmul_ref(input, weights, output, in_dim, out_dim); + return; + } + var i : usize = 0; + while (i < out_dim) { + output[i] = 0.0; + i = i + 1; + } + var j : usize = 0; + while (j < in_dim) { + const xj : f64 = input[j]; + if (xj != 0.0) { + const row : usize = j * out_dim; + var c : usize = 0; + while (c < out_dim) { + const wf : f64 = @floatFromInt(@intFromEnum(weights[row + c])); + output[c] = output[c] + xj * wf; + c = c + 1; + } + } + j = j + 1; + } +} + +/// nn/hslm.t27's ternary_matmul, verbatim +fn tr_matmul_ref(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 wt = weights[j * out_dim + i]; + if (wt == Trit.pos) { + acc = acc + input[j]; + } else if (wt == Trit.neg) { + acc = acc - input[j]; + } + j = j + 1; + } + output[i] = acc; + i = i + 1; + } +} + +/// dx[j] += sum_i w[j * out_dim + i] * dy[i]. Speed 0: nn/hslm.t27's branchy sum. Otherwise four +/// partial sums of trit * dy (the only kernel whose rounding differs from the reference). +fn ternary_matmul_t_add(dy: []const f64, weights: []const Trit, dx: []f64, in_dim: usize, out_dim: usize) -> void { + var j : usize = 0; + while (j < in_dim) { + const row : usize = j * out_dim; + if (speed_level == 0) { + var acc : f64 = 0.0; + var i : usize = 0; + while (i < out_dim) { + const wt = weights[row + i]; + if (wt == Trit.pos) { + acc = acc + dy[i]; + } else if (wt == Trit.neg) { + acc = acc - dy[i]; + } + i = i + 1; + } + dx[j] = dx[j] + acc; + } else { + var a0 : f64 = 0.0; + var a1 : f64 = 0.0; + var a2 : f64 = 0.0; + var a3 : f64 = 0.0; + var i : usize = 0; + while (i + 4 <= out_dim) { + const w0 : f64 = @floatFromInt(@intFromEnum(weights[row + i])); + const w1 : f64 = @floatFromInt(@intFromEnum(weights[row + i + 1])); + const w2 : f64 = @floatFromInt(@intFromEnum(weights[row + i + 2])); + const w3 : f64 = @floatFromInt(@intFromEnum(weights[row + i + 3])); + a0 = a0 + w0 * dy[i]; + a1 = a1 + w1 * dy[i + 1]; + a2 = a2 + w2 * dy[i + 2]; + a3 = a3 + w3 * dy[i + 3]; + i = i + 4; + } + while (i < out_dim) { + const wr : f64 = @floatFromInt(@intFromEnum(weights[row + i])); + a0 = a0 + wr * dy[i]; + i = i + 1; + } + dx[j] = dx[j] + ((a0 + a1) + (a2 + a3)); + } + j = j + 1; + } +} + +/// gw[j * out_dim + i] += x[j] * dy[i], with x[j] read once per row; above speed 0 a row whose +/// x[j] is 0 (half the ReLU outputs) is skipped, which adds nothing either way +fn outer_add(x: []const f64, dy: []const f64, gw: []f64, in_dim: usize, out_dim: usize) -> void { + var j : usize = 0; + while (j < in_dim) { + const xj : f64 = x[j]; + if (speed_level == 0 or xj != 0.0) { + const row : usize = j * out_dim; + var i : usize = 0; + while (i < out_dim) { + gw[row + i] = gw[row + i] + xj * dy[i]; + i = i + 1; + } + } + j = j + 1; + } +} + +/// -1 / 0 / +1 of a trit; above speed 0 the enum's own value, without branches +fn trit_int(t: Trit) -> i64 { + if (speed_level == 0) { + if (t == Trit.pos) { + return 1; + } + if (t == Trit.neg) { + return 0 - 1; + } + return 0; + } + return @as(i64, @intFromEnum(t)); +} + +/// The last block at a position before the last: only its key and value are ever read (by the +/// last position's attention). Same calls, in the same order, as attention_forward. +fn tr_kv_only(block_idx: usize, position: usize, input: []const f64) -> void { + const bw = &model.blocks[block_idx]; + const ko : usize = (block_idx * CONTEXT_LEN + position) * EMBED_DIM; + const rms : f64 = rms_norm_into(input, &bw.norm1_gamma, &kv_n1); + const kk = kv_k[ko..ko + EMBED_DIM]; + const vv = kv_v[ko..ko + EMBED_DIM]; + ternary_matmul(&kv_n1, &bw.w_k, kk, EMBED_DIM, EMBED_DIM); + nudge(MAT_K, block_idx, &kv_n1, kk, EMBED_DIM); + scale_into(kk, bw.alpha_k, kk); + ternary_matmul(&kv_n1, &bw.w_v, vv, EMBED_DIM, EMBED_DIM); + nudge(MAT_V, block_idx, &kv_n1, vv, EMBED_DIM); + scale_into(vv, bw.alpha_v, vv); + rope(kk, position, false); + if (rms < 0.0) { + io_fail = io_fail + 1; + } +} + +/// hslm_forward at a position before the last, without the work nobody reads +fn tr_prefix_forward(token: usize, position: usize) -> void { + embed_into(token, position, &bw_tmp); + if (vsa_on) { + vsa_embed(token, position, vsa_in[position * VSA_DIM..(position + 1) * VSA_DIM]); + } + hslm_block_forward(0, position, &bw_tmp, act_out[0..EMBED_DIM]); + var b : usize = 1; + while (b + 1 < NUM_BLOCKS) { + const prev : usize = (b - 1) * EMBED_DIM; + const cur : usize = b * EMBED_DIM; + hslm_block_forward(b, position, act_out[prev..prev + EMBED_DIM], act_out[cur..cur + EMBED_DIM]); + b = b + 1; + } + const last_in : usize = (NUM_BLOCKS - 2) * EMBED_DIM; + tr_kv_only(NUM_BLOCKS - 1, position, act_out[last_in..last_in + EMBED_DIM]); +} + +/// One 27-token window: every position in order (each attends to the cache of the earlier ones); +/// the logits of the last position land in `logits`, its activations stay for hslm_backward +fn tr_window_forward(tokens: []const u16, start: usize, logits: []f64) -> void { + var p : usize = 0; + while (p < CONTEXT_LEN) { + const tok : usize = @intCast(tokens[start + p]); + if (speed_level < 2 or p + 1 == CONTEXT_LEN) { + hslm_forward(tok, p, logits); + } else { + tr_prefix_forward(tok, p); + } + p = p + 1; + } +} + +/// Forward, loss and backward of one window; gradients accumulate in nn/hslm's gw_* and g_* +fn tr_window_grad(tokens: []const u16, start: usize) -> f64 { + tr_window_forward(tokens, start, &win_logits); + const target : usize = @intCast(tokens[start + CONTEXT_LEN]); + const loss : f64 = cross_entropy(&win_logits, target); + cross_entropy_backward(&win_logits, target, &win_dl); + hslm_backward(&win_dl, &win_gin); + return loss; +} + +/// Loss of one window, forward only +fn tr_window_loss(tokens: []const u16, start: usize) -> f64 { + tr_window_forward(tokens, start, &win_logits); + const target : usize = @intCast(tokens[start + CONTEXT_LEN]); + return cross_entropy(&win_logits, target); +} + +; --------------------------------------------------------------------------------------------- +; 2. Parameters: the flat vector, its segments, init, requantize +; --------------------------------------------------------------------------------------------- + +fn tr_kind_len(kind: u8) -> usize { + if (kind <= K_O) { + return EMBED_DIM * EMBED_DIM; + } + if (kind == K_UP or kind == K_DOWN or kind == K_OUT or kind == K_EMBED) { + return EMBED_DIM * FF_DIM; + } + if (kind == K_BUP or kind == K_BOUT) { + return FF_DIM; + } + return EMBED_DIM; +} + +fn tr_kind_ternary(kind: u8) -> bool { + return kind <= K_DOWN or kind == K_OUT; +} + +/// Weight decay on the matrices and the embedding table, not on gains and biases +fn tr_kind_decay(kind: u8) -> bool { + return kind <= K_DOWN or kind == K_OUT or kind == K_EMBED; +} + +fn tr_seg(kind: u8, block_idx: usize) -> usize { + if (kind == K_OUT) { + return NUM_BLOCKS * KINDS_PER_BLOCK; + } + if (kind == K_BOUT) { + return NUM_BLOCKS * KINDS_PER_BLOCK + 1; + } + if (kind == K_EMBED) { + return NUM_BLOCKS * KINDS_PER_BLOCK + 2; + } + return block_idx * KINDS_PER_BLOCK + @as(usize, kind); +} + +fn tr_build_segments() -> usize { + var off : usize = 0; + var s : usize = 0; + var b : usize = 0; + while (b < NUM_BLOCKS) { + var k : u8 = 0; + while (k < 9) { + seg_kind[s] = k; + seg_block[s] = b; + seg_off[s] = off; + seg_len[s] = tr_kind_len(k); + off = off + seg_len[s]; + s = s + 1; + k = k + 1; + } + b = b + 1; + } + var tail : u8 = K_OUT; + while (tail <= K_EMBED) { + seg_kind[s] = tail; + seg_block[s] = 0; + seg_off[s] = off; + seg_len[s] = tr_kind_len(tail); + off = off + seg_len[s]; + s = s + 1; + tail = tail + 1; + } + return off; +} + +/// The gradient nn/hslm's backward accumulated for segment s +fn tr_grad_of(s: usize) -> []f64 { + const kind : u8 = seg_kind[s]; + const b : usize = seg_block[s]; + const ee : usize = EMBED_DIM * EMBED_DIM; + const ef : usize = EMBED_DIM * FF_DIM; + if (kind == K_Q) { + return gw_q[b * ee..(b + 1) * ee]; + } + if (kind == K_K) { + return gw_k[b * ee..(b + 1) * ee]; + } + if (kind == K_V) { + return gw_v[b * ee..(b + 1) * ee]; + } + if (kind == K_O) { + return gw_o[b * ee..(b + 1) * ee]; + } + if (kind == K_UP) { + return gw_up[b * ef..(b + 1) * ef]; + } + if (kind == K_DOWN) { + return gw_down[b * ef..(b + 1) * ef]; + } + if (kind == K_GAMMA) { + return g_gamma1[b * EMBED_DIM..(b + 1) * EMBED_DIM]; + } + if (kind == K_BUP) { + return g_b_up[b * FF_DIM..(b + 1) * FF_DIM]; + } + if (kind == K_BDOWN) { + return g_b_down[b * EMBED_DIM..(b + 1) * EMBED_DIM]; + } + if (kind == K_OUT) { + return gw_out[0..ef]; + } + if (kind == K_BOUT) { + return g_b_out[0..VOCAB_SIZE]; + } + return g_embed[0..ef]; +} + +fn tr_copy(src: []const f64, dst: []f64) -> void { + var i : usize = 0; + while (i < src.len) { + dst[i] = src[i]; + i = i + 1; + } +} + +/// Segment s of the flat vector into the model: TWN trits and alpha for a ternary matrix (the STE +/// requantize), a plain copy for the float parameters +fn tr_scatter(s: usize) -> void { + const kind : u8 = seg_kind[s]; + const off : usize = seg_off[s]; + const src = p_val[off..off + seg_len[s]]; + const bw = &model.blocks[seg_block[s]]; + if (kind == K_Q) { + bw.alpha_q = quantize_twn(src, &bw.w_q); + } else if (kind == K_K) { + bw.alpha_k = quantize_twn(src, &bw.w_k); + } else if (kind == K_V) { + bw.alpha_v = quantize_twn(src, &bw.w_v); + } else if (kind == K_O) { + bw.alpha_o = quantize_twn(src, &bw.w_o); + } else if (kind == K_UP) { + bw.alpha_up = quantize_twn(src, &bw.w_up); + } else if (kind == K_DOWN) { + bw.alpha_down = quantize_twn(src, &bw.w_down); + } else if (kind == K_GAMMA) { + tr_copy(src, &bw.norm1_gamma); + } else if (kind == K_BUP) { + tr_copy(src, &bw.b_up); + } else if (kind == K_BDOWN) { + tr_copy(src, &bw.b_down); + } else if (kind == K_OUT) { + model.alpha_out = quantize_twn(src, &model.w_out); + } else if (kind == K_BOUT) { + tr_copy(src, &model.b_out); + } else { + tr_copy(src, &model.embed); + } +} + +fn tr_scatter_all() -> void { + var s : usize = 0; + while (s < NSEG) { + tr_scatter(s); + s = s + 1; + } +} + +fn tr_xs(seed: u64) -> u64 { + var x : u64 = seed; + x = x ^ (x << 13); + x = x ^ (x >> 7); + x = x ^ (x << 17); + return x; +} + +/// A uniform number in [0, 1) from the top 53 bits of a state +fn tr_unit(x: u64) -> f64 { + const top : u64 = x >> 11; + return @as(f64, @floatFromInt(top)) * (1.0 / 9007199254740992.0); +} + +/// Fresh model from a seed: shadows and the embedding uniform in +-1/sqrt(fan_in), gains 1, +/// biases 0, LAMB moments 0; trits and alphas by TWN. The window PRNG starts at the seed. +fn tr_init(seed: u64) -> void { + const total : usize = tr_build_segments(); + var x : u64 = seed ^ INIT_SALT; + if (x == 0) { + x = INIT_SALT; + } + var s : usize = 0; + while (s < NSEG) { + const kind : u8 = seg_kind[s]; + const off : usize = seg_off[s]; + const n : usize = seg_len[s]; + var value : f64 = 0.0; + var bound : f64 = 0.0; + if (tr_kind_ternary(kind) or kind == K_EMBED) { + var fan_in : usize = EMBED_DIM; + if (kind == K_DOWN) { + fan_in = FF_DIM; + } + bound = 1.0 / sqrt(@as(f64, @floatFromInt(fan_in))); + } else if (kind == K_GAMMA) { + value = 1.0; + } + var i : usize = 0; + while (i < n) { + if (bound > 0.0) { + x = tr_xs(x); + p_val[off + i] = bound * (2.0 * tr_unit(x) - 1.0); + } else { + p_val[off + i] = value; + } + i = i + 1; + } + s = s + 1; + } + var k : usize = 0; + while (k < total) { + m_val[k] = 0.0; + v_val[k] = 0.0; + g_val[k] = 0.0; + k = k + 1; + } + beta1_pow = 1.0; + beta2_pow = 1.0; + rng_state = seed; + if (rng_state == 0) { + rng_state = 1; + } + nudge_which = MAT_NONE; + vsa_on = true; + gate_reset(); + tr_scatter_all(); +} + +; --------------------------------------------------------------------------------------------- +; 3. The optimizer: schedule, global clip, LAMB +; --------------------------------------------------------------------------------------------- + +/// Learning rate of update t (0-based): linear warmup to the peak over `warmup` updates, then a +/// cosine from the peak to the floor at update `total`, and the floor after it +fn tr_lr_at(t: u64, warmup: u64, total: u64) -> f64 { + if (t < warmup) { + return LR_PEAK * @as(f64, @floatFromInt(t + 1)) / @as(f64, @floatFromInt(warmup)); + } + if (total <= warmup) { + return LR_FLOOR; + } + var prog : f64 = @as(f64, @floatFromInt(t - warmup)) / @as(f64, @floatFromInt(total - warmup)); + if (prog > 1.0) { + prog = 1.0; + } + return LR_FLOOR + 0.5 * (LR_PEAK - LR_FLOOR) * (1.0 + cos(TR_PI * prog)); +} + +fn tr_norm(g: []const f64) -> f64 { + var s0 : f64 = 0.0; + var s1 : f64 = 0.0; + var s2 : f64 = 0.0; + var s3 : f64 = 0.0; + var i : usize = 0; + while (i + 4 <= g.len) { + s0 = s0 + g[i] * g[i]; + s1 = s1 + g[i + 1] * g[i + 1]; + s2 = s2 + g[i + 2] * g[i + 2]; + s3 = s3 + g[i + 3] * g[i + 3]; + i = i + 4; + } + while (i < g.len) { + s0 = s0 + g[i] * g[i]; + i = i + 1; + } + return sqrt((s0 + s1) + (s2 + s3)); +} + +/// Scale g down to L2 norm max_norm when it is above it; returns the norm before +fn tr_clip_global(g: []f64, max_norm: f64) -> f64 { + const norm : f64 = tr_norm(g); + if (norm > max_norm) { + const scale : f64 = max_norm / norm; + var i : usize = 0; + while (i < g.len) { + g[i] = g[i] * scale; + i = i + 1; + } + } + return norm; +} + +/// The averaged gradient of the batch into the flat vector +fn tr_gather(scale: f64) -> void { + var s : usize = 0; + while (s < NSEG) { + const src = tr_grad_of(s); + const off : usize = seg_off[s]; + var i : usize = 0; + while (i < src.len) { + g_val[off + i] = src[i] * scale; + i = i + 1; + } + s = s + 1; + } +} + +/// One LAMB update of every segment (autograd.zig lambStepSlice): Adam moments with bias +/// correction, r = m_hat / (sqrt(v_hat) + eps) + wd w, trust = ||w|| / ||r|| clamped to 10, +/// w -= lr trust r. The update direction r is kept in g_val. +fn tr_lamb(lr: f64) -> void { + beta1_pow = beta1_pow * LAMB_BETA1; + beta2_pow = beta2_pow * LAMB_BETA2; + const bc1 : f64 = 1.0 - beta1_pow; + const bc2 : f64 = 1.0 - beta2_pow; + var s : usize = 0; + while (s < NSEG) { + const off : usize = seg_off[s]; + const end : usize = off + seg_len[s]; + var wd : f64 = 0.0; + if (tr_kind_decay(seg_kind[s])) { + wd = WEIGHT_DECAY; + } + var i : usize = off; + while (i < end) { + const g : f64 = g_val[i]; + const m : f64 = LAMB_BETA1 * m_val[i] + (1.0 - LAMB_BETA1) * g; + const v : f64 = LAMB_BETA2 * v_val[i] + (1.0 - LAMB_BETA2) * g * g; + m_val[i] = m; + v_val[i] = v; + g_val[i] = (m / bc1) / (sqrt(v / bc2) + LAMB_EPS) + wd * p_val[i]; + i = i + 1; + } + const w_norm : f64 = tr_norm(p_val[off..end]); + const r_norm : f64 = tr_norm(g_val[off..end]); + var trust : f64 = 1.0; + if (w_norm > 0.0 and r_norm > 0.0) { + trust = min(w_norm / r_norm, TRUST_CLAMP); + } + const rate : f64 = lr * trust; + i = off; + while (i < end) { + p_val[i] = p_val[i] - rate * g_val[i]; + i = i + 1; + } + s = s + 1; + } +} + +/// A window start in [0, n - 28]: the 27 inputs and the target lie inside the n tokens +fn tr_next_start(n: usize) -> usize { + rng_state = tr_xs(rng_state); + const span : u64 = @intCast(n - CONTEXT_LEN); + return @intCast(rng_state % span); +} + +/// One training step: `batch` windows (sampled, or fixed_starts), mean loss, averaged gradient, +/// global clip, LAMB at update t's rate (or at lr_fixed when it is positive), requantize +fn tr_train_step(t: u64, batch: usize, lr_fixed: f64, fixed: bool) -> f64 { + zero_weight_gradients(); + var loss_sum : f64 = 0.0; + var wi : usize = 0; + while (wi < batch) { + var start : usize = 0; + if (fixed) { + start = fixed_starts[wi]; + } else { + start = tr_next_start(ntrain); + } + loss_sum = loss_sum + tr_window_grad(train_tok[0..ntrain], start); + wi = wi + 1; + } + const inv : f64 = 1.0 / @as(f64, @floatFromInt(batch)); + tr_gather(inv); + last_gnorm = tr_clip_global(g_val[0..NPARAM], CLIP_NORM); + var lr : f64 = lr_fixed; + if (lr <= 0.0) { + lr = tr_lr_at(t, cfg_warmup, cfg_total); + } + last_lr = lr; + tr_lamb(lr); + tr_scatter_all(); + last_loss = loss_sum * inv; + return last_loss; +} + +/// Mean loss of `count` fixed held-out windows, evenly spaced over valid.tok +fn tr_valid_loss(count: usize) -> f64 { + var stride : usize = (nvalid - CONTEXT_LEN - 1) / count; + if (stride == 0) { + stride = 1; + } + var sum : f64 = 0.0; + var used : usize = 0; + var k : usize = 0; + while (k < count) { + const start : usize = k * stride; + if (start + CONTEXT_LEN < nvalid) { + sum = sum + tr_window_loss(valid_tok[0..nvalid], start); + used = used + 1; + } + k = k + 1; + } + return sum / @as(f64, @floatFromInt(used)); +} + +; --------------------------------------------------------------------------------------------- +; 4. Host I/O: whole reads and writes, the clock, stdout lines, the config +; --------------------------------------------------------------------------------------------- + +fn tr_note(rc: i32) -> void { + if (rc < 0) { + io_fail = io_fail + 1; + } +} + +/// Read until buf is full or the file ends; the byte count, or a negative errno +fn tr_read_full(fd: i32, buf: []u8) -> i64 { + var got : usize = 0; + var done : bool = false; + while (!done and got < buf.len) { + const r : i32 = host_read(fd, buf[got..buf.len]); + if (r < 0) { + return @as(i64, r); + } + if (r == 0) { + done = true; + } else { + got = got + @as(usize, @intCast(r)); + } + } + return @as(i64, @intCast(got)); +} + +/// Write all of buf; 0, or a negative errno +fn tr_write_all(fd: i32, buf: []const u8) -> i32 { + var put : usize = 0; + while (put < buf.len) { + const r : i32 = host_write(fd, buf[put..buf.len]); + if (r <= 0) { + return 0 - 1; + } + put = put + @as(usize, @intCast(r)); + } + return 0; +} + +/// Seconds since boot from /proc/uptime (10 ms resolution): the wall clock of the timings +fn tr_now() -> f64 { + const fd : i32 = host_open("/proc/uptime", 0); + if (fd < 0) { + return 0.0; + } + const n : i32 = host_read(fd, &up_buf); + tr_note(host_close(fd)); + if (n <= 0) { + return 0.0; + } + const len : usize = @intCast(n); + var secs : f64 = 0.0; + var frac : f64 = 0.0; + var i : usize = 0; + while (i < len) { + const ch : u8 = up_buf[i]; + if (ch == 46) { + frac = 0.1; + } else if (ch >= 48 and ch <= 57) { + const d : f64 = @floatFromInt(ch - 48); + if (frac == 0.0) { + secs = secs * 10.0 + d; + } else { + secs = secs + d * frac; + frac = frac * 0.1; + } + } else { + i = len; + } + i = i + 1; + } + return secs; +} + +fn ln_reset() -> void { + ln_n = 0; +} + +fn ln_byte(ch: u8) -> void { + if (ln_n < 511) { + ln_buf[ln_n] = ch; + ln_n = ln_n + 1; + } +} + +fn ln_str(text: []const u8) -> void { + var i : usize = 0; + while (i < text.len) { + ln_byte(text[i]); + i = i + 1; + } +} + +fn ln_u64(value: u64) -> void { + if (value == 0) { + ln_byte(48); + return; + } + var digits : [24]u8 = undefined; + var n : usize = 0; + var x : u64 = value; + while (x > 0) { + digits[n] = @intCast(48 + x % 10); + x = x / 10; + n = n + 1; + } + while (n > 0) { + n = n - 1; + ln_byte(digits[n]); + } +} + +/// x with `places` decimals, rounded half up; nan and values above 1e15 by name +fn ln_f64(x: f64, places: u32) -> void { + if (x != x) { + ln_str("nan"); + return; + } + var y : f64 = x; + if (y < 0.0) { + ln_byte(45); + y = 0.0 - y; + } + if (y > 1000000000000000.0) { + ln_str("inf"); + return; + } + var scale : f64 = 1.0; + var unit : u64 = 1; + var k : u32 = 0; + while (k < places) { + scale = scale * 10.0; + unit = unit * 10; + k = k + 1; + } + const all : u64 = @intFromFloat(y * scale + 0.5); + ln_u64(all / unit); + if (places > 0) { + ln_byte(46); + var rest : u64 = all % unit; + var div : u64 = unit / 10; + while (div > 0) { + ln_byte(@intCast(48 + rest / div)); + rest = rest % div; + div = div / 10; + } + } +} + +fn ln_flush() -> void { + ln_byte(10); + tr_note(tr_write_all(1, ln_buf[0..ln_n])); + ln_n = 0; +} + +fn tr_key_is(off: usize, len: usize, key: []const u8) -> bool { + if (len != key.len) { + return false; + } + var i : usize = 0; + while (i < len) { + if (cfg_buf[off + i] != key[i]) { + return false; + } + i = i + 1; + } + return true; +} + +fn tr_cfg_set(off: usize, len: usize, value: u64) -> void { + if (tr_key_is(off, len, "steps")) { + cfg_steps = value; + } else if (tr_key_is(off, len, "total")) { + cfg_total = value; + } else if (tr_key_is(off, len, "warmup")) { + cfg_warmup = value; + } else if (tr_key_is(off, len, "eval_every")) { + cfg_eval_every = value; + } else if (tr_key_is(off, len, "ckpt_every")) { + cfg_ckpt_every = value; + } else if (tr_key_is(off, len, "seed")) { + cfg_seed = value; + } else if (tr_key_is(off, len, "speed")) { + cfg_speed = value; + } else if (tr_key_is(off, len, "resume")) { + cfg_resume = value; + } else if (tr_key_is(off, len, "synth")) { + cfg_synth = value; + } else if (tr_key_is(off, len, "eval_windows")) { + cfg_eval_windows = value; + } +} + +/// key=value lines of /tmp/hslm/train.cfg, if the file exists; unknown keys are ignored +fn tr_read_cfg() -> void { + const fd : i32 = host_open("/tmp/hslm/train.cfg", 0); + if (fd < 0) { + return; + } + const got : i64 = tr_read_full(fd, &cfg_buf); + tr_note(host_close(fd)); + if (got <= 0) { + return; + } + const len : usize = @intCast(got); + var i : usize = 0; + while (i < len) { + const key_off : usize = i; + while (i < len and cfg_buf[i] != 61 and cfg_buf[i] != 10) { + i = i + 1; + } + if (i < len and cfg_buf[i] == 61) { + const key_len : usize = i - key_off; + i = i + 1; + var value : u64 = 0; + while (i < len and cfg_buf[i] >= 48 and cfg_buf[i] <= 57) { + value = value * 10 + @as(u64, cfg_buf[i] - 48); + i = i + 1; + } + tr_cfg_set(key_off, key_len, value); + } + while (i < len and cfg_buf[i] != 10) { + i = i + 1; + } + i = i + 1; + } +} + +; --------------------------------------------------------------------------------------------- +; 5. Tokens: HTOK files, and a synthetic corpus written by this binary when asked (synth=1) +; --------------------------------------------------------------------------------------------- + +fn tr_open_tok(which: u8, mode: i32) -> i32 { + if (which == 0) { + return host_open("/tmp/hslm/train.tok", mode); + } + if (which == 1) { + return host_open("/tmp/hslm/valid.tok", mode); + } + if (which == 2) { + return host_open("/tmp/hslm/synth_train.tok", mode); + } + return host_open("/tmp/hslm/synth_valid.tok", mode); +} + +/// Read an HTOK file into dst; the token count, or -1 (no file), -2 (bad header or size), +/// -3 (an id of 729 or more), -4 (more tokens than dst holds) +fn tr_read_tok(which: u8, dst: []u16) -> i64 { + const fd : i32 = tr_open_tok(which, 0); + if (fd < 0) { + return 0 - 1; + } + const hg : i64 = tr_read_full(fd, io_hdr[0..12]); + if (hg != 12 or read_u32_le(&io_hdr, 0) != HTOK_MAGIC or read_u32_le(&io_hdr, 4) != HTOK_VERSION) { + tr_note(host_close(fd)); + return 0 - 2; + } + const count : usize = @intCast(read_u32_le(&io_hdr, 8)); + if (count > dst.len) { + tr_note(host_close(fd)); + return 0 - 4; + } + var idx : usize = 0; + var bad : bool = false; + while (idx < count and !bad) { + var want : usize = 2 * (count - idx); + if (want > IO_CHUNK) { + want = IO_CHUNK; + } + const got : i64 = tr_read_full(fd, io_chunk[0..want]); + if (got != @as(i64, @intCast(want))) { + tr_note(host_close(fd)); + return 0 - 2; + } + var b : usize = 0; + while (b < want) { + const lo : u16 = @as(u16, io_chunk[b]); + const hi : u16 = @as(u16, io_chunk[b + 1]); + const id : u16 = lo | (hi << 8); + if (id >= 729) { + bad = true; + } + dst[idx] = id; + idx = idx + 1; + b = b + 2; + } + } + const extra : i64 = tr_read_full(fd, io_hdr[0..1]); + tr_note(host_close(fd)); + if (bad) { + return 0 - 3; + } + if (extra != 0) { + return 0 - 2; + } + return @as(i64, @intCast(count)); +} + +/// Write src as an HTOK file (only the synthetic names: the real corpus is never written here) +fn tr_write_tok(which: u8, src: []const u16) -> i32 { + const fd : i32 = tr_open_tok(which, 1); + if (fd < 0) { + return 0 - 1; + } + write_u32_le(&io_hdr, 0, HTOK_MAGIC); + write_u32_le(&io_hdr, 4, HTOK_VERSION); + write_u32_le(&io_hdr, 8, @intCast(src.len)); + var rc : i32 = tr_write_all(fd, io_hdr[0..12]); + var idx : usize = 0; + while (rc == 0 and idx < src.len) { + var n : usize = 0; + while (n + 2 <= IO_CHUNK and idx < src.len) { + io_chunk[n] = @intCast(src[idx] & 255); + io_chunk[n + 1] = @intCast(src[idx] >> 8); + n = n + 2; + idx = idx + 1; + } + rc = tr_write_all(fd, io_chunk[0..n]); + } + const c : i32 = host_close(fd); + if (rc != 0 or c != 0) { + return 0 - 2; + } + return 0; +} + +/// A learnable synthetic stream: next = (7 cur + 3) mod 729, one token in ten drawn at random +fn tr_synth(dst: []u16, seed: u64) -> void { + var x : u64 = seed; + if (x == 0) { + x = 1; + } + var cur : u64 = 1; + var i : usize = 0; + while (i < dst.len) { + x = tr_xs(x); + var nxt : u64 = (7 * cur + 3) % 729; + if (x % 10 == 0) { + nxt = (x >> 20) % 729; + } + dst[i] = @intCast(nxt); + cur = nxt; + i = i + 1; + } +} + +/// train.tok and valid.tok; without them and with synth=1, the synthetic pair (written once, +/// then read back through the same reader) +fn tr_load_data() -> i32 { + var nt : i64 = tr_read_tok(0, &train_tok); + var nv : i64 = 0; + if (nt == 0 - 1 and cfg_synth == 1) { + data_synthetic = true; + nt = tr_read_tok(2, &train_tok); + if (nt < 0) { + tr_synth(train_tok[0..SYN_TRAIN], 11); + tr_synth(valid_tok[0..SYN_VALID], 13); + if (tr_write_tok(2, train_tok[0..SYN_TRAIN]) != 0 or tr_write_tok(3, valid_tok[0..SYN_VALID]) != 0) { + return RC_SYNTH_WRITE; + } + nt = tr_read_tok(2, &train_tok); + } + nv = tr_read_tok(3, &valid_tok); + } else { + nv = tr_read_tok(1, &valid_tok); + } + if (nt < @as(i64, @intCast(CONTEXT_LEN + 2))) { + return RC_TRAIN_TOK; + } + if (nv < @as(i64, @intCast(CONTEXT_LEN + 2))) { + return RC_VALID_TOK; + } + ntrain = @intCast(nt); + nvalid = @intCast(nv); + return RC_OK; +} + +; --------------------------------------------------------------------------------------------- +; 6. Checkpoints: MSLH (nn/hslm_data) and the resumable state, each read back and compared +; --------------------------------------------------------------------------------------------- + +/// /tmp/hslm/ckpt_.bin (kind 0) or .state (kind 1), NUL-terminated, in path_buf +fn tr_path(step: u64, kind: u8) -> void { + ln_reset(); + ln_str("/tmp/hslm/ckpt_"); + ln_u64(step); + if (kind == 0) { + ln_str(".bin"); + } else { + ln_str(".state"); + } + var i : usize = 0; + while (i < ln_n) { + path_buf[i] = ln_buf[i]; + i = i + 1; + } + path_buf[ln_n] = 0; + ln_n = 0; +} + +fn tr_open_path(mode: i32) -> i32 { + return host_open(@ptrCast(&path_buf), mode); +} + +fn tr_put_f32(buf: []u8, off: usize, x: f64) -> void { + const f : f32 = @floatCast(x); + const bits : u32 = @bitCast(f); + write_u32_le(buf, off, bits); +} + +fn tr_get_f32(buf: []const u8, off: usize) -> f64 { + const bits : u32 = read_u32_le(buf, off); + const f : f32 = @bitCast(bits); + return @floatCast(f); +} + +fn tr_put_u64(buf: []u8, off: usize, value: u64) -> void { + const lo : u32 = @intCast(value & 4294967295); + const hi : u32 = @intCast(value >> 32); + write_u32_le(buf, off, lo); + write_u32_le(buf, off + 4, hi); +} + +fn tr_get_u64(buf: []const u8, off: usize) -> u64 { + const lo : u64 = read_u32_le(buf, off); + const hi : u64 = read_u32_le(buf, off + 4); + return (hi << 32) | lo; +} + +fn tr_put_f64(buf: []u8, off: usize, x: f64) -> void { + const bits : u64 = @bitCast(x); + tr_put_u64(buf, off, bits); +} + +fn tr_get_f64(buf: []const u8, off: usize) -> f64 { + const bits : u64 = tr_get_u64(buf, off); + const x : f64 = @bitCast(bits); + return x; +} + +fn tr_put_seg_f32(buf: []u8, off: usize, s: usize) -> void { + const base : usize = seg_off[s]; + var i : usize = 0; + while (i < seg_len[s]) { + tr_put_f32(buf, off + 4 * i, p_val[base + i]); + i = i + 1; + } +} + +fn tr_get_seg_f32(buf: []const u8, off: usize, s: usize) -> void { + const base : usize = seg_off[s]; + var i : usize = 0; + while (i < seg_len[s]) { + p_val[base + i] = tr_get_f32(buf, off + 4 * i); + i = i + 1; + } +} + +/// The MSLH image: header, then the sections in nn/hslm_data's order and offsets +fn tr_pack_mslh(buf: []u8, step: u32, loss: f64) -> void { + write_header(buf, step, loss); + tr_put_seg_f32(buf, output_offset(), tr_seg(K_OUT, 0)); + var b : usize = 0; + while (b < NUM_BLOCKS) { + tr_put_seg_f32(buf, shadow_up_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, b), tr_seg(K_UP, b)); + tr_put_seg_f32(buf, shadow_down_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, b), tr_seg(K_DOWN, b)); + tr_put_seg_f32(buf, bias_up_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, b), tr_seg(K_BUP, b)); + tr_put_seg_f32(buf, bias_down_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, b), tr_seg(K_BDOWN, b)); + b = b + 1; + } +} + +/// The sections of an MSLH image back into the flat vector (as f32 values) +fn tr_unpack_mslh(buf: []const u8) -> void { + tr_get_seg_f32(buf, output_offset(), tr_seg(K_OUT, 0)); + var b : usize = 0; + while (b < NUM_BLOCKS) { + tr_get_seg_f32(buf, shadow_up_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, b), tr_seg(K_UP, b)); + tr_get_seg_f32(buf, shadow_down_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, b), tr_seg(K_DOWN, b)); + tr_get_seg_f32(buf, bias_up_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, b), tr_seg(K_BUP, b)); + tr_get_seg_f32(buf, bias_down_offset(VOCAB_SIZE, EMBED_DIM, FF_DIM, b), tr_seg(K_BDOWN, b)); + b = b + 1; + } +} + +fn tr_count_diff(a: []const u8, b: []const u8, n: usize) -> usize { + var diff : usize = 0; + var i : usize = 0; + while (i < n) { + if (a[i] != b[i]) { + diff = diff + 1; + } + i = i + 1; + } + return diff; +} + +/// Write ckpt_.bin, read it back whole, compare every byte and the header +fn tr_write_mslh(step: u64) -> i32 { + tr_pack_mslh(&ck_buf, @intCast(step), last_loss); + tr_path(step, 0); + const fd : i32 = tr_open_path(1); + if (fd < 0) { + return RC_CKPT_WRITE; + } + const w : i32 = tr_write_all(fd, &ck_buf); + const c : i32 = host_close(fd); + if (w != 0 or c != 0) { + return RC_CKPT_WRITE; + } + const rd : i32 = tr_open_path(0); + if (rd < 0) { + return RC_CKPT_VERIFY; + } + const got : i64 = tr_read_full(rd, &ck_back); + const extra : i64 = tr_read_full(rd, io_hdr[0..1]); + tr_note(host_close(rd)); + if (got != @as(i64, @intCast(MSLH_BYTES)) or extra != 0) { + return RC_CKPT_VERIFY; + } + if (tr_count_diff(&ck_buf, &ck_back, MSLH_BYTES) != 0) { + return RC_CKPT_VERIFY; + } + if (parse_header(&ck_back) != HEADER_OK or @as(u64, header_step(&ck_back)) != step) { + return RC_CKPT_VERIFY; + } + return RC_OK; +} + +fn tr_state_get(e: usize) -> f64 { + if (e < NPARAM) { + return p_val[e]; + } + if (e < 2 * NPARAM) { + return m_val[e - NPARAM]; + } + return v_val[e - 2 * NPARAM]; +} + +fn tr_state_set(e: usize, x: f64) -> void { + if (e < NPARAM) { + p_val[e] = x; + } else if (e < 2 * NPARAM) { + m_val[e - NPARAM] = x; + } else { + v_val[e - 2 * NPARAM] = x; + } +} + +/// count f64 of the state stream (p, then m, then v) from element `first`, into buf +fn tr_state_encode(first: usize, count: usize, buf: []u8) -> void { + var k : usize = 0; + while (k < count) { + tr_put_f64(buf, 8 * k, tr_state_get(first + k)); + k = k + 1; + } +} + +fn tr_state_decode(first: usize, count: usize, buf: []const u8) -> void { + var k : usize = 0; + while (k < count) { + tr_state_set(first + k, tr_get_f64(buf, 8 * k)); + k = k + 1; + } +} + +fn tr_state_header(buf: []u8, step: u64) -> void { + var z : usize = 0; + while (z < STATE_HEADER) { + buf[z] = 0; + z = z + 1; + } + write_u32_le(buf, 0, STATE_MAGIC); + write_u32_le(buf, 4, STATE_VERSION); + write_u32_le(buf, 8, @intCast(step)); + write_u32_le(buf, 12, @intCast(NPARAM)); + tr_put_u64(buf, 16, rng_state); + tr_put_f64(buf, 24, beta1_pow); + tr_put_f64(buf, 32, beta2_pow); + tr_put_f64(buf, 40, last_loss); + write_u32_le(buf, 48, @intCast(NSEG)); +} + +/// Write ckpt_.state, then read it back and compare it chunk by chunk with the encoding +fn tr_write_state(step: u64) -> i32 { + tr_path(step, 1); + const fd : i32 = tr_open_path(1); + if (fd < 0) { + return RC_STATE_WRITE; + } + tr_state_header(&io_hdr, step); + var rc : i32 = tr_write_all(fd, &io_hdr); + const total : usize = 3 * NPARAM; + const per : usize = IO_CHUNK / 8; + var e : usize = 0; + while (rc == 0 and e < total) { + var n : usize = total - e; + if (n > per) { + n = per; + } + tr_state_encode(e, n, &io_chunk); + rc = tr_write_all(fd, io_chunk[0..8 * n]); + e = e + n; + } + const c : i32 = host_close(fd); + if (rc != 0 or c != 0) { + return RC_STATE_WRITE; + } + const rd : i32 = tr_open_path(0); + if (rd < 0) { + return RC_STATE_VERIFY; + } + var same : bool = tr_read_full(rd, &io_hdr_back) == @as(i64, @intCast(STATE_HEADER)); + if (same and tr_count_diff(&io_hdr, &io_hdr_back, STATE_HEADER) != 0) { + same = false; + } + e = 0; + while (same and e < total) { + var n : usize = total - e; + if (n > per) { + n = per; + } + tr_state_encode(e, n, &io_chunk); + const got : i64 = tr_read_full(rd, io_back[0..8 * n]); + if (got != @as(i64, @intCast(8 * n)) or tr_count_diff(&io_chunk, &io_back, 8 * n) != 0) { + same = false; + } + e = e + n; + } + if (same and tr_read_full(rd, io_hdr_back[0..1]) != 0) { + same = false; + } + tr_note(host_close(rd)); + if (!same) { + return RC_STATE_VERIFY; + } + return RC_OK; +} + +/// Restart from ckpt_.state: parameters, moments, beta powers, PRNG, last loss +fn tr_load_state(step: u64) -> i32 { + tr_path(step, 1); + const rd : i32 = tr_open_path(0); + if (rd < 0) { + return RC_RESUME; + } + const hg : i64 = tr_read_full(rd, &io_hdr_back); + var ok : bool = hg == @as(i64, @intCast(STATE_HEADER)); + if (ok) { + ok = read_u32_le(&io_hdr_back, 0) == STATE_MAGIC and read_u32_le(&io_hdr_back, 4) == STATE_VERSION and @as(u64, read_u32_le(&io_hdr_back, 8)) == step and @as(usize, read_u32_le(&io_hdr_back, 12)) == NPARAM and @as(usize, read_u32_le(&io_hdr_back, 48)) == NSEG; + } + const total : usize = 3 * NPARAM; + const per : usize = IO_CHUNK / 8; + var e : usize = 0; + while (ok and e < total) { + var n : usize = total - e; + if (n > per) { + n = per; + } + const got : i64 = tr_read_full(rd, io_back[0..8 * n]); + if (got != @as(i64, @intCast(8 * n))) { + ok = false; + } else { + tr_state_decode(e, n, &io_back); + } + e = e + n; + } + tr_note(host_close(rd)); + if (!ok) { + return RC_RESUME; + } + rng_state = tr_get_u64(&io_hdr_back, 16); + beta1_pow = tr_get_f64(&io_hdr_back, 24); + beta2_pow = tr_get_f64(&io_hdr_back, 32); + last_loss = tr_get_f64(&io_hdr_back, 40); + tr_scatter_all(); + return RC_OK; +} + +; --------------------------------------------------------------------------------------------- +; 7. main +; --------------------------------------------------------------------------------------------- + +fn tr_eval_line(step: u64, have_train: bool, avg_loss: f64, valid: f64, train_secs: f64, train_steps: u64) -> void { + ln_reset(); + ln_str("step "); + ln_u64(step); + if (have_train) { + ln_str(" train_loss "); + ln_f64(last_loss, 4); + ln_str(" train_ppl "); + ln_f64(exp(last_loss), 2); + ln_str(" avg_loss "); + ln_f64(avg_loss, 4); + } else { + ln_str(" train_loss - train_ppl - avg_loss -"); + } + ln_str(" valid_loss "); + ln_f64(valid, 4); + ln_str(" valid_ppl "); + ln_f64(exp(valid), 2); + if (train_steps > 0 and train_secs > 0.0) { + const sps : f64 = train_secs / @as(f64, @floatFromInt(train_steps)); + ln_str(" tok_per_s "); + ln_f64(@as(f64, @floatFromInt(BATCH * CONTEXT_LEN)) / sps, 1); + ln_str(" s_per_step "); + ln_f64(sps, 4); + } else { + ln_str(" tok_per_s - s_per_step -"); + } + ln_str(" lr "); + ln_f64(last_lr, 8); + ln_str(" gnorm "); + ln_f64(last_gnorm, 4); + ln_flush(); +} + +fn tr_main() -> i32 { + const t_begin : f64 = tr_now(); + tr_read_cfg(); + speed_level = @intCast(cfg_speed % 3); + const rc_data : i32 = tr_load_data(); + if (rc_data != RC_OK) { + ln_reset(); + ln_str("train_hslm: token files missing or bad, code "); + ln_u64(@intCast(rc_data)); + ln_flush(); + return rc_data; + } + tr_init(cfg_seed); + var step : u64 = 0; + if (cfg_resume > 0) { + const rc_res : i32 = tr_load_state(cfg_resume); + if (rc_res != RC_OK) { + return rc_res; + } + step = cfg_resume; + } + ln_reset(); + ln_str("train_hslm: train_tokens "); + ln_u64(@intCast(ntrain)); + ln_str(" valid_tokens "); + ln_u64(@intCast(nvalid)); + if (data_synthetic) { + ln_str(" (synthetic)"); + } + ln_str(" nparam "); + ln_u64(@intCast(NPARAM)); + ln_str(" batch "); + ln_u64(@intCast(BATCH)); + ln_str(" steps "); + ln_u64(cfg_steps); + ln_str(" total "); + ln_u64(cfg_total); + ln_str(" warmup "); + ln_u64(cfg_warmup); + ln_str(" seed "); + ln_u64(cfg_seed); + ln_str(" speed "); + ln_u64(@as(u64, speed_level)); + ln_str(" from_step "); + ln_u64(step); + ln_flush(); + var eval_count : u64 = cfg_eval_windows; + if (eval_count == 0) { + eval_count = 1; + } + tr_eval_line(step, false, 0.0, tr_valid_loss(@intCast(eval_count)), 0.0, 0); + var train_secs : f64 = 0.0; + var all_secs : f64 = 0.0; + var interval_steps : u64 = 0; + var interval_loss : f64 = 0.0; + var run_steps : u64 = 0; + var first_loss : f64 = 0.0; + var rc : i32 = RC_OK; + while (rc == RC_OK and step < cfg_steps) { + const t0 : f64 = tr_now(); + const loss : f64 = tr_train_step(step, BATCH, 0.0, false); + const t1 : f64 = tr_now(); + step = step + 1; + train_secs = train_secs + (t1 - t0); + all_secs = all_secs + (t1 - t0); + interval_steps = interval_steps + 1; + interval_loss = interval_loss + loss; + if (run_steps == 0) { + first_loss = loss; + } + run_steps = run_steps + 1; + if (loss != loss) { + rc = RC_NAN; + } + if (rc == RC_OK and (step % cfg_eval_every == 0 or step == cfg_steps)) { + const avg : f64 = interval_loss / @as(f64, @floatFromInt(interval_steps)); + tr_eval_line(step, true, avg, tr_valid_loss(@intCast(eval_count)), train_secs, interval_steps); + train_secs = 0.0; + interval_steps = 0; + interval_loss = 0.0; + } + if (rc == RC_OK and (step % cfg_ckpt_every == 0 or step == cfg_steps)) { + rc = tr_write_mslh(step); + if (rc == RC_OK) { + rc = tr_write_state(step); + } + ln_reset(); + ln_str("checkpoint "); + ln_u64(step); + if (rc == RC_OK) { + ln_str(" written and read back byte-equal: /tmp/hslm/ckpt_"); + } else { + ln_str(" FAILED, code "); + ln_u64(@intCast(rc)); + ln_str(" /tmp/hslm/ckpt_"); + } + ln_u64(step); + ln_str(".bin (4971796 bytes) + .state"); + ln_flush(); + } + } + ln_reset(); + ln_str("done steps "); + ln_u64(run_steps); + ln_str(" train_s "); + ln_f64(all_secs, 2); + if (run_steps > 0 and all_secs > 0.0) { + const sps : f64 = all_secs / @as(f64, @floatFromInt(run_steps)); + ln_str(" s_per_step "); + ln_f64(sps, 4); + ln_str(" tok_per_s "); + ln_f64(@as(f64, @floatFromInt(BATCH * CONTEXT_LEN)) / sps, 1); + ln_str(" first_loss "); + ln_f64(first_loss, 4); + ln_str(" last_loss "); + ln_f64(last_loss, 4); + } + ln_str(" wall_s "); + ln_f64(tr_now() - t_begin, 2); + ln_str(" io_fail "); + ln_u64(@as(u64, io_fail)); + ln_str(" rc "); + ln_u64(@intCast(rc)); + ln_flush(); + return rc; +} + +/// The program `t27c compile --backend ml` builds: train, evaluate, checkpoint; any checkpoint +/// that does not read back byte-equal, a NaN loss or missing tokens fail the run +pub fn main() -> void { + assert tr_main() == RC_OK; +} + +; --------------------------------------------------------------------------------------------- +; 8. Test helpers (each returns the number its test checks) and the tests +; --------------------------------------------------------------------------------------------- + +/// A digest of `count` window starts drawn from `seed` over n tokens +fn tr_sample_digest(seed: u64, count: usize, n: usize) -> u64 { + rng_state = seed; + var h : u64 = 17; + var k : usize = 0; + while (k < count) { + const s : u64 = @intCast(tr_next_start(n)); + h = (h * 31 + s) % 4294967291; + k = k + 1; + } + return h; +} + +/// The largest of `count` starts drawn from `seed` over n tokens +fn tr_sample_max(seed: u64, count: usize, n: usize) -> usize { + rng_state = seed; + var mx : usize = 0; + var k : usize = 0; + while (k < count) { + const s : usize = tr_next_start(n); + if (s > mx) { + mx = s; + } + k = k + 1; + } + return mx; +} + +/// The global norm of a fixed 1000-entry vector times `scale`, after clipping to 1.0 +fn tr_clip_norm_after(scale: f64) -> f64 { + var i : usize = 0; + while (i < 1000) { + clip_buf[i] = scale * (@as(f64, @floatFromInt((i * 37) % 101)) - 50.0); + i = i + 1; + } + const before : f64 = tr_clip_global(&clip_buf, CLIP_NORM); + if (before <= CLIP_NORM) { + return 0.0 - 1.0; + } + return tr_norm(&clip_buf); +} + +/// How many entries clipping changes in a vector whose norm is already below 1.0 +fn tr_clip_changed(scale: f64) -> usize { + var i : usize = 0; + while (i < 1000) { + clip_buf[i] = scale * (@as(f64, @floatFromInt((i * 37) % 101)) - 50.0); + clip_ref[i] = clip_buf[i]; + i = i + 1; + } + const before : f64 = tr_clip_global(&clip_buf, CLIP_NORM); + var changed : usize = 0; + i = 0; + while (i < 1000) { + if (clip_buf[i] != clip_ref[i]) { + changed = changed + 1; + } + i = i + 1; + } + if (before > CLIP_NORM) { + changed = changed + 1000; + } + return changed; +} + +/// A small synthetic corpus in train_tok and valid_tok, and a fresh model +fn tr_tiny_setup(seed: u64) -> void { + tr_synth(train_tok[0..4096], 99); + tr_synth(valid_tok[0..1024], 98); + ntrain = 4096; + nvalid = 1024; + tr_init(seed); +} + +fn tr_fixed_loss(count: usize) -> f64 { + var sum : f64 = 0.0; + var k : usize = 0; + while (k < count) { + sum = sum + tr_window_loss(train_tok[0..ntrain], fixed_starts[k]); + k = k + 1; + } + return sum / @as(f64, @floatFromInt(count)); +} + +/// Two LAMB steps at lr 0.01 on two fixed windows; the loss on them after (before: tiny_before) +fn tr_tiny_after() -> f64 { + speed_level = 2; + tr_tiny_setup(7); + fixed_starts[0] = 100; + fixed_starts[1] = 2000; + tiny_before = tr_fixed_loss(2); + var t : u64 = 0; + while (t < 2) { + const seen : f64 = tr_train_step(t, 2, 0.01, true); + if (seen != seen) { + return seen; + } + t = t + 1; + } + return tr_fixed_loss(2); +} + +fn tr_tiny_before() -> f64 { + return tiny_before; +} + +/// The largest |difference| between the logits of one window at speed `fast` and at speed 0 +fn tr_speed_logit_diff(fast: u8) -> f64 { + tr_tiny_setup(5); + speed_level = 0; + tr_window_forward(train_tok[0..ntrain], 500, &ref_logits); + speed_level = fast; + tr_window_forward(train_tok[0..ntrain], 500, &win_logits); + var mx : f64 = 0.0; + var i : usize = 0; + while (i < VOCAB_SIZE) { + const d : f64 = abs(win_logits[i] - ref_logits[i]); + if (d > mx or d != d) { + mx = d; + } + i = i + 1; + } + return mx; +} + +/// The relative difference of the input gradient (it flows back through every +/// ternary_matmul_t_add) between speed 2 and speed 0 +fn tr_speed_grad_reldiff() -> f64 { + tr_tiny_setup(5); + speed_level = 0; + zero_weight_gradients(); + const l0 : f64 = tr_window_grad(train_tok[0..ntrain], 700); + tr_copy(&win_gin, &ref_gin); + speed_level = 2; + zero_weight_gradients(); + const l2 : f64 = tr_window_grad(train_tok[0..ntrain], 700); + var num : f64 = 0.0; + var den : f64 = 0.0; + var i : usize = 0; + while (i < EMBED_DIM) { + const d : f64 = win_gin[i] - ref_gin[i]; + num = num + d * d; + den = den + ref_gin[i] * ref_gin[i]; + i = i + 1; + } + if (l0 != l2 or den == 0.0) { + return 1.0; + } + return sqrt(num / den); +} + +/// Bytes that differ between an MSLH image and the image packed again from its own decoded floats +fn tr_mslh_roundtrip_diff() -> usize { + tr_init(3); + last_loss = 5.5; + tr_pack_mslh(&ck_buf, 7, last_loss); + tr_init(4); + tr_unpack_mslh(&ck_buf); + tr_pack_mslh(&ck_back, 7, 5.5); + return tr_count_diff(&ck_buf, &ck_back, MSLH_BYTES); +} + +fn tr_mslh_header_status() -> u32 { + return parse_header(&ck_back); +} + +fn tr_mslh_header_step() -> u32 { + return header_step(&ck_back); +} + +fn tr_mslh_header_loss() -> f64 { + return header_loss(&ck_back); +} + +/// One flipped byte in the copy is found: the diff count is 1 +fn tr_mslh_flip_diff() -> usize { + ck_back[123457] = ck_back[123457] ^ 1; + return tr_count_diff(&ck_buf, &ck_back, MSLH_BYTES); +} + +/// State stream round trip across the p / m boundary: encode a chunk, clear it, decode it; +/// the number of f64 whose bits differ from before +fn tr_state_roundtrip_diff() -> usize { + tr_init(9); + var k : usize = 0; + while (k < NPARAM) { + m_val[k] = 0.001 * p_val[k]; + v_val[k] = p_val[k] * p_val[k]; + k = k + 1; + } + const first : usize = NPARAM - 1000; + const per : usize = IO_CHUNK / 8; + k = 0; + while (k < per) { + st_save[k] = tr_state_get(first + k); + k = k + 1; + } + tr_state_encode(first, per, &io_chunk); + k = 0; + while (k < per) { + tr_state_set(first + k, 0.0); + k = k + 1; + } + tr_state_decode(first, per, &io_chunk); + var diff : usize = 0; + k = 0; + while (k < per) { + const a : u64 = @bitCast(st_save[k]); + const b : u64 = @bitCast(tr_state_get(first + k)); + if (a != b) { + diff = diff + 1; + } + k = k + 1; + } + return diff; +} + +fn tr_mslh_size() -> usize { + return real_file_bytes(); +} + +fn tr_param_count() -> usize { + return tr_build_segments(); +} + +test window_sampling_is_deterministic_for_a_seed + given first = tr_sample_digest(27, 1000, 100000) + and again = tr_sample_digest(27, 1000, 100000) + and other = tr_sample_digest(28, 1000, 100000) + then first == again + and first != other + +test window_starts_leave_room_for_the_target + given top = tr_sample_max(27, 5000, 100) + then top <= 72 + and top >= 60 + +test lr_at_step_zero_is_the_first_warmup_slice + given lr0 = tr_lr_at(0, 2000, 32000) + when diff = abs(lr0 - 0.001 / 2000.0) + then diff < 1e-18 + +test lr_at_the_end_of_warmup_is_the_peak + given last_warm = tr_lr_at(1999, 2000, 32000) + and first_cos = tr_lr_at(2000, 2000, 32000) + then abs(last_warm - 0.001) < 1e-15 + and abs(first_cos - 0.001) < 1e-15 + +test lr_at_the_end_is_the_floor + given lr_end = tr_lr_at(32000, 2000, 32000) + and lr_mid = tr_lr_at(17000, 2000, 32000) + and lr_after = tr_lr_at(40000, 2000, 32000) + then abs(lr_end - 0.000001) < 1e-15 + and abs(lr_mid - 0.0005005) < 1e-12 + and abs(lr_after - 0.000001) < 1e-15 + +test clip_bounds_the_global_norm + given after = tr_clip_norm_after(3.0) + then after <= 1.0 + 1e-12 + and after >= 1.0 - 1e-12 + +test clip_leaves_a_small_gradient_alone + given changed = tr_clip_changed(0.0001) + then changed == 0 + +test the_parameter_vector_is_the_whole_model + given n = tr_param_count() + then n == 2130138 + +test a_checkpoint_round_trip_on_a_buffer_is_byte_equal + given diff = tr_mslh_roundtrip_diff() + and status = tr_mslh_header_status() + and step = tr_mslh_header_step() + and loss = tr_mslh_header_loss() + and size = tr_mslh_size() + and flipped = tr_mslh_flip_diff() + then diff == 0 + and status == 0 + and step == 7 + and abs(loss - 5.5) < 1e-6 + and size == 4971796 + and flipped == 1 + +test the_state_stream_round_trips_bit_exactly + given diff = tr_state_roundtrip_diff() + then diff == 0 + +test speed_one_gives_the_reference_logits_to_the_bit + given diff = tr_speed_logit_diff(1) + then diff == 0.0 + +test speed_two_gives_the_reference_logits_to_the_bit + given diff = tr_speed_logit_diff(2) + then diff == 0.0 + +test speed_two_input_gradient_matches_the_reference + given rel = tr_speed_grad_reldiff() + then rel < 1e-12 + +test a_few_steps_on_a_tiny_fixed_slice_lower_the_loss + given after = tr_tiny_after() + and before = tr_tiny_before() + then before > 6.0 + and after < before From 004f875782a4f7a90c23854bddd26a4b30d01826 Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Sun, 18 Oct 2026 17:33:49 +0700 Subject: [PATCH 2/5] spec(ml): train_hslm -- drop lone ; lines (Refs #10372) Co-Authored-By: Claude Opus 5.5 --- specs/ml/train_hslm.t27 | 7 ------- 1 file changed, 7 deletions(-) diff --git a/specs/ml/train_hslm.t27 b/specs/ml/train_hslm.t27 index d77b8bb328..336dd54f45 100644 --- a/specs/ml/train_hslm.t27 +++ b/specs/ml/train_hslm.t27 @@ -3,20 +3,17 @@ ; tokens: straight-through shadow weights, LAMB, a warmup-cosine schedule, global-norm clipping, ; held-out evaluation, and MSLH checkpoints that are read back and compared byte for byte. ; Issue: gHashTag/t27#10372. -; ; WHAT IS REUSED, NOT COPIED: the forward pass (hslm_forward, hslm_block_forward, the sacred ; attention, the FFN, the VSA stream and its gate), the backward pass (hslm_backward: the exact ; gradient of the last position's cross-entropy, embedding included), cross_entropy, the TWN ; quantizer (nn/ternary_ste quantize_twn) and the checkpoint header and layout (nn/hslm_data ; write_header, parse_header, the section offsets). They are reached through `use`. -; ; THE MODEL: ctx 27, vocab 729, embed 243, ff 729, 3 blocks, 3 heads of 81. Every ternary matrix ; (Q, K, V, O, up, down per block, and the output head: 19 matrices, 1771470 trits) has an f64 ; shadow; after every update the shadow is requantized with TWN (delta = 0.7 mean |w|, alpha = mean ; |w| over the kept entries), one alpha per matrix. The float parameters (embedding table 729 x 243, ; RMSNorm gains, FFN biases, output bias) are trained directly. NPARAM = 2130138 f64 in one flat ; vector, 30 segments. -; ; THE STEP: 66 windows of 27 tokens, starts drawn by a xorshift64 PRNG (pinned seed); per window, ; 27 forward positions, cross-entropy of the last position against the 28th token, hslm_backward. ; Gradients are averaged over the 66 windows, clipped to global L2 norm 1.0, and handed to LAMB @@ -24,7 +21,6 @@ ; none on gains and biases, trust ratio ||w|| / ||r|| per segment, clamped to 10, 1 when either norm ; is 0 -- trinity-training autograd.zig lambStepSlice). Learning rate: 1e-3, linear warmup over ; 2000 updates, then cosine down to 1e-6 at TOTAL updates. -; ; THE FILES (all little endian): ; /tmp/hslm/train.tok, valid.tok "HTOK", u32 version 1, u32 N, N u16 ids below 729 (input) ; /tmp/hslm/train.cfg optional, key=value lines of integers: steps, total, warmup, @@ -41,12 +37,10 @@ ; Both files are read back after they are written and compared with what was meant to be on disk; ; any difference makes main() fail (assert, a nonzero exit). resume= in train.cfg restarts ; from ckpt_.state and continues bit-identically (same windows, same updates). -; ; EVERY EVAL_EVERY STEPS one line goes to stdout: step, the step's minibatch loss and its exp, the ; mean minibatch loss since the last line, the loss on EVAL_WINDOWS fixed held-out windows (evenly ; spaced starts in valid.tok, last position) and its exp, tokens/s (66 x 27 tokens per step over the ; training time only), s/step, the learning rate and the gradient norm before clipping. -; ; SPEED, HONESTLY: speed=0 runs the reference kernels of nn/hslm.t27 as they are. speed=1 replaces ; four kernels by local ones with the same results: ternary_matmul as a branch-free axpy over the ; rows of W (bit-identical: same summation order per output, x * 0.0 adds +0.0 to a sum that is never @@ -55,7 +49,6 @@ ; whose result nobody reads: the head at positions 0..25 and, in the last block at those positions, ; everything but the key and value that the last position attends to. The tests check that speed 1 ; and speed 2 produce the same logits as speed 0 to the bit. The model is not reduced. -; ; WHAT DIFFERS FROM THE ZIG TRAINER (trinity-training @ 3e9648c, run R5): f64, not f32; LAMB, not ; AdamW; one global clip, not per tensor; linear warmup + cosine to 1e-6, not the phi schedule; the ; embedding table and RMSNorm gains are trained (the trainer froze the table); the gradient is the From 47aff93dcfbdf7399dd05e03ae88f3629682bc3b Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Sun, 18 Oct 2026 17:47:29 +0700 Subject: [PATCH 3/5] spec(ml): train_hslm -- no duplicated bodies: read float params from the flat vector, xorshift guards zero (Refs #10372) Co-Authored-By: Claude Opus 5.5 --- specs/ml/train_hslm.t27 | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/specs/ml/train_hslm.t27 b/specs/ml/train_hslm.t27 index 336dd54f45..88080e4785 100644 --- a/specs/ml/train_hslm.t27 +++ b/specs/ml/train_hslm.t27 @@ -574,10 +574,11 @@ fn tr_grad_of(s: usize) -> []f64 { return g_embed[0..ef]; } -fn tr_copy(src: []const f64, dst: []f64) -> void { +/// dst.len float parameters of the flat vector, from element off, into dst +fn tr_from_flat(off: usize, dst: []f64) -> void { var i : usize = 0; - while (i < src.len) { - dst[i] = src[i]; + while (i < dst.len) { + dst[i] = p_val[off + i]; i = i + 1; } } @@ -602,17 +603,17 @@ fn tr_scatter(s: usize) -> void { } else if (kind == K_DOWN) { bw.alpha_down = quantize_twn(src, &bw.w_down); } else if (kind == K_GAMMA) { - tr_copy(src, &bw.norm1_gamma); + tr_from_flat(off, &bw.norm1_gamma); } else if (kind == K_BUP) { - tr_copy(src, &bw.b_up); + tr_from_flat(off, &bw.b_up); } else if (kind == K_BDOWN) { - tr_copy(src, &bw.b_down); + tr_from_flat(off, &bw.b_down); } else if (kind == K_OUT) { model.alpha_out = quantize_twn(src, &model.w_out); } else if (kind == K_BOUT) { - tr_copy(src, &model.b_out); + tr_from_flat(off, &model.b_out); } else { - tr_copy(src, &model.embed); + tr_from_flat(off, &model.embed); } } @@ -624,8 +625,12 @@ fn tr_scatter_all() -> void { } } +/// One xorshift64 step; a zero state (a fixed point of xorshift) restarts at 1 fn tr_xs(seed: u64) -> u64 { var x : u64 = seed; + if (x == 0) { + x = 1; + } x = x ^ (x << 13); x = x ^ (x >> 7); x = x ^ (x << 17); @@ -1849,7 +1854,11 @@ fn tr_speed_grad_reldiff() -> f64 { speed_level = 0; zero_weight_gradients(); const l0 : f64 = tr_window_grad(train_tok[0..ntrain], 700); - tr_copy(&win_gin, &ref_gin); + var c : usize = 0; + while (c < EMBED_DIM) { + ref_gin[c] = win_gin[c]; + c = c + 1; + } speed_level = 2; zero_weight_gradients(); const l2 : f64 = tr_window_grad(train_tok[0..ntrain], 700); From a6d8af0f4ff81b5569bcc0b461c7adcb4ad5eb80 Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Sun, 18 Oct 2026 20:05:51 +0700 Subject: [PATCH 4/5] spec(ml): eval_hslm -- held-out loss of an HSLM checkpoint, last position and full sequence (Refs #10372) main() reads /tmp/hslm/eval.cfg (step, windows, stride, check), loads ckpt_.state and valid.tok through train_hslm's own readers, and scores non-overlapping 28-token windows two ways: the last position only (the trainer's valid_loss) and every position (27 terms per window). It prints and writes eval step windows positions

last_loss .. last_ppl .. full_loss .. full_ppl .. to /tmp/hslm/eval_.txt. train_hslm.t27 is reached through use and is not touched. 10 tests, sealed. Co-Authored-By: Claude Sonnet 5.5 --- .trinity/seals/ml_eval_hslm.json | 37 ++ specs/ml/eval_hslm.t27 | 719 +++++++++++++++++++++++++++++++ 2 files changed, 756 insertions(+) create mode 100644 .trinity/seals/ml_eval_hslm.json create mode 100644 specs/ml/eval_hslm.t27 diff --git a/.trinity/seals/ml_eval_hslm.json b/.trinity/seals/ml_eval_hslm.json new file mode 100644 index 0000000000..2fac724feb --- /dev/null +++ b/.trinity/seals/ml_eval_hslm.json @@ -0,0 +1,37 @@ +{ + "built_by": "t27c-bootstrap@0.5.2+fd332dfd5", + "closure": [ + { + "spec": "specs/base/types.t27", + "spec_hash": "sha256:3c2fb2a2edf20e90bc76798fb5d12690b2f60fc0f5bbf2f6052936e22d6e62e2" + }, + { + "spec": "specs/ml/train_hslm.t27", + "spec_hash": "sha256:03fcf23d2966542b6d572676ec373afa136b35a2062fb57eb850f7324c7c4147" + } + ], + "config": "gen=c,rust,verilog,zig;tests=zig", + "gen_hash_c": "sha256:e91abbab32e18fcfa9e21d34d008f03403bc89cbd0854b67bb8abee8b9837fc5", + "gen_hash_rust": "sha256:b6de0a736dcf8e345589be6023120266c76f7027b0228bad13acbf539ac919ae", + "gen_hash_verilog": "sha256:dcd0ce71c7ad6490a0e832129ee77cfbecd5deabbecdd9ee622e898663341a3c", + "gen_hash_zig": "sha256:f4985e4700d6a5df8adb59b772c04f664429adf19fb3ee05abd0b0cbf16d2c0b", + "module": "eval_hslm", + "ring": 12, + "seal_schema": 2, + "sealed_at": "2026-10-18T13:05:14Z", + "sealed_by": "t27c-bootstrap@0.5.2", + "spec_hash": "sha256:09c6729a1509a6ec95616fd50fccf4eb6a3d65dfd22c9aec359e6079a52acec1", + "spec_path": "specs/ml/eval_hslm.t27", + "tests": { + "failed": 0, + "failing": [], + "forced": false, + "passed": 10, + "total": 10, + "vacuous": 0 + }, + "toolchain": { + "test_runner": "sha256:1982ae948e24d00c5d221ccec12dfce1eb01204b53e37b0f857f93b16d64629d", + "zig": "0.16.0" + } +} \ No newline at end of file diff --git a/specs/ml/eval_hslm.t27 b/specs/ml/eval_hslm.t27 new file mode 100644 index 0000000000..918afdca65 --- /dev/null +++ b/specs/ml/eval_hslm.t27 @@ -0,0 +1,719 @@ +// SPDX-License-Identifier: Apache-2.0 +; specs/ml/eval_hslm.t27 -- main() measures a trained HSLM checkpoint on held-out tokens, two ways. +; Issue: gHashTag/t27#10372 (epic #10024, own_model cell 25). +; THE TWO MEASURES, on the same windows of /tmp/hslm/valid.tok (28 tokens each: 27 inputs, then +; the targets are tokens 1..27): +; last_loss the cross-entropy of the LAST position only (input token 26, target token 27), the +; number train_hslm prints as valid_loss and run R5 reported as PPL 2.96; +; full_loss the cross-entropy at EVERY position (27 terms per window, the mean over all of them). +; *_ppl exp of the loss. positions = 27 x windows. +; Position 26's term is the same number in both: the last position's loss is one of the 27 terms of +; the full one, so full_loss measures what the model does on the 26 positions the trainer never +; scores, and last_loss is exactly the trainer's yardstick. +; WHAT IS REUSED, NOT COPIED: the forward pass (hslm_forward and everything under it), +; cross_entropy, the checkpoint reader (tr_load_state reads ckpt_.state: the .bin of MSLH +; holds only the ternary shadows and biases, the state holds every parameter), the HTOK reader +; (tr_read_tok), the held-out loss of the trainer (tr_valid_loss, run as the cross-check), and the +; line formatters and file helpers of train_hslm. They are reached through `use ml::train_hslm`; +; nothing in train_hslm.t27 is changed by this module and no training path moves. +; WHAT IS LOCAL: the shape consts and the module-level state the spliced code binds to (use_resolve +; splices no `var`, see train_hslm.t27), two kernels (ternary_matmul, trit_int) because the name is +; declared by both nn/hslm and train_hslm and a name two imports declare is not resolved, and the +; evaluation itself. +; THE CONFIG /tmp/hslm/eval.cfg: key=value lines of integers, unknown keys ignored, a missing file +; or key means the default: +; step the checkpoint, /tmp/hslm/ckpt_.state (default 0) +; windows how many windows (default 243, the trainer's EVAL_WINDOWS) +; stride 0 (default): windows evenly spaced over the file exactly as tr_valid_loss spaces them, +; start k * ((N - 28) / windows); otherwise window k starts at k * stride, and a stride +; below 28 is raised to 28 so that windows never overlap +; check 1: after the evaluation also run tr_valid_loss(windows) and print both numbers +; (only meaningful with stride 0, otherwise the check line says it was skipped) +; A window that does not fit in the file (fewer than 28 tokens left at its start) is refused and +; not counted. +; THE OUTPUT, to stdout and to /tmp/hslm/eval_.txt (the file holds this one line): +; eval step windows positions

last_loss last_ppl full_loss full_ppl +; main() fails (assert, a nonzero exit) when the token file or the checkpoint cannot be read, no +; window fits, a loss is NaN, or the result file does not read back as written. +; phi^2 + 1/phi^2 = 3 | TRINITY + +module eval_hslm; + +use base::types; +use ml::train_hslm; + +; THE SHAPE. nn/hslm and nn/hslm_data both declare these; this module declares them once, with the +; same type and value as both (train_hslm.t27 does the same). +const CONTEXT_LEN : usize = 27; +const VOCAB_SIZE : usize = 729; +const EMBED_DIM : usize = 243; +const FF_DIM : usize = 729; +const NUM_BLOCKS : usize = 3; +const HEADER_BYTES : usize = 16; + +const EV_WINDOW : usize = 28; +const EV_DEFAULT_WINDOWS : u64 = 243; + +; Exit codes of ev_main +const EV_RC_OK : i32 = 0; +const EV_RC_DATA : i32 = 10; +const EV_RC_WRITE : i32 = 20; +const EV_RC_STATE : i32 = 30; +const EV_RC_NAN : i32 = 40; +const EV_RC_NONE : i32 = 50; + +; THE STATE OF nn/hslm.t27 (as in train_hslm.t27: names, types and initial values, nothing else). +var model : HSLMWeights = undefined; +var kv_k : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var kv_v : [NUM_BLOCKS * CONTEXT_LEN * EMBED_DIM]f64 = undefined; +var act_token : usize = 0; +var act_pos : [NUM_BLOCKS]usize = undefined; +var act_x : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_n1 : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_rms1 : [NUM_BLOCKS]f64 = undefined; +var act_q : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_attn : [NUM_BLOCKS * NUM_HEADS * CONTEXT_LEN]f64 = undefined; +var act_ctx : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_r : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_pre : [NUM_BLOCKS * FF_DIM]f64 = undefined; +var act_h : [NUM_BLOCKS * FF_DIM]f64 = undefined; +var act_out : [NUM_BLOCKS * EMBED_DIM]f64 = undefined; +var act_hn : [EMBED_DIM]f64 = undefined; +var act_hrms : f64 = 1.0; +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; +var bw_out : [EMBED_DIM]f64 = undefined; +var bw_h : [FF_DIM]f64 = undefined; +var bw_pre : [FF_DIM]f64 = undefined; +var bw_sf : [FF_DIM]f64 = undefined; +var bw_se : [EMBED_DIM]f64 = undefined; +var bw_n : [EMBED_DIM]f64 = undefined; +var bw_tmp : [EMBED_DIM]f64 = undefined; +var bw_r : [EMBED_DIM]f64 = undefined; +var bw_ctx : [EMBED_DIM]f64 = undefined; +var bw_q : [EMBED_DIM]f64 = undefined; +var bw_k : [EMBED_DIM]f64 = undefined; +var bw_v : [EMBED_DIM]f64 = undefined; +var bw_da : [CONTEXT_LEN]f64 = undefined; +var bw_x : [EMBED_DIM]f64 = undefined; +var bw_dz : [VOCAB_SIZE]f64 = undefined; +var nudge_which : u8 = MAT_NONE; +var nudge_block : usize = 0; +var nudge_index : usize = 0; +var nudge_delta : f64 = 0.0; +var vsa_on : bool = true; +var vsa_in : [NUM_BLOCKS * CONTEXT_LEN * VSA_DIM]Trit = undefined; +var vsa_out : [NUM_BLOCKS * VSA_DIM]Trit = undefined; +var vsa_ctx : [VSA_DIM]Trit = undefined; +var vsa_reasoned : [VSA_DIM]Trit = undefined; +var vsa_temp : [VSA_DIM]Trit = undefined; +var vsa_temp2 : [VSA_DIM]Trit = undefined; +var vsa_sim : [CONTEXT_LEN]f64 = undefined; +var vsa_accum : [VSA_DIM]i64 = undefined; +var vsa_max_sim : [NUM_BLOCKS]f64 = undefined; +var vsa_fired : [NUM_BLOCKS]bool = undefined; +var gate_total : [NUM_BLOCKS]usize = undefined; +var gate_conscious : [NUM_BLOCKS]usize = undefined; +var gate_ema : [NUM_BLOCKS]f64 = undefined; + +; THE STATE OF train_hslm.t27 that the spliced reader, forward driver and formatters bind to. +var p_val : [NPARAM]f64 = undefined; +var g_val : [NPARAM]f64 = undefined; +var m_val : [NPARAM]f64 = undefined; +var v_val : [NPARAM]f64 = undefined; +var seg_off : [NSEG]usize = undefined; +var seg_len : [NSEG]usize = undefined; +var seg_kind : [NSEG]u8 = undefined; +var seg_block : [NSEG]usize = undefined; +var beta1_pow : f64 = 1.0; +var beta2_pow : f64 = 1.0; +var rng_state : u64 = 1; +var speed_level : u8 = 2; +var last_loss : f64 = 0.0; +var io_fail : u32 = 0; +var valid_tok : [VALID_CAP]u16 = undefined; +var nvalid : usize = 0; +var win_logits : [VOCAB_SIZE]f64 = undefined; +var kv_n1 : [EMBED_DIM]f64 = undefined; +var io_chunk : [IO_CHUNK]u8 = undefined; +var io_back : [IO_CHUNK]u8 = undefined; +var io_hdr : [STATE_HEADER]u8 = undefined; +var io_hdr_back : [STATE_HEADER]u8 = undefined; +var path_buf : [96]u8 = undefined; +var cfg_buf : [4096]u8 = undefined; +var ln_buf : [512]u8 = undefined; +var ln_n : usize = 0; + +; THIS MODULE'S STATE +var ev_cfg_step : u64 = 0; +var ev_cfg_windows : u64 = EV_DEFAULT_WINDOWS; +var ev_cfg_stride : u64 = 0; +var ev_cfg_check : u64 = 0; +var ev_term : [CONTEXT_LEN]f64 = undefined; +var ev_sum_last : f64 = 0.0; +var ev_sum_full : f64 = 0.0; +var ev_nwin : u64 = 0; +var ev_npos : u64 = 0; + +; --------------------------------------------------------------------------------------------- +; 1. Kernels. The two names are nn/hslm.t27's and train_hslm.t27's; the local declaration wins. +; These are train_hslm.t27's speed-1 forms, which it tests to give the reference logits to the bit. +; --------------------------------------------------------------------------------------------- + +/// output[i] = sum_j input[j] * w[j * out_dim + i], as an axpy over the rows of W (per output the +/// terms are added in the order j = 0, 1, ..., the same bits as nn/hslm.t27's dot per output) +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) { + output[i] = 0.0; + i = i + 1; + } + var j : usize = 0; + while (j < in_dim) { + const xj : f64 = input[j]; + if (xj != 0.0) { + const row : usize = j * out_dim; + var c : usize = 0; + while (c < out_dim) { + const wf : f64 = @floatFromInt(@intFromEnum(weights[row + c])); + output[c] = output[c] + xj * wf; + c = c + 1; + } + } + j = j + 1; + } +} + +/// -1 / 0 / +1 of a trit: the enum's own value +fn trit_int(t: Trit) -> i64 { + return @as(i64, @intFromEnum(t)); +} + +; --------------------------------------------------------------------------------------------- +; 2. The config +; --------------------------------------------------------------------------------------------- + +fn ev_cfg_reset() -> void { + ev_cfg_step = 0; + ev_cfg_windows = EV_DEFAULT_WINDOWS; + ev_cfg_stride = 0; + ev_cfg_check = 0; +} + +fn ev_cfg_set(off: usize, len: usize, value: u64) -> void { + if (tr_key_is(off, len, "step")) { + ev_cfg_step = value; + } else if (tr_key_is(off, len, "windows")) { + ev_cfg_windows = value; + } else if (tr_key_is(off, len, "stride")) { + ev_cfg_stride = value; + } else if (tr_key_is(off, len, "check")) { + ev_cfg_check = value; + } +} + +/// The key=value lines in cfg_buf[0..len]; unknown keys are ignored +fn ev_parse_cfg(len: usize) -> void { + var i : usize = 0; + while (i < len) { + const key_off : usize = i; + while (i < len and cfg_buf[i] != 61 and cfg_buf[i] != 10) { + i = i + 1; + } + if (i < len and cfg_buf[i] == 61) { + const key_len : usize = i - key_off; + i = i + 1; + var value : u64 = 0; + while (i < len and cfg_buf[i] >= 48 and cfg_buf[i] <= 57) { + value = value * 10 + @as(u64, cfg_buf[i] - 48); + i = i + 1; + } + ev_cfg_set(key_off, key_len, value); + } + while (i < len and cfg_buf[i] != 10) { + i = i + 1; + } + i = i + 1; + } +} + +/// /tmp/hslm/eval.cfg, if the file exists +fn ev_read_cfg() -> void { + const fd : i32 = host_open("/tmp/hslm/eval.cfg", 0); + if (fd < 0) { + return; + } + const got : i64 = tr_read_full(fd, &cfg_buf); + tr_note(host_close(fd)); + if (got <= 0) { + return; + } + ev_parse_cfg(@intCast(got)); +} + +; --------------------------------------------------------------------------------------------- +; 3. The measure +; --------------------------------------------------------------------------------------------- + +fn ev_acc_reset() -> void { + ev_sum_last = 0.0; + ev_sum_full = 0.0; + ev_nwin = 0; + ev_npos = 0; +} + +/// The term of position p (target already chosen): kept in ev_term, added to the full sum; the +/// last position's term also goes to the last-position sum and closes the window +fn ev_add_term(p: usize, logits: []const f64, target: usize) -> void { + const ce : f64 = cross_entropy(logits, target); + ev_term[p] = ce; + ev_sum_full = ev_sum_full + ce; + ev_npos = ev_npos + 1; + if (p + 1 == CONTEXT_LEN) { + ev_sum_last = ev_sum_last + ce; + ev_nwin = ev_nwin + 1; + } +} + +fn ev_mean_last() -> f64 { + return ev_sum_last / @as(f64, @floatFromInt(ev_nwin)); +} + +fn ev_mean_full() -> f64 { + return ev_sum_full / @as(f64, @floatFromInt(ev_npos)); +} + +/// One window of 28 tokens starting at `start`: forward at every position in order (each attends +/// to the cache of the earlier ones), the term at position p is scored against token p + 1. A +/// window that does not fit is refused: false, nothing counted. +fn ev_try_window(tokens: []const u16, start: usize) -> bool { + if (start + EV_WINDOW > tokens.len) { + return false; + } + var p : usize = 0; + while (p < CONTEXT_LEN) { + const tok : usize = @intCast(tokens[start + p]); + hslm_forward(tok, p, &win_logits); + const target : usize = @intCast(tokens[start + p + 1]); + ev_add_term(p, &win_logits, target); + p = p + 1; + } + return true; +} + +/// Start spacing: the config's stride (raised to 28 so windows never overlap), or 0 = evenly +/// spaced over n tokens exactly as tr_valid_loss spaces `count` windows +fn ev_stride_for(count: usize, n: usize) -> usize { + if (ev_cfg_stride > 0) { + if (ev_cfg_stride < @as(u64, EV_WINDOW)) { + return EV_WINDOW; + } + return @intCast(ev_cfg_stride); + } + if (n <= CONTEXT_LEN + 1 or count == 0) { + return 1; + } + var stride : usize = (n - CONTEXT_LEN - 1) / count; + if (stride == 0) { + stride = 1; + } + return stride; +} + +/// Windows k = 0..count-1 at k * stride; the number that fit +fn ev_run(tokens: []const u16, count: usize, stride: usize) -> usize { + var used : usize = 0; + var k : usize = 0; + while (k < count) { + if (ev_try_window(tokens, k * stride)) { + used = used + 1; + } + k = k + 1; + } + return used; +} + +; --------------------------------------------------------------------------------------------- +; 4. Output, data, main +; --------------------------------------------------------------------------------------------- + +/// /tmp/hslm/eval_.txt into path_buf (NUL ended); clobbers ln_buf +fn ev_path(step: u64) -> void { + ln_reset(); + ln_str("/tmp/hslm/eval_"); + ln_u64(step); + ln_str(".txt"); + var i : usize = 0; + while (i < ln_n) { + path_buf[i] = ln_buf[i]; + i = i + 1; + } + path_buf[ln_n] = 0; + ln_n = 0; +} + +/// The result line, without its newline, in ln_buf +fn ev_compose_line(step: u64) -> void { + ln_reset(); + ln_str("eval step "); + ln_u64(step); + ln_str(" windows "); + ln_u64(ev_nwin); + ln_str(" positions "); + ln_u64(ev_npos); + ln_str(" last_loss "); + ln_f64(ev_mean_last(), 4); + ln_str(" last_ppl "); + ln_f64(exp(ev_mean_last()), 2); + ln_str(" full_loss "); + ln_f64(ev_mean_full(), 4); + ln_str(" full_ppl "); + ln_f64(exp(ev_mean_full()), 2); +} + +/// ln_buf + newline to stdout and to the file at path_buf, then the file is read back and compared +fn ev_emit() -> i32 { + ln_byte(10); + const len : usize = ln_n; + tr_note(tr_write_all(1, ln_buf[0..len])); + const fd : i32 = host_open(@ptrCast(&path_buf), 1); + if (fd < 0) { + return EV_RC_WRITE; + } + const wrote : i32 = tr_write_all(fd, ln_buf[0..len]); + tr_note(host_close(fd)); + if (wrote != 0) { + return EV_RC_WRITE; + } + const rd : i32 = host_open(@ptrCast(&path_buf), 0); + if (rd < 0) { + return EV_RC_WRITE; + } + const got : i64 = tr_read_full(rd, io_back[0..512]); + tr_note(host_close(rd)); + if (got != @as(i64, @intCast(len))) { + return EV_RC_WRITE; + } + var i : usize = 0; + while (i < len) { + if (io_back[i] != ln_buf[i]) { + return EV_RC_WRITE; + } + i = i + 1; + } + return EV_RC_OK; +} + +/// valid.tok into valid_tok / nvalid +fn ev_load_valid() -> i32 { + const nv : i64 = tr_read_tok(1, &valid_tok); + if (nv < @as(i64, @intCast(EV_WINDOW))) { + return EV_RC_DATA; + } + nvalid = @intCast(nv); + return EV_RC_OK; +} + +fn ev_say(text: []const u8) -> void { + ln_reset(); + ln_str(text); + ln_flush(); +} + +/// check=1: the trainer's own tr_valid_loss on the same windows, next to ours, and whether the +/// two f64 are the same bits +fn ev_check_line(count: usize) -> void { + ln_reset(); + if (ev_cfg_stride != 0) { + ln_str("check skipped: stride is set, tr_valid_loss spaces its windows evenly"); + ln_flush(); + return; + } + const theirs : f64 = tr_valid_loss(count); + const ours : f64 = ev_mean_last(); + const a : u64 = @bitCast(theirs); + const b : u64 = @bitCast(ours); + ln_str("check tr_valid_loss "); + ln_f64(theirs, 6); + ln_str(" last_loss "); + ln_f64(ours, 6); + ln_str(" same_bits "); + if (a == b) { + ln_str("1"); + } else { + ln_str("0"); + } + ln_flush(); +} + +fn ev_main() -> i32 { + ev_cfg_reset(); + ev_read_cfg(); + const rc_data : i32 = ev_load_valid(); + if (rc_data != EV_RC_OK) { + ev_say("eval_hslm: /tmp/hslm/valid.tok missing or bad"); + return rc_data; + } + tr_init(1); + if (tr_load_state(ev_cfg_step) != RC_OK) { + ev_say("eval_hslm: /tmp/hslm/ckpt_.state missing or bad (step from eval.cfg)"); + return EV_RC_STATE; + } + ev_acc_reset(); + var count : usize = @intCast(ev_cfg_windows); + if (count == 0) { + count = 1; + } + const used : usize = ev_run(valid_tok[0..nvalid], count, ev_stride_for(count, nvalid)); + if (used == 0) { + ev_say("eval_hslm: no window fits in valid.tok"); + return EV_RC_NONE; + } + const lf : f64 = ev_mean_last(); + const ff : f64 = ev_mean_full(); + if (lf != lf or ff != ff) { + ev_say("eval_hslm: a loss is NaN"); + return EV_RC_NAN; + } + ev_path(ev_cfg_step); + ev_compose_line(ev_cfg_step); + const rc : i32 = ev_emit(); + if (ev_cfg_check == 1) { + ev_check_line(count); + } + return rc; +} + +/// The program `t27c compile --backend ml` builds +pub fn main() -> void { + assert ev_main() == EV_RC_OK; +} + +; --------------------------------------------------------------------------------------------- +; 5. Test helpers (each returns the number its test checks) and the tests +; --------------------------------------------------------------------------------------------- + +/// A small synthetic held-out set in valid_tok and a fresh random model +fn ev_tiny_setup(seed: u64) -> void { + tr_synth(valid_tok[0..2048], 98); + nvalid = 2048; + tr_init(seed); + ev_cfg_reset(); + ev_acc_reset(); +} + +/// All-zero logits for every position of `windows` windows, through the real accumulator +fn ev_uniform_feed(windows: usize) -> void { + ev_acc_reset(); + var i : usize = 0; + while (i < VOCAB_SIZE) { + win_logits[i] = 0.0; + i = i + 1; + } + var w : usize = 0; + while (w < windows) { + var p : usize = 0; + while (p < CONTEXT_LEN) { + ev_add_term(p, &win_logits, (7 * w + 11 * p) % VOCAB_SIZE); + p = p + 1; + } + w = w + 1; + } +} + +fn ev_uniform_last_gap() -> f64 { + ev_uniform_feed(3); + return abs(ev_mean_last() - log(729.0)); +} + +fn ev_uniform_full_gap() -> f64 { + ev_uniform_feed(3); + return abs(ev_mean_full() - log(729.0)); +} + +fn ev_uniform_positions() -> u64 { + ev_uniform_feed(3); + return ev_npos; +} + +fn ev_uniform_windows() -> u64 { + ev_uniform_feed(3); + return ev_nwin; +} + +/// |term at position 26 of the full measure - the trainer's last-position loss| on one window +fn ev_p26_gap() -> f64 { + ev_tiny_setup(5); + if (!ev_try_window(valid_tok[0..nvalid], 300)) { + return 1.0; + } + const theirs : f64 = tr_window_loss(valid_tok[0..nvalid], 300); + return abs(ev_term[26] - theirs); +} + +/// |mean of the 27 terms - ev_mean_full| after one window, and the first term differs from the last +fn ev_full_is_the_mean_gap() -> f64 { + ev_tiny_setup(5); + if (!ev_try_window(valid_tok[0..nvalid], 300)) { + return 1.0; + } + var s : f64 = 0.0; + var p : usize = 0; + while (p < CONTEXT_LEN) { + s = s + ev_term[p]; + p = p + 1; + } + return abs(s / 27.0 - ev_mean_full()); +} + +fn ev_first_last_differ() -> bool { + ev_tiny_setup(5); + if (!ev_try_window(valid_tok[0..nvalid], 300)) { + return false; + } + return ev_term[0] != ev_term[26]; +} + +/// |our last-position mean over 7 evenly spaced windows - tr_valid_loss(7)| on a tiny model: the +/// trainer's own yardstick and ours agree to the bit (the trainer runs its speed-2 shortcut here) +fn ev_vs_train_gap() -> f64 { + ev_tiny_setup(5); + const used : usize = ev_run(valid_tok[0..nvalid], 7, ev_stride_for(7, nvalid)); + if (used != 7) { + return 1.0; + } + return abs(ev_mean_last() - tr_valid_loss(7)); +} + +/// A window of 27 tokens (one short) is refused and counts nothing; one of 28 is taken +fn ev_short_refused() -> u64 { + ev_tiny_setup(5); + var r : u64 = 0; + if (ev_try_window(valid_tok[0..27], 0)) { + r = r + 1; + } + if (ev_try_window(valid_tok[0..28], 1)) { + r = r + 10; + } + if (ev_try_window(valid_tok[0..0], 0)) { + r = r + 100; + } + if (ev_nwin != 0 or ev_npos != 0) { + r = r + 1000; + } + return r; +} + +fn ev_exact_fit_taken() -> u64 { + ev_tiny_setup(5); + if (!ev_try_window(valid_tok[0..28], 0)) { + return 0; + } + return ev_npos * 100 + ev_nwin; +} + +/// Parse `text` as an eval.cfg and read one value back: which = 0 step, 1 windows, 2 stride, 3 check +fn ev_cfg_value(text: []const u8, which: u8) -> u64 { + ev_cfg_reset(); + var i : usize = 0; + while (i < text.len) { + cfg_buf[i] = text[i]; + i = i + 1; + } + ev_parse_cfg(text.len); + if (which == 0) { + return ev_cfg_step; + } + if (which == 1) { + return ev_cfg_windows; + } + if (which == 2) { + return ev_cfg_stride; + } + return ev_cfg_check; +} + +/// ev_stride_for with the config's stride set to `cfg_stride` +fn ev_stride_with(cfg_stride: u64, count: usize, n: usize) -> usize { + ev_cfg_reset(); + ev_cfg_stride = cfg_stride; + return ev_stride_for(count, n); +} + +test uniform_logits_give_ln_729_for_the_last_position_measure + given gap = ev_uniform_last_gap() + and wins = ev_uniform_windows() + then gap < 1e-12 + and wins == 3 + +test uniform_logits_give_ln_729_for_the_full_measure + given gap = ev_uniform_full_gap() + and pos = ev_uniform_positions() + then gap < 1e-12 + and pos == 81 + +test position_26_of_the_full_measure_is_the_last_position_loss + given gap = ev_p26_gap() + then gap == 0.0 + +test the_full_loss_is_the_mean_of_27_terms_and_not_only_the_last + given gap = ev_full_is_the_mean_gap() + and differ = ev_first_last_differ() + then gap < 1e-12 + and differ + +test last_loss_equals_the_trainers_valid_loss_on_the_same_windows + given gap = ev_vs_train_gap() + then gap == 0.0 + +test a_window_shorter_than_28_tokens_is_refused + given r = ev_short_refused() + then r == 0 + +test a_window_of_exactly_28_tokens_is_taken + given r = ev_exact_fit_taken() + then r == 2701 + +test config_defaults_when_nothing_is_set + given step = ev_cfg_value("", 0) + and windows = ev_cfg_value("", 1) + and stride = ev_cfg_value("", 2) + and check = ev_cfg_value("", 3) + then step == 0 + and windows == 243 + and stride == 0 + and check == 0 + +test config_reads_its_four_keys_and_ignores_the_rest + given step = ev_cfg_value("step=2000\nwindows=50\nstride=100\ncheck=1\n", 0) + and windows = ev_cfg_value("step=2000\nwindows=50\nstride=100\ncheck=1\n", 1) + and stride = ev_cfg_value("step=2000\nwindows=50\nstride=100\ncheck=1\n", 2) + and check = ev_cfg_value("step=2000\nwindows=50\nstride=100\ncheck=1\n", 3) + and other = ev_cfg_value("speed=1\nsteps=9\nwindows=7", 1) + and keep = ev_cfg_value("speed=1\nsteps=9\nwindows=7", 0) + then step == 2000 + and windows == 50 + and stride == 100 + and check == 1 + and other == 7 + and keep == 0 + +test stride_zero_spaces_windows_as_the_trainer_does_and_a_short_stride_is_raised + given even = ev_stride_with(0, 243, 7689354) + and small = ev_stride_with(5, 243, 7689354) + and exact = ev_stride_with(28, 243, 7689354) + and wide = ev_stride_with(1000, 243, 7689354) + then even == 31643 + and small == 28 + and exact == 28 + and wide == 1000 From 2bc7da33d455ed7767cf69a01db723bce1039431 Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Mon, 19 Oct 2026 10:43:09 +0700 Subject: [PATCH 5/5] ml: ngram_hslm -- unigram, bigram and trigram baselines on HSLM's held-out windows; reseal eval_hslm specs/ml/ngram_hslm.t27 counts unigrams, bigrams and trigrams of /tmp/hslm/train.tok (interpolated absolute discounting, d 0.75) and scores /tmp/hslm/valid.tok two ways: the whole stream, and the same 19683 windows of 28 tokens that eval_hslm measures (last position and every position). 8/8 tests, 0 vacuous, on the Railway lab. Measured on the TinyStories-V2 valid split (7,689,354 tokens, 22,502,601 bytes), 19683 windows, 531,441 positions: unigram full ppl 333.98, stream bpb 2.8650 bigram full ppl 33.92, stream bpb 1.7380 trigram full ppl 12.52, stream bpb 1.2281 HSLM at step 8000 (eval_hslm, same windows): full ppl 25.41. The trigram is the better model on this data; HSLM beats the bigram. The eval_hslm seal is renewed against master's train_hslm (fa8490d8), which it reaches through `use`; 10/10, 0 vacuous. Refs #10372 Co-Authored-By: Claude Opus 5.5 --- .trinity/seals/ml_eval_hslm.json | 6 +- .trinity/seals/ml_ngram_hslm.json | 37 ++ specs/ml/ngram_hslm.t27 | 659 ++++++++++++++++++++++++++++++ 3 files changed, 699 insertions(+), 3 deletions(-) create mode 100644 .trinity/seals/ml_ngram_hslm.json create mode 100644 specs/ml/ngram_hslm.t27 diff --git a/.trinity/seals/ml_eval_hslm.json b/.trinity/seals/ml_eval_hslm.json index 2fac724feb..63c0f1320d 100644 --- a/.trinity/seals/ml_eval_hslm.json +++ b/.trinity/seals/ml_eval_hslm.json @@ -7,7 +7,7 @@ }, { "spec": "specs/ml/train_hslm.t27", - "spec_hash": "sha256:03fcf23d2966542b6d572676ec373afa136b35a2062fb57eb850f7324c7c4147" + "spec_hash": "sha256:fa8490d815bde2815fe9bfa89cb9badab241f580977e59be52fc5902f6f8bb03" } ], "config": "gen=c,rust,verilog,zig;tests=zig", @@ -18,7 +18,7 @@ "module": "eval_hslm", "ring": 12, "seal_schema": 2, - "sealed_at": "2026-10-18T13:05:14Z", + "sealed_at": "2026-10-19T03:39:21Z", "sealed_by": "t27c-bootstrap@0.5.2", "spec_hash": "sha256:09c6729a1509a6ec95616fd50fccf4eb6a3d65dfd22c9aec359e6079a52acec1", "spec_path": "specs/ml/eval_hslm.t27", @@ -31,7 +31,7 @@ "vacuous": 0 }, "toolchain": { - "test_runner": "sha256:1982ae948e24d00c5d221ccec12dfce1eb01204b53e37b0f857f93b16d64629d", + "test_runner": "sha256:8512b80da533b79289110fca4a65f3336abd7a059ca4b72ea595d5e33ff8ad10", "zig": "0.16.0" } } \ No newline at end of file diff --git a/.trinity/seals/ml_ngram_hslm.json b/.trinity/seals/ml_ngram_hslm.json new file mode 100644 index 0000000000..e5a01aca68 --- /dev/null +++ b/.trinity/seals/ml_ngram_hslm.json @@ -0,0 +1,37 @@ +{ + "built_by": "t27c-bootstrap@0.5.2+fd332dfd5", + "closure": [ + { + "spec": "specs/base/types.t27", + "spec_hash": "sha256:3c2fb2a2edf20e90bc76798fb5d12690b2f60fc0f5bbf2f6052936e22d6e62e2" + }, + { + "spec": "specs/ml/train_hslm.t27", + "spec_hash": "sha256:fa8490d815bde2815fe9bfa89cb9badab241f580977e59be52fc5902f6f8bb03" + } + ], + "config": "gen=c,rust,verilog,zig;tests=zig", + "gen_hash_c": "sha256:8e1f5f6ff5d42ae3972441869bd4c0d59c7133e3c2894a55392c0fa5b8553ddb", + "gen_hash_rust": "sha256:e8a4badf02f69ce4cfdba6dbf7e6f0ebe69d2e5948551ab85c97dfc4e3ab3b01", + "gen_hash_verilog": "sha256:acda1adadfdd0e36466319ce28b64350ac903de36fd1e7fefa07cdfb2d5a87f7", + "gen_hash_zig": "sha256:c1fc022ffe591f9ceed36a20feda6cad6fc81ee1a2b243298a052ac8e764926a", + "module": "ngram_hslm", + "ring": 12, + "seal_schema": 2, + "sealed_at": "2026-10-19T03:39:38Z", + "sealed_by": "t27c-bootstrap@0.5.2", + "spec_hash": "sha256:d3e31459bc0702d1b59357dad2b4d2ef679fbe95fcda2185adb612a2bdb6f6ed", + "spec_path": "specs/ml/ngram_hslm.t27", + "tests": { + "failed": 0, + "failing": [], + "forced": false, + "passed": 8, + "total": 8, + "vacuous": 0 + }, + "toolchain": { + "test_runner": "sha256:8512b80da533b79289110fca4a65f3336abd7a059ca4b72ea595d5e33ff8ad10", + "zig": "0.16.0" + } +} \ No newline at end of file diff --git a/specs/ml/ngram_hslm.t27 b/specs/ml/ngram_hslm.t27 new file mode 100644 index 0000000000..ca580bc0c7 --- /dev/null +++ b/specs/ml/ngram_hslm.t27 @@ -0,0 +1,659 @@ +// SPDX-License-Identifier: Apache-2.0 +; specs/ml/ngram_hslm.t27 -- main() fits unigram, bigram and trigram models on the HSLM training +; tokens and measures them on the held-out tokens, the yardstick a perplexity of the trained HSLM +; is read against. Issue: gHashTag/t27#10372 (epic #10024, own_model cell 25). +; WHY: a perplexity means nothing alone. Its tokenizer fixes the scale (vocab 729, about 2.9 +; bytes per token on TinyStories), so published TinyStories numbers with other tokenizers do not +; compare. What does compare is a count model on the SAME train.tok and valid.tok: if the network +; does not beat a trigram, it learned nothing a table could not. +; THE MODELS, counts over all of /tmp/hslm/train.tok, no parameter tuned on valid.tok: +; P1(w) = (c(w) + 1) / (N + 729) add-one unigram +; P2(w|v) = (max(c(v w) - D, 0) + D n(v .) P1(w)) / c(v .) interpolated absolute +; P3(w|u v) = (max(c(u v w) - D, 0) + D n(u v .) P2(w|v)) / c(u v .) discounting, D = 0.75 +; c(v .) and c(u v .) are context counts, n(v .) and n(u v .) the number of distinct words seen +; after the context. An unseen context falls back to the order below. Each P sums to 1 over w +; (the tests check it). Bigram counts are a dense 729 x 729 table; trigram counts live in an +; open-addressing hash of 2^24 slots, keyed by (u * 729 + v) * 729 + w. +; THE TWO MEASURES on /tmp/hslm/valid.tok: +; stream every token from the second on, predicted from the (up to) two tokens before it +; windows exactly eval_hslm's windows (NG_WINDOWS evenly spaced starts k * ((N - 28) / windows), +; 28 tokens each): position p predicts token p + 1 from the tokens of the window only, so +; position 0 has one token of context (the bigram), later positions two (the trigram). +; full = mean over the 27 positions, last = position 26 only, as eval_hslm prints them. +; bpb stream bits per byte = mean nats / ln 2 * tokens / bytes of /tmp/hslm/valid.txt. +; THE OUTPUT, stdout: one `ngram data` line, then one `ngram order ` line per order. +; main() fails (assert, a nonzero exit) when a token file cannot be read, valid.txt is missing, the +; hash fills, or a loss is NaN. +; WHAT IS REUSED: the HTOK reader (tr_read_tok), the file helpers and the line formatters of +; train_hslm, through `use ml::train_hslm`; the state they bind to is declared here (use_resolve +; splices no `var`). Nothing in train_hslm.t27 changes. +; phi^2 + 1/phi^2 = 3 | TRINITY + +module ngram_hslm; + +use base::types; +use ml::train_hslm; + +const NG_V : usize = 729; +const NG_D : f64 = 0.75; +const NG_HASH_MAX : usize = 16777216; +const NG_WINDOWS : usize = 19683; +const NG_WINDOW : usize = 28; +const NG_CTX : usize = 27; + +const NG_RC_OK : i32 = 0; +const NG_RC_DATA : i32 = 10; +const NG_RC_TEXT : i32 = 20; +const NG_RC_FULL : i32 = 30; +const NG_RC_NAN : i32 = 40; + +; THE STATE train_hslm's spliced helpers bind to +var io_fail : u32 = 0; +var io_hdr : [STATE_HEADER]u8 = undefined; +var io_chunk : [IO_CHUNK]u8 = undefined; +var cfg_buf : [4096]u8 = undefined; +var ln_buf : [512]u8 = undefined; +var ln_n : usize = 0; +var train_tok : [TRAIN_CAP]u16 = undefined; +var valid_tok : [VALID_CAP]u16 = undefined; +var ntrain : usize = 0; +var nvalid : usize = 0; + +; THE COUNTS +var ng_uni : [NG_V]u64 = undefined; +var ng_total : u64 = 0; +var ng_bi : [NG_V * NG_V]u32 = undefined; +var ng_bi_row : [NG_V]u64 = undefined; +var ng_bi_types : [NG_V]u64 = undefined; +var ng_tri_ctx : [NG_V * NG_V]u32 = undefined; +var ng_tri_types : [NG_V * NG_V]u32 = undefined; +var ng_key : [NG_HASH_MAX]u32 = undefined; +var ng_cnt : [NG_HASH_MAX]u32 = undefined; +var ng_hash_size : usize = NG_HASH_MAX; +var ng_hash_used : usize = 0; +var ng_full : bool = false; + +; THE MEASURES: index 0, 1, 2 = order 1, 2, 3 +var ng_s_stream : [3]f64 = undefined; +var ng_s_full : [3]f64 = undefined; +var ng_s_last : [3]f64 = undefined; +var ng_n_stream : u64 = 0; +var ng_n_pos : u64 = 0; +var ng_n_win : u64 = 0; +var ng_bytes : u64 = 0; + +; --------------------------------------------------------------------------------------------- +; 1. Counting +; --------------------------------------------------------------------------------------------- + +/// Zero every count; the hash uses its first `hash_size` slots (2^24 in main, small in the tests) +fn ng_reset(hash_size: usize) -> void { + var i : usize = 0; + while (i < NG_V) { + ng_uni[i] = 0; + ng_bi_row[i] = 0; + ng_bi_types[i] = 0; + i = i + 1; + } + i = 0; + while (i < NG_V * NG_V) { + ng_bi[i] = 0; + ng_tri_ctx[i] = 0; + ng_tri_types[i] = 0; + i = i + 1; + } + ng_hash_size = hash_size; + i = 0; + while (i < hash_size) { + ng_key[i] = 0; + ng_cnt[i] = 0; + i = i + 1; + } + ng_total = 0; + ng_hash_used = 0; + ng_full = false; +} + +fn ng_tri_key(u: usize, v: usize, w: usize) -> u32 { + return @intCast((u * NG_V + v) * NG_V + w); +} + +/// The slot of `key` (stored as key + 1, 0 = empty): its own, or the empty one where it would go +fn ng_slot(key: u32) -> usize { + const k : u64 = @as(u64, key) + 1; + var h : usize = @intCast(((k * 2654435761) >> 16) % @as(u64, @intCast(ng_hash_size))); + var probes : usize = 0; + while (probes < ng_hash_size) { + const s : u32 = ng_key[h]; + if (s == 0 or @as(u64, s) == k) { + return h; + } + h = h + 1; + if (h == ng_hash_size) { + h = 0; + } + probes = probes + 1; + } + return ng_hash_size; +} + +/// c(u v w) += 1; true when the trigram is new +fn ng_tri_add(key: u32) -> bool { + const h : usize = ng_slot(key); + if (h < ng_hash_size and ng_key[h] != 0) { + ng_cnt[h] = ng_cnt[h] + 1; + return false; + } + if (h == ng_hash_size or ng_hash_used * 4 >= ng_hash_size * 3) { + ng_full = true; + return false; + } + ng_key[h] = key + 1; + ng_cnt[h] = 1; + ng_hash_used = ng_hash_used + 1; + return true; +} + +fn ng_tri_count(key: u32) -> u32 { + const h : usize = ng_slot(key); + if (h == ng_hash_size or ng_key[h] == 0) { + return 0; + } + return ng_cnt[h]; +} + +/// All counts of tokens[0..n] +fn ng_fit(tokens: []const u16, n: usize) -> void { + var i : usize = 0; + while (i < n) { + const w : usize = @intCast(tokens[i]); + ng_uni[w] = ng_uni[w] + 1; + ng_total = ng_total + 1; + if (i >= 1) { + const v : usize = @intCast(tokens[i - 1]); + const b : usize = v * NG_V + w; + if (ng_bi[b] == 0) { + ng_bi_types[v] = ng_bi_types[v] + 1; + } + ng_bi[b] = ng_bi[b] + 1; + ng_bi_row[v] = ng_bi_row[v] + 1; + } + if (i >= 2) { + const u : usize = @intCast(tokens[i - 2]); + const v2 : usize = @intCast(tokens[i - 1]); + const c : usize = u * NG_V + v2; + if (ng_tri_add(ng_tri_key(u, v2, w))) { + ng_tri_types[c] = ng_tri_types[c] + 1; + } + ng_tri_ctx[c] = ng_tri_ctx[c] + 1; + } + i = i + 1; + } +} + +; --------------------------------------------------------------------------------------------- +; 2. The three models +; --------------------------------------------------------------------------------------------- + +fn ng_p1(w: usize) -> f64 { + const num : f64 = @floatFromInt(ng_uni[w] + 1); + const den : f64 = @floatFromInt(ng_total + @as(u64, NG_V)); + return num / den; +} + +fn ng_p2(v: usize, w: usize) -> f64 { + const row : u64 = ng_bi_row[v]; + if (row == 0) { + return ng_p1(w); + } + var kept : f64 = @as(f64, @floatFromInt(ng_bi[v * NG_V + w])) - NG_D; + if (kept < 0.0) { + kept = 0.0; + } + const types : f64 = @floatFromInt(ng_bi_types[v]); + return (kept + NG_D * types * ng_p1(w)) / @as(f64, @floatFromInt(row)); +} + +fn ng_p3(u: usize, v: usize, w: usize) -> f64 { + const c : usize = u * NG_V + v; + const ctx : u32 = ng_tri_ctx[c]; + if (ctx == 0) { + return ng_p2(v, w); + } + var kept : f64 = @as(f64, @floatFromInt(ng_tri_count(ng_tri_key(u, v, w)))) - NG_D; + if (kept < 0.0) { + kept = 0.0; + } + const types : f64 = @floatFromInt(ng_tri_types[c]); + return (kept + NG_D * types * ng_p2(v, w)) / @as(f64, @floatFromInt(ctx)); +} + +/// -ln P of token tokens[i] under order k (1, 2, 3), with the context limited to tokens[lo..i] +fn ng_nll(tokens: []const u16, lo: usize, i: usize, k: usize) -> f64 { + const w : usize = @intCast(tokens[i]); + if (k == 1 or i == lo) { + return 0.0 - log(ng_p1(w)); + } + const v : usize = @intCast(tokens[i - 1]); + if (k == 2 or i == lo + 1) { + return 0.0 - log(ng_p2(v, w)); + } + const u : usize = @intCast(tokens[i - 2]); + return 0.0 - log(ng_p3(u, v, w)); +} + +; --------------------------------------------------------------------------------------------- +; 3. The measures +; --------------------------------------------------------------------------------------------- + +fn ng_acc_reset() -> void { + var k : usize = 0; + while (k < 3) { + ng_s_stream[k] = 0.0; + ng_s_full[k] = 0.0; + ng_s_last[k] = 0.0; + k = k + 1; + } + ng_n_stream = 0; + ng_n_pos = 0; + ng_n_win = 0; +} + +/// Every token of tokens[1..n] from the tokens before it +fn ng_stream(tokens: []const u16, n: usize) -> void { + var i : usize = 1; + while (i < n) { + var k : usize = 0; + while (k < 3) { + ng_s_stream[k] = ng_s_stream[k] + ng_nll(tokens, 0, i, k + 1); + k = k + 1; + } + ng_n_stream = ng_n_stream + 1; + i = i + 1; + } +} + +/// One eval_hslm window at `start`: position p scores token start + p + 1 from start..start + p +fn ng_window(tokens: []const u16, n: usize, start: usize) -> bool { + if (start + NG_WINDOW > n) { + return false; + } + var p : usize = 0; + while (p < NG_CTX) { + var k : usize = 0; + while (k < 3) { + const t : f64 = ng_nll(tokens, start, start + p + 1, k + 1); + ng_s_full[k] = ng_s_full[k] + t; + if (p + 1 == NG_CTX) { + ng_s_last[k] = ng_s_last[k] + t; + } + k = k + 1; + } + ng_n_pos = ng_n_pos + 1; + p = p + 1; + } + ng_n_win = ng_n_win + 1; + return true; +} + +/// eval_hslm's spacing for `count` windows over n tokens (its ev_stride_for with stride 0) +fn ng_stride_for(count: usize, n: usize) -> usize { + if (n <= NG_CTX + 1 or count == 0) { + return 1; + } + var stride : usize = (n - NG_CTX - 1) / count; + if (stride == 0) { + stride = 1; + } + return stride; +} + +/// The windows that fit +fn ng_windows(tokens: []const u16, n: usize, count: usize) -> usize { + const stride : usize = ng_stride_for(count, n); + var used : usize = 0; + var k : usize = 0; + while (k < count) { + if (ng_window(tokens, n, k * stride)) { + used = used + 1; + } + k = k + 1; + } + return used; +} + +fn ng_mean(sum: f64, n: u64) -> f64 { + return sum / @as(f64, @floatFromInt(n)); +} + +; --------------------------------------------------------------------------------------------- +; 4. Data, output, main +; --------------------------------------------------------------------------------------------- + +/// The byte count of /tmp/hslm/valid.txt, or 0 when it cannot be read +fn ng_count_bytes() -> u64 { + const fd : i32 = host_open("/tmp/hslm/valid.txt", 0); + if (fd < 0) { + return 0; + } + var total : u64 = 0; + var done : bool = false; + while (!done) { + const got : i64 = tr_read_full(fd, &io_chunk); + if (got <= 0) { + done = true; + } else { + total = total + @as(u64, @intCast(got)); + if (got < @as(i64, @intCast(IO_CHUNK))) { + done = true; + } + } + } + tr_note(host_close(fd)); + return total; +} + +fn ng_say(text: []const u8) -> void { + ln_reset(); + ln_str(text); + ln_flush(); +} + +fn ng_order_line(k: usize) -> void { + const sl : f64 = ng_mean(ng_s_stream[k], ng_n_stream); + const fl : f64 = ng_mean(ng_s_full[k], ng_n_pos); + const ll : f64 = ng_mean(ng_s_last[k], ng_n_win); + const bpb : f64 = sl / log(2.0) * @as(f64, @floatFromInt(nvalid)) / @as(f64, @floatFromInt(ng_bytes)); + ln_reset(); + ln_str("ngram order "); + ln_u64(@intCast(k + 1)); + ln_str(" stream_tokens "); + ln_u64(ng_n_stream); + ln_str(" stream_loss "); + ln_f64(sl, 4); + ln_str(" stream_ppl "); + ln_f64(exp(sl), 2); + ln_str(" stream_bpb "); + ln_f64(bpb, 4); + ln_str(" windows "); + ln_u64(ng_n_win); + ln_str(" positions "); + ln_u64(ng_n_pos); + ln_str(" full_loss "); + ln_f64(fl, 4); + ln_str(" full_ppl "); + ln_f64(exp(fl), 2); + ln_str(" last_loss "); + ln_f64(ll, 4); + ln_str(" last_ppl "); + ln_f64(exp(ll), 2); + ln_flush(); +} + +fn ng_data_line() -> void { + ln_reset(); + ln_str("ngram data train_tokens "); + ln_u64(@intCast(ntrain)); + ln_str(" valid_tokens "); + ln_u64(@intCast(nvalid)); + ln_str(" valid_bytes "); + ln_u64(ng_bytes); + ln_str(" trigram_types "); + ln_u64(@intCast(ng_hash_used)); + ln_str(" discount "); + ln_f64(NG_D, 2); + ln_flush(); +} + +fn ng_main() -> i32 { + const nt : i64 = tr_read_tok(0, &train_tok); + const nv : i64 = tr_read_tok(1, &valid_tok); + if (nt < @as(i64, @intCast(NG_WINDOW)) or nv < @as(i64, @intCast(NG_WINDOW))) { + ng_say("ngram_hslm: /tmp/hslm/train.tok or valid.tok missing or bad"); + return NG_RC_DATA; + } + ntrain = @intCast(nt); + nvalid = @intCast(nv); + ng_bytes = ng_count_bytes(); + if (ng_bytes == 0) { + ng_say("ngram_hslm: /tmp/hslm/valid.txt missing"); + return NG_RC_TEXT; + } + ng_reset(NG_HASH_MAX); + ng_fit(train_tok[0..ntrain], ntrain); + if (ng_full) { + ng_say("ngram_hslm: the trigram hash is full"); + return NG_RC_FULL; + } + ng_acc_reset(); + ng_stream(valid_tok[0..nvalid], nvalid); + if (ng_windows(valid_tok[0..nvalid], nvalid, NG_WINDOWS) == 0) { + ng_say("ngram_hslm: no window fits in valid.tok"); + return NG_RC_DATA; + } + ng_data_line(); + var k : usize = 0; + while (k < 3) { + const a : f64 = ng_s_stream[k]; + const b : f64 = ng_s_full[k]; + if (a != a or b != b) { + ng_say("ngram_hslm: a loss is NaN"); + return NG_RC_NAN; + } + ng_order_line(k); + k = k + 1; + } + return NG_RC_OK; +} + +/// The program `t27c compile --backend ml` builds +pub fn main() -> void { + assert ng_main() == NG_RC_OK; +} + +; --------------------------------------------------------------------------------------------- +; 5. Test helpers (each returns the number its test checks) and the tests +; --------------------------------------------------------------------------------------------- + +/// Counts of tr_synth's learnable stream (next = 7 cur + 3 mod 729, one token in ten random) in a +/// small hash +fn ng_tiny_fit() -> void { + tr_synth(train_tok[0..20000], 11); + tr_synth(valid_tok[0..3000], 13); + ntrain = 20000; + nvalid = 3000; + ng_reset(65536); + ng_fit(train_tok[0..ntrain], ntrain); +} + +/// |1 - sum over w of P_k(w | context)| for order k, context (u, v) +fn ng_sum_gap(k: usize, u: usize, v: usize) -> f64 { + var s : f64 = 0.0; + var w : usize = 0; + while (w < NG_V) { + if (k == 1) { + s = s + ng_p1(w); + } else if (k == 2) { + s = s + ng_p2(v, w); + } else { + s = s + ng_p3(u, v, w); + } + w = w + 1; + } + return abs(1.0 - s); +} + +/// The largest sum gap over the three orders, a seen context (u v = 1, 10) and an unseen one +fn ng_worst_sum_gap() -> f64 { + ng_tiny_fit(); + var worst : f64 = 0.0; + var k : usize = 1; + while (k <= 3) { + const g1 : f64 = ng_sum_gap(k, 1, 10); + const g2 : f64 = ng_sum_gap(k, 700, 3); + if (g1 > worst) { + worst = g1; + } + if (g2 > worst) { + worst = g2; + } + k = k + 1; + } + return worst; +} + +fn ng_seen_context() -> u64 { + ng_tiny_fit(); + return @as(u64, ng_tri_ctx[1 * NG_V + 10]) * 1000 + @as(u64, ng_tri_ctx[700 * NG_V + 3]); +} + +/// The stream losses on the synthetic held-out stream, order k (1..3) +fn ng_tiny_stream(k: usize) -> f64 { + ng_tiny_fit(); + ng_acc_reset(); + ng_stream(valid_tok[0..nvalid], nvalid); + return ng_mean(ng_s_stream[k - 1], ng_n_stream); +} + +/// A second-order stream over 27 ids: next = (cur + 5 prev + 3) mod 27, one token in ten random. +/// Given cur alone the next id is uniform over the 27 values of prev, so only the trigram sees it. +fn ng_synth2(dst: []u16, seed: u64) -> void { + var x : u64 = seed; + var prev : u64 = 1; + var cur : u64 = 2; + var i : usize = 0; + while (i < dst.len) { + x = tr_xs(x); + var nxt : u64 = (cur + 5 * prev + 3) % 27; + if (x % 10 == 0) { + nxt = (x >> 20) % 27; + } + dst[i] = @intCast(nxt); + prev = cur; + cur = nxt; + i = i + 1; + } +} + +/// The stream loss of order k on ng_synth2 held-out tokens, after fitting on ng_synth2 tokens +fn ng_tiny2_stream(k: usize) -> f64 { + ng_synth2(train_tok[0..20000], 11); + ng_synth2(valid_tok[0..3000], 13); + ntrain = 20000; + nvalid = 3000; + ng_reset(65536); + ng_fit(train_tok[0..ntrain], ntrain); + ng_acc_reset(); + ng_stream(valid_tok[0..nvalid], nvalid); + return ng_mean(ng_s_stream[k - 1], ng_n_stream); +} + +/// With no training tokens every order is the uniform 1/729 +fn ng_empty_gap() -> f64 { + ng_reset(4096); + tr_synth(valid_tok[0..500], 13); + nvalid = 500; + ng_acc_reset(); + ng_stream(valid_tok[0..nvalid], nvalid); + var worst : f64 = 0.0; + var k : usize = 0; + while (k < 3) { + const g : f64 = abs(ng_mean(ng_s_stream[k], ng_n_stream) - log(729.0)); + if (g > worst) { + worst = g; + } + k = k + 1; + } + return worst; +} + +/// A trigram seen twice counts 2, one seen once counts 1, an unseen one 0 +fn ng_hash_counts() -> u64 { + ng_reset(4096); + const a : bool = ng_tri_add(ng_tri_key(5, 6, 7)); + const b : bool = ng_tri_add(ng_tri_key(5, 6, 7)); + const c : bool = ng_tri_add(ng_tri_key(728, 728, 728)); + var r : u64 = @as(u64, ng_tri_count(ng_tri_key(5, 6, 7))) * 100; + r = r + @as(u64, ng_tri_count(ng_tri_key(728, 728, 728))) * 10; + r = r + @as(u64, ng_tri_count(ng_tri_key(6, 5, 7))); + if (!a or b or !c) { + r = r + 1000; + } + return r + @as(u64, @intCast(ng_hash_used)) * 10000; +} + +/// One window: position 0 is the bigram term, position 26 the trigram term, full is the mean +fn ng_window_gap() -> f64 { + ng_tiny_fit(); + ng_acc_reset(); + if (!ng_window(valid_tok[0..nvalid], nvalid, 100)) { + return 1.0; + } + var s : f64 = 0.0; + var p : usize = 0; + while (p < NG_CTX) { + const w : usize = @intCast(valid_tok[101 + p]); + const v : usize = @intCast(valid_tok[100 + p]); + if (p == 0) { + s = s - log(ng_p2(v, w)); + } else { + const u : usize = @intCast(valid_tok[99 + p]); + s = s - log(ng_p3(u, v, w)); + } + p = p + 1; + } + return abs(s - ng_s_full[2]) + abs(ng_mean(ng_s_full[2], ng_n_pos) - s / 27.0); +} + +fn ng_short_window_refused() -> u64 { + ng_tiny_fit(); + ng_acc_reset(); + var r : u64 = 0; + if (ng_window(valid_tok[0..nvalid], 27, 0)) { + r = r + 1; + } + if (ng_window(valid_tok[0..nvalid], 28, 0)) { + r = r + 10; + } + return r + ng_n_pos * 100; +} + +test every_order_sums_to_one_over_the_vocabulary + given gap = ng_worst_sum_gap() + and seen = ng_seen_context() + then gap < 1e-9 + and seen > 1000 + +test no_training_tokens_gives_ln_729_for_every_order + given gap = ng_empty_gap() + then gap < 1e-12 + +test the_hash_counts_repeats_and_keeps_keys_apart + given r = ng_hash_counts() + then r == 20210 + +test a_first_order_stream_is_learned_by_the_bigram + given l1 = ng_tiny_stream(1) + and l2 = ng_tiny_stream(2) + then l2 < l1 - 3.0 + +test a_second_order_stream_needs_the_trigram + given l2 = ng_tiny2_stream(2) + and l3 = ng_tiny2_stream(3) + then l3 < l2 - 2.0 + +test the_window_measure_is_bigram_at_position_0_and_trigram_after + given gap = ng_window_gap() + then gap < 1e-9 + +test a_window_of_27_tokens_is_refused_and_28_is_taken + given r = ng_short_window_refused() + then r == 2710 + +test the_window_spacing_is_eval_hslms + given s1 = ng_stride_for(243, 7689354) + and s2 = ng_stride_for(19683, 7689354) + and s3 = ng_stride_for(5, 20) + then s1 == 31643 + and s2 == 390 + and s3 == 1