diff --git a/.trinity/seals/AdamW.json b/.trinity/seals/AdamW.json index e4c9084256..add3f8d705 100644 --- a/.trinity/seals/AdamW.json +++ b/.trinity/seals/AdamW.json @@ -1,14 +1,19 @@ { - "gen_hash_c": "sha256:c01ec3973f0a92519fd5bf1490d014ed9651d4cb0f4a89669b288661a85aa9ec", - "gen_hash_rust": "sha256:0010ce0e41a962f8db2e94baf5515a99b5312d3ef74f54ade99c42ebb65a6274", - "gen_hash_verilog": "sha256:ebecfc4c499455c39a38af49f4a9ac544e70c4452685d9583db2392044160964", - "gen_hash_zig": "sha256:4ca45035d84b11fef4daa132ea6fde13935fa74f55ef8bdc9c322f7b460d26e0", + "gen_hash_c": "sha256:5e0e6de65832ccf6c4f64d7410fe24af093b2d3a137a17bf37edc8a10c75a75a", + "gen_hash_rust": "sha256:93895575b4e59283c27fc742a1261d832cc79706db435e7ca5ffc1466a63e0f0", + "gen_hash_verilog": "sha256:c7f2055f86cd0f47f00558cfd67c2b5170d8cefa72db2a9a22798f57bf5c2878", + "gen_hash_zig": "sha256:a5225b8c6ab539841efc6bfe8bb875fac265068d0fc8019370f1c14eb7728685", "module": "AdamW", "ring": 12, - "sealed_at": "2026-10-11T10:43:20Z", - "spec_hash": "sha256:1b8cb03ec9edccd4b26b996d4f789f3ab4fc3feaca5df97d75ee003dd26070b0", + "sealed_at": "2026-10-17T18:31:13Z", + "spec_hash": "sha256:b58a44f505aeb0bbbc664b8708d13c9399f73b714c6b22d0b46cc5d33e853be7", "spec_path": "specs/ml/optimizer/adamw.t27", "tests": { - "blocked": "does not compile: spec.zig:54:13: error: use of undeclared identifier 'allocate_array'" + "failed": 0, + "failing": [], + "forced": false, + "passed": 12, + "total": 12, + "vacuous": 0 } } \ No newline at end of file diff --git a/.trinity/seals/optimizer_AdamW.json b/.trinity/seals/optimizer_AdamW.json index 86c41ce255..3f0af3b7c0 100644 --- a/.trinity/seals/optimizer_AdamW.json +++ b/.trinity/seals/optimizer_AdamW.json @@ -1,5 +1,5 @@ { - "built_by": "t27c-bootstrap@0.5.2+cc19d1fcd", + "built_by": "t27c-bootstrap@0.5.2+efe3c94", "closure": [ { "spec": "specs/base/types.t27", @@ -10,24 +10,33 @@ "spec_hash": "sha256:57187e7afda1d88167b2af0382ad7f0623aa6b269ba020a78e3ffeebba1161bb" }, { - "spec": "specs/numeric/gf16.t27", - "spec_hash": "sha256:2a33d77db23966331716955924d08e30d12a3af361067da0a2c7072b7f4b7b0d" + "spec": "specs/ml/arena_heap.t27", + "spec_hash": "sha256:dadb837793dbfaa2f313bb3dfebe92764fe23c96228e11c10b933958e9feb0bb" + }, + { + "spec": "specs/ml/optimizer/adam.t27", + "spec_hash": "sha256:75f3d191b54ab786de3aa085164717813022eb8f403da12fe5b55c7bda8aa73f" } ], "config": "gen=c,rust,verilog,zig;tests=zig", - "gen_hash_c": "sha256:c01ec3973f0a92519fd5bf1490d014ed9651d4cb0f4a89669b288661a85aa9ec", - "gen_hash_rust": "sha256:0010ce0e41a962f8db2e94baf5515a99b5312d3ef74f54ade99c42ebb65a6274", - "gen_hash_verilog": "sha256:ebecfc4c499455c39a38af49f4a9ac544e70c4452685d9583db2392044160964", - "gen_hash_zig": "sha256:4ca45035d84b11fef4daa132ea6fde13935fa74f55ef8bdc9c322f7b460d26e0", + "gen_hash_c": "sha256:5e0e6de65832ccf6c4f64d7410fe24af093b2d3a137a17bf37edc8a10c75a75a", + "gen_hash_rust": "sha256:93895575b4e59283c27fc742a1261d832cc79706db435e7ca5ffc1466a63e0f0", + "gen_hash_verilog": "sha256:c7f2055f86cd0f47f00558cfd67c2b5170d8cefa72db2a9a22798f57bf5c2878", + "gen_hash_zig": "sha256:a5225b8c6ab539841efc6bfe8bb875fac265068d0fc8019370f1c14eb7728685", "module": "AdamW", "ring": 12, "seal_schema": 2, - "sealed_at": "2026-10-11T10:43:20Z", + "sealed_at": "2026-10-17T18:31:13Z", "sealed_by": "t27c-bootstrap@0.5.2", - "spec_hash": "sha256:1b8cb03ec9edccd4b26b996d4f789f3ab4fc3feaca5df97d75ee003dd26070b0", + "spec_hash": "sha256:b58a44f505aeb0bbbc664b8708d13c9399f73b714c6b22d0b46cc5d33e853be7", "spec_path": "specs/ml/optimizer/adamw.t27", "tests": { - "blocked": "does not compile: spec.zig:54:13: error: use of undeclared identifier 'allocate_array'" + "failed": 0, + "failing": [], + "forced": false, + "passed": 12, + "total": 12, + "vacuous": 0 }, "toolchain": { "test_runner": "sha256:1982ae948e24d00c5d221ccec12dfce1eb01204b53e37b0f857f93b16d64629d", diff --git a/docs/reports/t27b_expectations.json b/docs/reports/t27b_expectations.json index 537ba8aeca..cfbb74965e 100644 --- a/docs/reports/t27b_expectations.json +++ b/docs/reports/t27b_expectations.json @@ -12,7 +12,7 @@ "reference-bug": "the reference passes it for the wrong reason; t27b is right not to", "n/a": "outside what t27b is for (say why in the PR that sets it)" }, - "max_not_pass": 28, + "max_not_pass": 31, "entries": [ {"path": "specs/a/b_c.t27", "t27b": "pass"}, {"path": "specs/a_b/c.t27", "t27b": "pass"}, @@ -413,6 +413,8 @@ {"path": "specs/ml/activation/silu_swish_activation.t27", "t27b": "blocked", "blocker": "ExprCall(undeclared fn)", "reason": "unimplemented"}, {"path": "specs/ml/activation/softmax.t27", "t27b": "blocked", "blocker": "ExprCall(undeclared fn)", "reason": "unimplemented"}, {"path": "specs/ml/activation/tanh_activation.t27", "t27b": "blocked", "blocker": "ExprCall(undeclared fn)", "reason": "unimplemented"}, + {"path": "specs/ml/arena_heap.t27", "t27b": "blocked", "blocker": "ExprCall(@mod)", "reason": "unimplemented"}, + {"path": "specs/ml/file_io_contract.t27", "t27b": "blocked", "blocker": "ExprCall(@mod)", "reason": "unimplemented"}, {"path": "specs/ml/igla_champion_capsule.t27", "t27b": "blocked", "blocker": "ExprIdentifier(undeclared)", "reason": "unimplemented"}, {"path": "specs/ml/layers/avgpool2d_layer.t27", "t27b": "pass"}, {"path": "specs/ml/layers/batchnorm_layer.t27", "t27b": "blocked", "blocker": "ExprCall(undeclared fn)", "reason": "unimplemented"}, @@ -432,6 +434,7 @@ {"path": "specs/ml/loss/mse_loss.t27", "t27b": "pass"}, {"path": "specs/ml/optimizer/adagrad.t27", "t27b": "blocked", "blocker": "ExprCall(undeclared fn)", "reason": "unimplemented"}, {"path": "specs/ml/optimizer/adam.t27", "t27b": "blocked", "blocker": "ExprCall(undeclared fn)", "reason": "unimplemented"}, + {"path": "specs/ml/optimizer/adamw.t27", "t27b": "blocked", "blocker": "ExprCall(undeclared fn)", "reason": "unimplemented"}, {"path": "specs/ml/optimizer/lamb.t27", "t27b": "blocked", "blocker": "ExprCall(undeclared fn)", "reason": "unimplemented"}, {"path": "specs/ml/optimizer/race_config.t27", "t27b": "blocked", "blocker": "ExprIdentifier(undeclared)", "reason": "unimplemented"}, {"path": "specs/ml/optimizer/rmsprop.t27", "t27b": "pass"}, diff --git a/docs/reports/type_conflicts.json b/docs/reports/type_conflicts.json index 2fbb40d67d..59cd172373 100644 --- a/docs/reports/type_conflicts.json +++ b/docs/reports/type_conflicts.json @@ -53,7 +53,6 @@ "MigrationStep", "ModuleInterface", "Node", - "OptimizerStepResult", "OscillatorState", "Outputs", "ParseError", diff --git a/specs/igla/own_model.t27 b/specs/igla/own_model.t27 index c412ada08d..47466c51f5 100644 --- a/specs/igla/own_model.t27 +++ b/specs/igla/own_model.t27 @@ -31,6 +31,15 @@ ; vocab-729 head and cross-entropy), its backward reaches the embedding, and every gradient is ; checked against a central difference; a mutant that drops one alpha fails 10 of the 46. ; Not in it then: the per-block VSA path, f32, and the trainer's STE update (shadow weights). +; - re-test on 2026-10-18, same t27c, lab checkout brought to master, after the PR that gave +; t27 a pure arena allocator and a file I/O contract (specs/ml/arena_heap.t27 10/10, +; specs/ml/file_io_contract.t27 12/12) and rewrote optimizer/adamw (#10077): adamw 12/12 +; pass, 0 vacuous, its moment buffers carved from the arena. cross_entropy_loss 5/5, sgd 4/4, +; adam 4/4, lamb 4/4, nn/hslm (61/61 after the VSA PR) and hslm/forward_pass 14/14 still pass, so 7 of the 17 +; compile. nn/attention ('PHI_CUBED_INV') and the nine coder specs stay BLOCKED. Cells 24 and +; 25 stay gaps: the arena and the file contract are specs of what t27c must lower, and t27c +; lowers neither (no heap, no file I/O, no `--backend ml`; that work waits for the owner's +; owner-approved-foreign label). ; - test on 2026-10-18, same t27c, after the PR that added the training step (STE, backward, data): ; nn/ternary_ste 26/26, nn/hslm_backward 15/15, nn/hslm_data 39/39 pass, 0 vacuous. The STE is ; the TWN quantizer (delta 0.7 * mean abs, alpha = mean abs of the kept weights), the @@ -39,7 +48,7 @@ ; (ctx 27, head_dim 3): its trit, bias and input gradients match a central difference to 1e-6 ; and 20 steps cut the loss below half. hslm_data holds the curriculum, next-token batches and ; the MSLH header and layout, whose byte count equals hslm_step_10000.bin (4971796). Not in -; them: the STE on the whole 3-block model, file I/O, heap, the tokenizer, shuffle, AdamW. +; them: the STE on the whole 3-block model, the tokenizer, shuffle; AdamW, heap and file I/O are in the entry above. ; - re-test on 2026-10-18, same t27c, after the PR that added the VSA path: nn/hslm 61/61 tests ; pass, 12 invariants proved, 0 vacuous. Each block now runs the VSA stream of model.zig ; forwardTrain (bind, rotate, the PHI_INV gate, +0.1 * project(context) when it fires, the @@ -58,11 +67,12 @@ ; cell weighted the same; that weighting is a convention, not a measurement. ; THE ANSWER this map gives: HSLM was trained before, in Zig only (64 files); its best runs -; (PPL 2.96 at ctx 27) were wiped by EXP-002 and survive only in the ledger. In t27 the loss, three +; (PPL 2.96 at ctx 27) were wiped by EXP-002 and survive only in the ledger. In t27 the loss, four ; optimizers and the HSLM spec's forward and backward passes pass, and that spec is now the trained ; model, its VSA path included (see MEASURED); the ternary STE, the training step and the data ; and checkpoint layout now pass too, on a model slice and on bytes; the STE over the whole model, -; heap and file I/O and `t27c --backend ml` do not exist yet. +; heap and file I/O exist as pure t27 contracts (ml/arena_heap, ml/file_io_contract) that t27c does not +; lower, and `t27c --backend ml` does not exist. ; Not claimed: no model quality beyond the numbers in SOURCES; no GIF or MP4, stills only. ; phi^2 + 1/phi^2 = 3 | TRINITY @@ -158,7 +168,7 @@ pub const SOURCES : [27]str = [ "trinity-training src/hslm, 64 Zig files", "trinity data/wave9/worker-2/hslm_step_10000.bin", "t27 specs/ml/loss/cross_entropy_loss.t27", - "t27 specs/ml/optimizer/{sgd,adam,lamb}.t27", + "t27 specs/ml/optimizer/{sgd,adam,adamw,lamb}.t27", "t27 specs/nn/hslm.t27 and specs/hslm/forward_pass.t27", "t27 specs/nn/hslm.t27 backward and nn/hslm_backward.t27 training step on a head slice, finite-difference checked", "t27 specs/nn/ternary_ste.t27, 26 tests; whole-model shadow weights not yet", @@ -167,7 +177,7 @@ pub const SOURCES : [27]str = [ "trinity-fpga b2bd5a30, 7682304/7682304 receipts", "board-production-2dec.up.railway.app/state", "t27 specs/port/trinity/fpga/openxc7-synth/ternary_matvec.t27", - "t27c test-report on 17 ML specs: 6 compile", + "t27c test-report on 17 ML specs: 7 compile", "", "", "", @@ -183,7 +193,7 @@ pub const MISSING : [27]str = [ "", "", "", "", "", "", "", - "heap allocation and file I/O lowered by t27c: neither exists", + "heap allocation and file I/O lowered by t27c: neither exists; ml/arena_heap.t27 and ml/file_io_contract.t27 are the contract it must meet", "a t27c backend that lowers a model spec to a training binary: none exists", "embeddings, head and weights in DDR3 on the board: they stay on the host", "" @@ -200,11 +210,11 @@ pub const STATUS : [27]u8 = [4, 4, 4, 3, 3, 3, 3, 0, 0, 4, 4, 4, 4, 4, 4, 3, 3, pub const LABELS : [27]str = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15", "16", "17", "18", "19", "20", "21", "22", "23", "24", "25", "26", "27"]; /// Percent of each part that exists, estimated from its source (see the header). Never a status. -pub const PCT : [27]u8 = [60, 70, 80, 80, 60, 10, 50, 0, 0, 100, 100, 20, 80, 65, 95, 75, 55, 40, 70, 15, 60, 40, 35, 0, 0, 0, 50]; +pub const PCT : [27]u8 = [60, 70, 80, 80, 60, 10, 50, 0, 0, 100, 100, 20, 80, 65, 95, 75, 55, 40, 70, 15, 60, 40, 41, 0, 0, 0, 50]; /// The t27c test-report counts the PRACTICE cells quote. pub const ML_SPECS_MEASURED : i32 = 17; -pub const ML_SPECS_COMPILING : i32 = 6; +pub const ML_SPECS_COMPILING : i32 = 7; pub const CODER_SPECS : i32 = 9; pub const CODER_SPECS_COMPILING : i32 = 0; @@ -223,7 +233,7 @@ pub const ROWS : [20]str = [ "64 Zig files; best PPL 2.96 at ctx 27", "Those checkpoints were wiped", "In t27: loss, optimizers, HSLM pass", - "6 of 17 ML specs compile on the lab", + "7 of 17 ML specs compile on the lab", "0 of 9 IGLA CODER specs compile", "1 published, 12 verified, 9 drafts", "5 gaps, about 48 percent overall", @@ -326,7 +336,7 @@ test the_compile_counts_are_the_measured_ones { assert TEST_PASS[5] == false; // the compile cell's percent is the measured share, truncated const share : i32 = share_pct(ML_SPECS_COMPILING, ML_SPECS_MEASURED); - assert share == 35; + assert share == 41; const p22 : i32 = @intCast(PCT[22]); assert p22 == share; assert share_pct(CODER_SPECS_COMPILING, CODER_SPECS) == 0; diff --git a/specs/igla/own_model_board.t27 b/specs/igla/own_model_board.t27 index 99eaec3f87..211110831b 100644 --- a/specs/igla/own_model_board.t27 +++ b/specs/igla/own_model_board.t27 @@ -61,7 +61,7 @@ pub const KEY_PCT : str = "pct"; pub const LENS_PCT : [3]u8 = [45, 70, 30]; pub const CELL_STATUS : [27]u8 = [4, 4, 4, 3, 3, 3, 3, 0, 0, 4, 4, 4, 4, 4, 4, 3, 3, 3, 4, 4, 5, 4, 3, 0, 0, 0, 3]; -pub const CELL_PCT : [27]u8 = [60, 70, 80, 80, 60, 10, 50, 0, 0, 100, 100, 20, 80, 65, 95, 75, 55, 40, 70, 15, 60, 40, 35, 0, 0, 0, 50]; +pub const CELL_PCT : [27]u8 = [60, 70, 80, 80, 60, 10, 50, 0, 0, 100, 100, 20, 80, 65, 95, 75, 55, 40, 70, 15, 60, 40, 41, 0, 0, 0, 50]; /// The size of the compact metrics object, in bytes, computed from the keys and values above. pub const METRICS_BYTES : i32 = 624; diff --git a/specs/ml/arena_heap.t27 b/specs/ml/arena_heap.t27 new file mode 100644 index 0000000000..61d9978e1c --- /dev/null +++ b/specs/ml/arena_heap.t27 @@ -0,0 +1,393 @@ +// SPDX-License-Identifier: Apache-2.0 +; specs/ml/arena_heap.t27 -- a bump (arena) allocator model over a fixed backing array. +; Part of gHashTag/t27#10024 (our own model), cell 24 "Heap and file I/O in t27c"; issue #10077. + +; WHY: training an HSLM needs buffers whose size is known only when the run starts (Adam moments, +; activations, a data batch). t27 has no heap, and the lab shows it: specs/ml/optimizer/adamw.t27 +; calls allocate_array and is BLOCKED ("use of undeclared identifier 'allocate_array'"). The HSLM +; spec (specs/nn/hslm.t27) gets around it with module storage (`var model : HSLMWeights`), which +; fixes every size at compile time. An arena is the smallest thing that keeps one backing array +; and hands out non-overlapping pieces of it, so a trainer can pick its sizes at run time. + +; WHAT THIS MODELS: the allocator is a value. `Arena` holds four integers and no memory; the +; memory is a caller-owned `[]f32` backing array, and an allocation is an OFFSET into it, in +; f32 slots (not bytes). alloc never traps: out of memory, a bad size and a bad alignment come +; back as an error code with the arena unchanged. reset forgets every allocation but keeps the +; high-water mark, which is how a trainer learns what capacity the next run needs. + +; WHAT t27c WOULD HAVE TO DO (the contract, not an implementation): lower `Arena` to the struct +; below, lower `arena_alloc` to this exact arithmetic, and hand a trainer a backing array it +; owns. specs/ml/file_io_contract.t27 is the twin for files. The Rust work is its own issue and +; waits for the owner's owner-approved-foreign label; nothing here needs it. + +; Not claimed: no thread safety, no free of a single block (use arena_mark and arena_release for +; scratch space), no bytes: sizes and offsets count f32 slots. +; phi^2 + 1/phi^2 = 3 | TRINITY + +module arena_heap; + +pub const ARENA_OK : i32 = 0; +pub const ARENA_ERR_OOM : i32 = 1; +pub const ARENA_ERR_SIZE : i32 = 2; +pub const ARENA_ERR_ALIGN : i32 = 3; +pub const ARENA_ERR_MARK : i32 = 4; + +/// The offset an allocation returns when it fails; no real offset is negative. +pub const ARENA_NO_OFFSET : i32 = -1; + +/// cap: slots in the backing array; top: first free slot; high: the largest top ever reached; +/// resets: how many times reset ran. +pub const Arena = struct { + cap : i32, + top : i32, + high : i32, + resets : i32, +}; + +/// offset is ARENA_NO_OFFSET whenever code is not ARENA_OK, and then arena is the arena that went in. +pub const AllocResult = struct { + arena : Arena, + offset : i32, + code : i32, +}; + +/// Two equal-sized buffers for an optimizer's first and second moments. +pub const MomentsResult = struct { + arena : Arena, + m_off : i32, + v_off : i32, + code : i32, +}; + +pub fn arena_new(cap: i32) -> Arena { + return Arena{ .cap = cap, .top = 0, .high = 0, .resets = 0 }; +} + +pub fn arena_used(a: Arena) -> i32 { + return a.top; +} + +pub fn arena_left(a: Arena) -> i32 { + return a.cap - a.top; +} + +pub fn arena_high_water(a: Arena) -> i32 { + return a.high; +} + +/// Is n a positive power of two? Alignment must be one. +pub fn is_pow2(n: i32) -> bool { + if (n < 1) { + return false; + } + var m : i32 = n; + while (m > 1) { + if (@mod(m, 2) != 0) { + return false; + } + m = @divTrunc(m, 2); + } + return true; +} + +/// Smallest multiple of alignment that is >= x. alignment must be a positive power of two. +pub fn align_up(x: i32, alignment: i32) -> i32 { + return @divTrunc(x + alignment - 1, alignment) * alignment; +} + +fn failed(a: Arena, code: i32) -> AllocResult { + return AllocResult{ .arena = a, .offset = ARENA_NO_OFFSET, .code = code }; +} + +/// Take size slots, the first at a multiple of alignment. Never traps; on any error the arena is +/// returned unchanged. The size test is against what is left, not top + size, so a size near +/// the i32 maximum cannot overflow into a "fits". +pub fn arena_alloc(a: Arena, size: i32, alignment: i32) -> AllocResult { + if (size < 1) { + return failed(a, ARENA_ERR_SIZE); + } + if (is_pow2(alignment) == false) { + return failed(a, ARENA_ERR_ALIGN); + } + const start : i32 = align_up(a.top, alignment); + if (start > a.cap) { + return failed(a, ARENA_ERR_OOM); + } + if (size > a.cap - start) { + return failed(a, ARENA_ERR_OOM); + } + const end : i32 = start + size; + var high : i32 = a.high; + if (end > high) { + high = end; + } + return AllocResult{ + .arena = Arena{ .cap = a.cap, .top = end, .high = high, .resets = a.resets }, + .offset = start, + .code = ARENA_OK + }; +} + +/// Forget every allocation. The high-water mark stays: it measures the whole run. +pub fn arena_reset(a: Arena) -> Arena { + return Arena{ .cap = a.cap, .top = 0, .high = a.high, .resets = a.resets + 1 }; +} + +/// A mark is the current top. Release to a mark frees everything allocated after it. +pub fn arena_mark(a: Arena) -> i32 { + return a.top; +} + +/// Free back to a mark taken earlier. A mark above top (or below 0) is ARENA_ERR_MARK and the +/// arena comes back unchanged; releasing to top is a no-op. +pub fn arena_release(a: Arena, mark: i32) -> AllocResult { + if (mark < 0) { + return failed(a, ARENA_ERR_MARK); + } + if (mark > a.top) { + return failed(a, ARENA_ERR_MARK); + } + return AllocResult{ + .arena = Arena{ .cap = a.cap, .top = mark, .high = a.high, .resets = a.resets }, + .offset = mark, + .code = ARENA_OK + }; +} + +pub fn fill_slots(backing: []f32, from: i32, n: i32, value: f32) -> void { + var i : i32 = 0; + while (i < n) { + backing[@intCast(from + i)] = value; + i = i + 1; + } +} + +/// The optimizer use: take two buffers of n slots for adam/adamw moments, both or neither. +/// Both start at zero in the backing array. If the second does not fit, the arena and the +/// backing array are exactly as they came in. +pub fn moments_alloc(a: Arena, backing: []f32, n: i32) -> MomentsResult { + const first : AllocResult = arena_alloc(a, n, 1); + if (first.code != ARENA_OK) { + return MomentsResult{ .arena = a, .m_off = ARENA_NO_OFFSET, .v_off = ARENA_NO_OFFSET, .code = first.code }; + } + const second : AllocResult = arena_alloc(first.arena, n, 1); + if (second.code != ARENA_OK) { + return MomentsResult{ .arena = a, .m_off = ARENA_NO_OFFSET, .v_off = ARENA_NO_OFFSET, .code = second.code }; + } + fill_slots(backing, first.offset, n, 0.0); + fill_slots(backing, second.offset, n, 0.0); + return MomentsResult{ .arena = second.arena, .m_off = first.offset, .v_off = second.offset, .code = ARENA_OK }; +} + +// ---- helpers that exist only so the tests go through function calls ---- + +fn alloc_offset(cap: i32, size: i32, alignment: i32) -> i32 { + return arena_alloc(arena_new(cap), size, alignment).offset; +} + +fn alloc_code(cap: i32, size: i32, alignment: i32) -> i32 { + return arena_alloc(arena_new(cap), size, alignment).code; +} + +/// Three allocations of 10, 5 (aligned to 4) and 3: returns the arena after them. +fn three_allocs() -> Arena { + const r1 : AllocResult = arena_alloc(arena_new(100), 10, 1); + const r2 : AllocResult = arena_alloc(r1.arena, 5, 4); + const r3 : AllocResult = arena_alloc(r2.arena, 3, 1); + return r3.arena; +} + +/// Moments of n=4 in a backing array of 16 slots that held 9.0 everywhere. Writes 2.5 into +/// m[1] and 7.5 into v[2]; returns the sum of every slot of m and v, so a zero-fill that missed +/// a slot, or an overlap between m and v, changes the number. +fn moments_sum() -> f32 { + var backing : [16]f32 = undefined; + fill_slots(&backing, 0, 16, 9.0); + const r : MomentsResult = moments_alloc(arena_new(16), &backing, 4); + backing[@intCast(r.m_off + 1)] = 2.5; + backing[@intCast(r.v_off + 2)] = 7.5; + var sum : f32 = 0.0; + var i : i32 = 0; + while (i < 4) { + sum = sum + backing[@intCast(r.m_off + i)] + backing[@intCast(r.v_off + i)]; + i = i + 1; + } + return sum; +} + +/// A failed moments_alloc must not touch the backing array: slot 0 stays 9.0. +fn backing_after_failed_moments() -> f32 { + var backing : [16]f32 = undefined; + fill_slots(&backing, 0, 16, 9.0); + const r : MomentsResult = moments_alloc(arena_new(7), &backing, 4); + return backing[0] + backing[4]; +} + +/// The code and the top of a moments_alloc of n slots into capacity cap, as code * 1000 + top. +fn moments_outcome(cap: i32, n: i32) -> i32 { + var backing : [16]f32 = undefined; + fill_slots(&backing, 0, 16, 9.0); + const r : MomentsResult = moments_alloc(arena_new(cap), &backing, n); + return r.code * 1000 + r.arena.top; +} + +test a_fresh_arena_is_empty { + const a : Arena = arena_new(100); + assert arena_used(a) == 0; + assert arena_left(a) == 100; + assert arena_high_water(a) == 0; + assert a.resets == 0; +} + +test alloc_returns_offsets_that_do_not_overlap { + const a : Arena = three_allocs(); + // 10 at 0, then 5 aligned to 4 starts at 12 (slots 10 and 11 are padding), ends at 17, + // then 3 at 17 ends at 20 + assert a.top == 20; + assert arena_left(a) == 80; + assert arena_high_water(a) == 20; + const r1 : AllocResult = arena_alloc(arena_new(100), 10, 1); + assert r1.offset == 0; + const r2 : AllocResult = arena_alloc(r1.arena, 5, 4); + assert r2.offset == 12; + assert r2.arena.top == 17; + const r3 : AllocResult = arena_alloc(r2.arena, 3, 1); + assert r3.offset == 17; +} + +test alignment_pads_to_a_multiple { + assert align_up(0, 8) == 0; + assert align_up(1, 8) == 8; + assert align_up(8, 8) == 8; + assert align_up(9, 8) == 16; + assert align_up(13, 1) == 13; + assert alloc_offset(100, 4, 16) == 0; + const r1 : AllocResult = arena_alloc(arena_new(100), 4, 1); + const r2 : AllocResult = arena_alloc(r1.arena, 4, 16); + assert r2.offset == 16; + assert r2.arena.top == 20; +} + +test an_exact_fit_succeeds_and_one_more_is_out_of_memory { + const r1 : AllocResult = arena_alloc(arena_new(100), 83, 1); + const r2 : AllocResult = arena_alloc(r1.arena, 17, 1); + assert r2.code == ARENA_OK; + assert r2.arena.top == 100; + assert arena_left(r2.arena) == 0; + const r3 : AllocResult = arena_alloc(r2.arena, 1, 1); + assert r3.code == ARENA_ERR_OOM; + assert r3.offset == ARENA_NO_OFFSET; + // the failed call returns the arena that went in + assert r3.arena.top == 100; + assert r3.arena.high == 100; +} + +test out_of_memory_is_a_code_not_a_trap { + assert alloc_code(100, 101, 1) == ARENA_ERR_OOM; + assert alloc_offset(100, 101, 1) == ARENA_NO_OFFSET; + // the largest i32 must not wrap around into "fits" + assert alloc_code(100, 2147483647, 1) == ARENA_ERR_OOM; + // padding to the alignment can itself push the start past the end + const r1 : AllocResult = arena_alloc(arena_new(10), 9, 1); + const r2 : AllocResult = arena_alloc(r1.arena, 1, 16); + assert r2.code == ARENA_ERR_OOM; + assert r2.arena.top == 9; + // an empty arena of capacity 0 refuses everything + assert alloc_code(0, 1, 1) == ARENA_ERR_OOM; +} + +test a_bad_size_or_alignment_is_a_code { + assert alloc_code(100, 0, 1) == ARENA_ERR_SIZE; + assert alloc_code(100, -5, 1) == ARENA_ERR_SIZE; + assert alloc_code(100, 4, 0) == ARENA_ERR_ALIGN; + assert alloc_code(100, 4, 3) == ARENA_ERR_ALIGN; + assert alloc_code(100, 4, -2) == ARENA_ERR_ALIGN; + assert alloc_code(100, 4, 12) == ARENA_ERR_ALIGN; + assert alloc_code(100, 4, 64) == ARENA_OK; + assert is_pow2(1); + assert is_pow2(2); + assert is_pow2(1024); + assert is_pow2(6) == false; + assert is_pow2(0) == false; +} + +test the_high_water_mark_survives_reset { + const a : Arena = three_allocs(); + assert a.high == 20; + const b : Arena = arena_reset(a); + assert b.top == 0; + assert b.high == 20; + assert b.resets == 1; + assert arena_left(b) == 100; + // after a reset the first allocation is at offset 0 again + const r : AllocResult = arena_alloc(b, 6, 1); + assert r.offset == 0; + // a smaller second run does not lower the mark + assert r.arena.high == 20; + // a larger one raises it + const big : AllocResult = arena_alloc(b, 64, 1); + assert big.arena.high == 64; + assert big.arena.resets == 1; +} + +test release_to_a_mark_frees_scratch_space { + const r1 : AllocResult = arena_alloc(arena_new(100), 10, 1); + const m : i32 = arena_mark(r1.arena); + assert m == 10; + const scratch : AllocResult = arena_alloc(r1.arena, 30, 1); + assert scratch.offset == 10; + assert scratch.arena.top == 40; + const back : AllocResult = arena_release(scratch.arena, m); + assert back.code == ARENA_OK; + assert back.arena.top == 10; + // the mark of the run stays at 40 + assert back.arena.high == 40; + // the freed space is handed out again at the same offset + const again : AllocResult = arena_alloc(back.arena, 30, 1); + assert again.offset == 10; + // a mark above top, or negative, is a code and changes nothing + const bad : AllocResult = arena_release(r1.arena, 11); + assert bad.code == ARENA_ERR_MARK; + assert bad.arena.top == 10; + assert arena_release(r1.arena, -1).code == ARENA_ERR_MARK; + // releasing to top is a no-op + assert arena_release(r1.arena, 10).arena.top == 10; +} + +test moments_are_two_zeroed_non_overlapping_buffers { + const a : Arena = arena_new(16); + var backing : [16]f32 = undefined; + fill_slots(&backing, 0, 16, 9.0); + const r : MomentsResult = moments_alloc(a, &backing, 4); + assert r.code == ARENA_OK; + assert r.m_off == 0; + assert r.v_off == 4; + assert r.arena.top == 8; + // zeroed where allocated, untouched elsewhere + assert backing[0] == 0.0; + assert backing[3] == 0.0; + assert backing[4] == 0.0; + assert backing[7] == 0.0; + assert backing[8] == 9.0; + // 2.5 + 7.5 over zeros: a missed zero (9.0) or an overlap would change the sum + assert moments_sum() == 10.0; +} + +test moments_are_both_or_neither { + // 7 slots cannot hold two buffers of 4: the first fits, the second does not + // code 1 (out of memory), top 0 + assert moments_outcome(7, 4) == 1000; + // an exact fit: code 0, top 8 + assert moments_outcome(8, 4) == 8; + // a size of 0 is a size error, code 2 + assert moments_outcome(16, 0) == 2000; + // and the backing array was not touched: 9.0 + 9.0 + assert backing_after_failed_moments() == 18.0; +} + +invariant top_never_exceeds_cap { + const a : Arena = three_allocs(); + assert a.top <= a.cap; + assert a.high >= a.top; + assert a.top >= 0; +} diff --git a/specs/ml/file_io_contract.t27 b/specs/ml/file_io_contract.t27 new file mode 100644 index 0000000000..07969972ec --- /dev/null +++ b/specs/ml/file_io_contract.t27 @@ -0,0 +1,837 @@ +// SPDX-License-Identifier: Apache-2.0 +; specs/ml/file_io_contract.t27 -- the file I/O contract t27c would have to implement, written as +; pure t27 over a byte array, plus a checkpoint format whose round trip is tested. +; Part of gHashTag/t27#10024 (our own model), cell 24 "Heap and file I/O in t27c" and cell 18 +; "Data loader and checkpoint I/O"; issue #10077. + +; WHY: a trainer has to save its weights and read them back, and read a token stream. t27 has no +; open, read, write or close, so there is nothing for a checkpoint spec to call. This file says +; what those four calls MEAN, as functions a test can run today, so that when t27c lowers them to +; the host (open/read/write/close on a real descriptor) there is an executable spec to hold the +; lowering against: the same tests, run against the real thing. + +; THE MODEL: a "volume" is a caller-owned byte array `data` cut into NFILES slots of FILE_CAP +; bytes, and `lens` holds each slot's length, -1 while the slot has never been written. A slot +; number stands for a path. An IoHandle is a value: slot, mode, position. Every call takes the +; handle and returns a new one in an IoResult together with a count and a code; nothing is +; global and nothing traps. Rules the tests pin down: +; - open for read of a slot that does not exist: IO_ERR_NOENT; open for write creates the slot +; or truncates it to length 0; +; - read returns min(n, len - pos) bytes; at the end it returns 0 bytes and IO_OK; +; - write is all or nothing: a write that would pass FILE_CAP is IO_ERR_FULL, writes no byte +; and leaves the length as it was; +; - read on a write handle and write on a read handle are IO_ERR_MODE; any call on a closed +; handle, including a second close, is IO_ERR_BADHANDLE. + +; THE CHECKPOINT (format "T27C", defined here, not the trainer's old .bin): 16 header bytes, then +; count f32 weights, then a 4-byte checksum. All integers are 4 bytes little endian. +; [0..4) magic 84 50 55 67 ("T27C") +; [4..8) version +; [8..12) step +; [12..16) count (number of f32 weights) +; [16..16+4*count) weights, each the f32 bit pattern (@bitCast), so a round trip is exact +; [16+4*count..+4) checksum: the sum of every byte before it +; An f32 goes through @bitCast to u32 and back; the lab shows t27c compiles that today, so the +; round trip needs no new float lowering. + +; THE TOKEN STREAM and THE MSLH HEADER (added after #10088): a token stream is u32 ids, 4 bytes +; little endian each, no header, read whole into a caller buffer with every size problem a code. +; The MSLH header is 16 bytes with the T27C field layout and the magic "MSLH" (77 83 76 72). No +; MSLH file format is defined anywhere else in this repository, so that header is what t27 would +; write and read, not a checked match for an existing file. + +; WHAT IS NOT CLAIMED: this is a model, not the host. It does not touch a disk; it does not model +; two handles on one slot, seek, rename, permissions, or a partial write after a crash (a write +; that fails part-way leaves the bytes it wrote, and a reader catches that through SHORT or SUM). +; Mapping the one HSLM checkpoint that survives (trinity data/wave9/worker-2/hslm_step_10000.bin) +; onto this format is a separate step and is not done here. +; phi^2 + 1/phi^2 = 3 | TRINITY + +module file_io_contract; + +pub const NFILES : i32 = 4; +pub const FILE_CAP : i32 = 256; +pub const VOLUME_BYTES : i32 = 1024; + +pub const MODE_CLOSED : i32 = 0; +pub const MODE_READ : i32 = 1; +pub const MODE_WRITE : i32 = 2; + +pub const IO_OK : i32 = 0; +pub const IO_ERR_NOENT : i32 = 1; +pub const IO_ERR_BADHANDLE : i32 = 2; +pub const IO_ERR_MODE : i32 = 3; +pub const IO_ERR_FULL : i32 = 4; +pub const IO_ERR_ARG : i32 = 5; + +pub const CKPT_OK : i32 = 0; +pub const CKPT_ERR_MAGIC : i32 = 20; +pub const CKPT_ERR_SHORT : i32 = 21; +pub const CKPT_ERR_SIZE : i32 = 22; +pub const CKPT_ERR_SUM : i32 = 23; + +pub const TOK_OK : i32 = 0; +pub const TOK_ERR_ALIGN : i32 = 30; +pub const TOK_ERR_SIZE : i32 = 31; + +pub const MSLH_OK : i32 = 0; +pub const MSLH_ERR_MAGIC : i32 = 40; +pub const MSLH_ERR_SHORT : i32 = 41; +pub const MSLH_ERR_COUNT : i32 = 42; + +pub const HEADER_BYTES : i32 = 16; +pub const MAGIC0 : u8 = 84; +pub const MAGIC1 : u8 = 50; +pub const MAGIC2 : u8 = 55; +pub const MAGIC3 : u8 = 67; + +pub const IoHandle = struct { + slot : i32, + mode : i32, + pos : i32, +}; + +pub const IoResult = struct { + handle : IoHandle, + n : i32, + code : i32, +}; + +pub const CkptInfo = struct { + code : i32, + version : i32, + step : i32, + count : i32, +}; + +pub const TokenInfo = struct { + code : i32, + count : i32, +}; + +pub const MslhHeader = struct { + code : i32, + version : i32, + step : i32, + count : i32, +}; + +fn closed_handle() -> IoHandle { + return IoHandle{ .slot = -1, .mode = MODE_CLOSED, .pos = 0 }; +} + +fn io_fail(h: IoHandle, code: i32) -> IoResult { + return IoResult{ .handle = h, .n = 0, .code = code }; +} + +/// A new volume: every slot absent. +pub fn vol_init(lens: []i32) -> void { + var i : usize = 0; + while (i < lens.len) { + lens[i] = -1; + i = i + 1; + } +} + +pub fn io_open(lens: []i32, slot: i32, mode: i32) -> IoResult { + if (slot < 0) { + return io_fail(closed_handle(), IO_ERR_ARG); + } + if (slot >= NFILES) { + return io_fail(closed_handle(), IO_ERR_ARG); + } + if (mode == MODE_READ) { + if (lens[@intCast(slot)] < 0) { + return io_fail(closed_handle(), IO_ERR_NOENT); + } + return IoResult{ .handle = IoHandle{ .slot = slot, .mode = MODE_READ, .pos = 0 }, .n = 0, .code = IO_OK }; + } + if (mode == MODE_WRITE) { + lens[@intCast(slot)] = 0; + return IoResult{ .handle = IoHandle{ .slot = slot, .mode = MODE_WRITE, .pos = 0 }, .n = 0, .code = IO_OK }; + } + return io_fail(closed_handle(), IO_ERR_ARG); +} + +/// Append n bytes of src at the handle's position. All or nothing. +pub fn io_write(data: []u8, lens: []i32, h: IoHandle, src: []const u8, n: i32) -> IoResult { + if (h.mode == MODE_CLOSED) { + return io_fail(h, IO_ERR_BADHANDLE); + } + if (h.mode != MODE_WRITE) { + return io_fail(h, IO_ERR_MODE); + } + if (n < 0) { + return io_fail(h, IO_ERR_ARG); + } + if (n > FILE_CAP - h.pos) { + return io_fail(h, IO_ERR_FULL); + } + const base : i32 = h.slot * FILE_CAP + h.pos; + var i : i32 = 0; + while (i < n) { + data[@intCast(base + i)] = src[@intCast(i)]; + i = i + 1; + } + lens[@intCast(h.slot)] = h.pos + n; + return IoResult{ .handle = IoHandle{ .slot = h.slot, .mode = MODE_WRITE, .pos = h.pos + n }, .n = n, .code = IO_OK }; +} + +/// Read up to n bytes into dst: min(n, len - pos). At the end it returns 0 bytes and IO_OK. +pub fn io_read(data: []u8, lens: []i32, h: IoHandle, dst: []u8, n: i32) -> IoResult { + if (h.mode == MODE_CLOSED) { + return io_fail(h, IO_ERR_BADHANDLE); + } + if (h.mode != MODE_READ) { + return io_fail(h, IO_ERR_MODE); + } + if (n < 0) { + return io_fail(h, IO_ERR_ARG); + } + var take : i32 = lens[@intCast(h.slot)] - h.pos; + if (take > n) { + take = n; + } + const base : i32 = h.slot * FILE_CAP + h.pos; + var i : i32 = 0; + while (i < take) { + dst[@intCast(i)] = data[@intCast(base + i)]; + i = i + 1; + } + return IoResult{ .handle = IoHandle{ .slot = h.slot, .mode = MODE_READ, .pos = h.pos + take }, .n = take, .code = IO_OK }; +} + +pub fn io_close(h: IoHandle) -> IoResult { + if (h.mode == MODE_CLOSED) { + return io_fail(h, IO_ERR_BADHANDLE); + } + return IoResult{ .handle = closed_handle(), .n = 0, .code = IO_OK }; +} + +pub fn io_size(lens: []const i32, slot: i32) -> i32 { + return lens[@intCast(slot)]; +} + +// ---- 4-byte little endian integers ---- + +pub fn put_u32(buf: []u8, off: i32, v: u32) -> void { + var x : u32 = v; + var k : i32 = 0; + while (k < 4) { + buf[@intCast(off + k)] = @intCast(@mod(x, 256)); + x = @divTrunc(x, 256); + k = k + 1; + } +} + +pub fn get_u32(buf: []const u8, off: i32) -> u32 { + var x : u32 = 0; + var k : i32 = 3; + while (k >= 0) { + const b : u32 = @intCast(buf[@intCast(off + k)]); + x = x * 256 + b; + k = k - 1; + } + return x; +} + +fn byte_sum(buf: []const u8, n: i32) -> u32 { + var s : u32 = 0; + var i : i32 = 0; + while (i < n) { + const b : u32 = @intCast(buf[@intCast(i)]); + s = s + b; + i = i + 1; + } + return s; +} + +// ---- the checkpoint ---- + +/// Write a checkpoint of the first count weights into slot. Returns CKPT_OK, or the IO code of +/// the call that failed (the slot is then left with the bytes written so far). +pub fn ckpt_write(data: []u8, lens: []i32, slot: i32, version: i32, step: i32, weights: []const f32, count: i32) -> i32 { + const o : IoResult = io_open(lens, slot, MODE_WRITE); + if (o.code != IO_OK) { + return o.code; + } + var h : IoHandle = o.handle; + var hdr : [16]u8 = undefined; + hdr[0] = MAGIC0; + hdr[1] = MAGIC1; + hdr[2] = MAGIC2; + hdr[3] = MAGIC3; + put_u32(&hdr, 4, @intCast(version)); + put_u32(&hdr, 8, @intCast(step)); + put_u32(&hdr, 12, @intCast(count)); + var sum : u32 = byte_sum(&hdr, HEADER_BYTES); + const wh : IoResult = io_write(data, lens, h, &hdr, HEADER_BYTES); + if (wh.code != IO_OK) { + return wh.code; + } + h = wh.handle; + var word : [4]u8 = undefined; + var i : i32 = 0; + while (i < count) { + const bits : u32 = @bitCast(weights[@intCast(i)]); + put_u32(&word, 0, bits); + sum = sum + byte_sum(&word, 4); + const ww : IoResult = io_write(data, lens, h, &word, 4); + if (ww.code != IO_OK) { + return ww.code; + } + h = ww.handle; + i = i + 1; + } + put_u32(&word, 0, sum); + const ws : IoResult = io_write(data, lens, h, &word, 4); + if (ws.code != IO_OK) { + return ws.code; + } + const c : IoResult = io_close(ws.handle); + return c.code; +} + +fn info_fail(code: i32) -> CkptInfo { + return CkptInfo{ .code = code, .version = 0, .step = 0, .count = 0 }; +} + +/// Read the checkpoint in slot into out (room for max weights). On any error out may hold the +/// weights read so far and the code says why; on CKPT_OK the file was complete and its checksum matched. +pub fn ckpt_read(data: []u8, lens: []i32, slot: i32, out: []f32, max: i32) -> CkptInfo { + const o : IoResult = io_open(lens, slot, MODE_READ); + if (o.code != IO_OK) { + return info_fail(o.code); + } + var hdr : [16]u8 = undefined; + const rh : IoResult = io_read(data, lens, o.handle, &hdr, HEADER_BYTES); + if (rh.code != IO_OK) { + return info_fail(rh.code); + } + if (rh.n < HEADER_BYTES) { + return info_fail(CKPT_ERR_SHORT); + } + if (hdr[0] != MAGIC0 or hdr[1] != MAGIC1 or hdr[2] != MAGIC2 or hdr[3] != MAGIC3) { + return info_fail(CKPT_ERR_MAGIC); + } + const count_u : u32 = get_u32(&hdr, 12); + const max_u : u32 = @intCast(max); + if (count_u > max_u) { + return info_fail(CKPT_ERR_SIZE); + } + const count : i32 = @intCast(count_u); + var sum : u32 = byte_sum(&hdr, HEADER_BYTES); + var h : IoHandle = rh.handle; + var word : [4]u8 = undefined; + var i : i32 = 0; + while (i < count) { + const rw : IoResult = io_read(data, lens, h, &word, 4); + if (rw.code != IO_OK) { + return info_fail(rw.code); + } + if (rw.n < 4) { + return info_fail(CKPT_ERR_SHORT); + } + h = rw.handle; + sum = sum + byte_sum(&word, 4); + out[@intCast(i)] = @bitCast(get_u32(&word, 0)); + i = i + 1; + } + const rs : IoResult = io_read(data, lens, h, &word, 4); + if (rs.code != IO_OK) { + return info_fail(rs.code); + } + if (rs.n < 4) { + return info_fail(CKPT_ERR_SHORT); + } + if (get_u32(&word, 0) != sum) { + return info_fail(CKPT_ERR_SUM); + } + const c : IoResult = io_close(rs.handle); + if (c.code != IO_OK) { + return info_fail(c.code); + } + return CkptInfo{ .code = CKPT_OK, .version = @intCast(get_u32(&hdr, 4)), .step = @intCast(get_u32(&hdr, 8)), .count = count }; +} + +// ---- a token stream: u32 token ids, 4 bytes little endian each, no header ---- + +/// Write the first n token ids into slot. Returns TOK_OK or the IO code that failed (a stream +/// that does not fit is IO_ERR_FULL and the slot keeps the whole tokens written before it). +pub fn tokens_write(data: []u8, lens: []i32, slot: i32, toks: []const u32, n: i32) -> i32 { + const o : IoResult = io_open(lens, slot, MODE_WRITE); + if (o.code != IO_OK) { + return o.code; + } + var h : IoHandle = o.handle; + var word : [4]u8 = undefined; + var i : i32 = 0; + while (i < n) { + put_u32(&word, 0, toks[@intCast(i)]); + const w : IoResult = io_write(data, lens, h, &word, 4); + if (w.code != IO_OK) { + return w.code; + } + h = w.handle; + i = i + 1; + } + const c : IoResult = io_close(h); + return c.code; +} + +/// Read the whole stream in slot into out (room for max ids). A file whose length is not a +/// multiple of 4 is TOK_ERR_ALIGN, one holding more than max ids is TOK_ERR_SIZE; out is not +/// touched in either case. +pub fn tokens_read(data: []u8, lens: []i32, slot: i32, out: []u32, max: i32) -> TokenInfo { + const o : IoResult = io_open(lens, slot, MODE_READ); + if (o.code != IO_OK) { + return TokenInfo{ .code = o.code, .count = 0 }; + } + const size : i32 = io_size(lens, slot); + if (@mod(size, 4) != 0) { + return TokenInfo{ .code = TOK_ERR_ALIGN, .count = 0 }; + } + const count : i32 = @divTrunc(size, 4); + if (count > max) { + return TokenInfo{ .code = TOK_ERR_SIZE, .count = 0 }; + } + var h : IoHandle = o.handle; + var word : [4]u8 = undefined; + var i : i32 = 0; + while (i < count) { + const r : IoResult = io_read(data, lens, h, &word, 4); + if (r.code != IO_OK) { + return TokenInfo{ .code = r.code, .count = 0 }; + } + h = r.handle; + out[@intCast(i)] = get_u32(&word, 0); + i = i + 1; + } + const c : IoResult = io_close(h); + return TokenInfo{ .code = c.code, .count = count }; +} + +// ---- the MSLH header: magic "MSLH", version, step, count (parameters that follow) ---- +// Defined here as a 16-byte header with the same field layout as T27C and a different magic. +// Nothing in this repository defines an MSLH file, so this is the header t27 would write and +// read, NOT a verified match for any existing MSLH file. + +pub fn mslh_header_put(buf: []u8, version: i32, step: i32, count: i32) -> void { + buf[0] = 77; + buf[1] = 83; + buf[2] = 76; + buf[3] = 72; + put_u32(buf, 4, @intCast(version)); + put_u32(buf, 8, @intCast(step)); + put_u32(buf, 12, @intCast(count)); +} + +/// Parse a header already read into buf (n bytes were available). +pub fn mslh_header_get(buf: []const u8, n: i32) -> MslhHeader { + if (n < HEADER_BYTES) { + return MslhHeader{ .code = MSLH_ERR_SHORT, .version = 0, .step = 0, .count = 0 }; + } + if (buf[0] != 77 or buf[1] != 83 or buf[2] != 76 or buf[3] != 72) { + return MslhHeader{ .code = MSLH_ERR_MAGIC, .version = 0, .step = 0, .count = 0 }; + } + const count_u : u32 = get_u32(buf, 12); + if (count_u > 1073741823) { + return MslhHeader{ .code = MSLH_ERR_COUNT, .version = 0, .step = 0, .count = 0 }; + } + return MslhHeader{ .code = MSLH_OK, .version = @intCast(get_u32(buf, 4)), .step = @intCast(get_u32(buf, 8)), .count = @intCast(count_u) }; +} + +/// Write just the header into slot (a loader reads it first to size its buffers). +pub fn mslh_write_header(data: []u8, lens: []i32, slot: i32, version: i32, step: i32, count: i32) -> i32 { + const o : IoResult = io_open(lens, slot, MODE_WRITE); + if (o.code != IO_OK) { + return o.code; + } + var hdr : [16]u8 = undefined; + mslh_header_put(&hdr, version, step, count); + const w : IoResult = io_write(data, lens, o.handle, &hdr, HEADER_BYTES); + if (w.code != IO_OK) { + return w.code; + } + const c : IoResult = io_close(w.handle); + return c.code; +} + +pub fn mslh_read_header(data: []u8, lens: []i32, slot: i32) -> MslhHeader { + const o : IoResult = io_open(lens, slot, MODE_READ); + if (o.code != IO_OK) { + return MslhHeader{ .code = o.code, .version = 0, .step = 0, .count = 0 }; + } + var hdr : [16]u8 = undefined; + const r : IoResult = io_read(data, lens, o.handle, &hdr, HEADER_BYTES); + if (r.code != IO_OK) { + return MslhHeader{ .code = r.code, .version = 0, .step = 0, .count = 0 }; + } + return mslh_header_get(&hdr, r.n); +} + +// ---- helpers that exist only so the tests go through function calls ---- + +fn zero_volume(data: []u8, lens: []i32) -> void { + var i : usize = 0; + while (i < data.len) { + data[i] = 0; + i = i + 1; + } + vol_init(lens); +} + +/// Twelve weights chosen to be awkward: fractions that are not exact in binary, a negative, a +/// tiny and a huge value, zero, and a value near the f32 maximum. +fn fill_weights(w: []f32) -> void { + const src : [12]f32 = [0.1, -2.5, 1.0e-8, 3.0e30, 0.0, 0.333333, -0.7, 1.5, 65504.0, -1.0e-30, 123456.789, 3.4e38]; + var i : usize = 0; + while (i < 12) { + w[i] = src[i]; + i = i + 1; + } +} + +/// How many of the first n weights differ in BIT PATTERN between a and b. 0 means an exact copy. +fn bit_mismatches(a: []const f32, b: []const f32, n: i32) -> i32 { + var bad : i32 = 0; + var i : i32 = 0; + while (i < n) { + const x : u32 = @bitCast(a[@intCast(i)]); + const y : u32 = @bitCast(b[@intCast(i)]); + if (x != y) { + bad = bad + 1; + } + i = i + 1; + } + return bad; +} + +/// Write five bytes, read 3 then 10 then 10 again; returns 100*first + 10*second + third, the +/// byte counts, so the end-of-file behaviour is one number: 3, 2, 0 -> 320. +fn read_pattern() -> i32 { + var data : [1024]u8 = undefined; + var lens : [4]i32 = undefined; + zero_volume(&data, &lens); + var src : [5]u8 = [1, 2, 3, 4, 5]; + var dst : [10]u8 = undefined; + const o : IoResult = io_open(&lens, 2, MODE_WRITE); + const w : IoResult = io_write(&data, &lens, o.handle, &src, 5); + const closed : IoResult = io_close(w.handle); + const r : IoResult = io_open(&lens, 2, MODE_READ); + const a : IoResult = io_read(&data, &lens, r.handle, &dst, 3); + const b : IoResult = io_read(&data, &lens, a.handle, &dst, 10); + const c : IoResult = io_read(&data, &lens, b.handle, &dst, 10); + return 100 * a.n + 10 * b.n + c.n; +} + +test little_endian_words_round_trip { + var buf : [8]u8 = undefined; + put_u32(&buf, 2, 305419896); + // 0x12345678 little endian + assert buf[2] == 120; + assert buf[3] == 86; + assert buf[4] == 52; + assert buf[5] == 18; + assert get_u32(&buf, 2) == 305419896; + put_u32(&buf, 0, 4294967295); + assert get_u32(&buf, 0) == 4294967295; + put_u32(&buf, 4, 0); + assert get_u32(&buf, 4) == 0; +} + +test open_read_write_close_is_a_state_machine { + var data : [1024]u8 = undefined; + var lens : [4]i32 = undefined; + zero_volume(&data, &lens); + // an absent slot cannot be opened for reading + assert io_open(&lens, 1, MODE_READ).code == IO_ERR_NOENT; + // a slot outside the volume, and an unknown mode, are argument errors + assert io_open(&lens, 4, MODE_WRITE).code == IO_ERR_ARG; + assert io_open(&lens, -1, MODE_WRITE).code == IO_ERR_ARG; + assert io_open(&lens, 1, 7).code == IO_ERR_ARG; + // open for write creates the slot, empty + const o : IoResult = io_open(&lens, 1, MODE_WRITE); + assert o.code == IO_OK; + assert o.handle.mode == MODE_WRITE; + assert o.handle.pos == 0; + assert io_size(&lens, 1) == 0; + var src : [3]u8 = [7, 8, 9]; + const w : IoResult = io_write(&data, &lens, o.handle, &src, 3); + assert w.code == IO_OK; + assert w.n == 3; + assert w.handle.pos == 3; + assert io_size(&lens, 1) == 3; + // the bytes sit in slot 1: offset 1 * FILE_CAP + assert data[256] == 7; + assert data[258] == 9; + assert data[0] == 0; + const c : IoResult = io_close(w.handle); + assert c.code == IO_OK; + assert c.handle.mode == MODE_CLOSED; +} + +test read_stops_at_the_end_without_an_error { + // 5 bytes written; reads of 3, 10, 10 return 3, 2, 0 + assert read_pattern() == 320; +} + +test the_wrong_mode_or_a_closed_handle_is_an_error { + var data : [1024]u8 = undefined; + var lens : [4]i32 = undefined; + zero_volume(&data, &lens); + var src : [2]u8 = [1, 2]; + var dst : [2]u8 = undefined; + const ow : IoResult = io_open(&lens, 0, MODE_WRITE); + // read on a write handle + assert io_read(&data, &lens, ow.handle, &dst, 1).code == IO_ERR_MODE; + const cw : IoResult = io_close(ow.handle); + // a closed handle, and a second close + assert io_write(&data, &lens, cw.handle, &src, 1).code == IO_ERR_BADHANDLE; + assert io_close(cw.handle).code == IO_ERR_BADHANDLE; + assert io_read(&data, &lens, cw.handle, &dst, 1).code == IO_ERR_BADHANDLE; + // write on a read handle + const orr : IoResult = io_open(&lens, 0, MODE_READ); + assert orr.code == IO_OK; + assert io_write(&data, &lens, orr.handle, &src, 1).code == IO_ERR_MODE; + // a negative count is an argument error + assert io_read(&data, &lens, orr.handle, &dst, -1).code == IO_ERR_ARG; +} + +test a_write_past_the_slot_is_all_or_nothing { + var data : [1024]u8 = undefined; + var lens : [4]i32 = undefined; + zero_volume(&data, &lens); + var chunk : [200]u8 = undefined; + var i : usize = 0; + while (i < 200) { + chunk[i] = 5; + i = i + 1; + } + const o : IoResult = io_open(&lens, 0, MODE_WRITE); + const w1 : IoResult = io_write(&data, &lens, o.handle, &chunk, 200); + assert w1.code == IO_OK; + // 200 + 57 > 256: refused, nothing written, length unchanged + const w2 : IoResult = io_write(&data, &lens, w1.handle, &chunk, 57); + assert w2.code == IO_ERR_FULL; + assert w2.n == 0; + assert io_size(&lens, 0) == 200; + assert data[200] == 0; + // exactly the 56 bytes that remain fit + const w3 : IoResult = io_write(&data, &lens, w1.handle, &chunk, 56); + assert w3.code == IO_OK; + assert io_size(&lens, 0) == 256; + // and slot 1 was not touched by the first slot filling up + assert data[256] == 0; +} + +test a_checkpoint_round_trips_exactly { + var data : [1024]u8 = undefined; + var lens : [4]i32 = undefined; + zero_volume(&data, &lens); + var w : [12]f32 = undefined; + fill_weights(&w); + const code : i32 = ckpt_write(&data, &lens, 0, 3, 10000, &w, 12); + assert code == CKPT_OK; + // 16 header + 12 * 4 weights + 4 checksum + assert io_size(&lens, 0) == 68; + // the header is in the file: magic "T27C", then version 3 + assert data[0] == 84; + assert data[3] == 67; + assert data[4] == 3; + var back : [12]f32 = undefined; + const info : CkptInfo = ckpt_read(&data, &lens, 0, &back, 12); + assert info.code == CKPT_OK; + assert info.version == 3; + assert info.step == 10000; + assert info.count == 12; + assert bit_mismatches(&w, &back, 12) == 0; + // a spot value, to show the comparison is not trivially true + assert back[1] == -2.5; + assert back[10] == 123456.789; +} + +test the_bit_comparison_can_fail { + // a negative control: a one-bit difference is seen + var a : [12]f32 = undefined; + var b : [12]f32 = undefined; + fill_weights(&a); + fill_weights(&b); + assert bit_mismatches(&a, &b, 12) == 0; + b[4] = -0.0; + // 0.0 == -0.0 as numbers, but the bit patterns differ + assert b[4] == a[4]; + assert bit_mismatches(&a, &b, 12) == 1; + b[7] = 1.5000001; + assert bit_mismatches(&a, &b, 12) == 2; +} + +test a_second_write_truncates_and_slots_are_independent { + var data : [1024]u8 = undefined; + var lens : [4]i32 = undefined; + zero_volume(&data, &lens); + var w : [12]f32 = undefined; + fill_weights(&w); + assert ckpt_write(&data, &lens, 0, 1, 100, &w, 12) == CKPT_OK; + assert ckpt_write(&data, &lens, 1, 2, 200, &w, 5) == CKPT_OK; + // overwrite slot 0 with three weights: the file is now 16 + 12 + 4 = 32 bytes + assert ckpt_write(&data, &lens, 0, 9, 300, &w, 3) == CKPT_OK; + assert io_size(&lens, 0) == 32; + var back : [12]f32 = undefined; + const i0 : CkptInfo = ckpt_read(&data, &lens, 0, &back, 12); + assert i0.code == CKPT_OK; + assert i0.version == 9; + assert i0.step == 300; + assert i0.count == 3; + assert bit_mismatches(&w, &back, 3) == 0; + const i1 : CkptInfo = ckpt_read(&data, &lens, 1, &back, 12); + assert i1.code == CKPT_OK; + assert i1.version == 2; + assert i1.step == 200; + assert i1.count == 5; + assert bit_mismatches(&w, &back, 5) == 0; +} + +test a_damaged_checkpoint_is_refused { + var data : [1024]u8 = undefined; + var lens : [4]i32 = undefined; + zero_volume(&data, &lens); + var w : [12]f32 = undefined; + fill_weights(&w); + assert ckpt_write(&data, &lens, 0, 1, 100, &w, 12) == CKPT_OK; + var back : [12]f32 = undefined; + // a flipped payload byte: checksum mismatch + data[20] = data[20] + 1; + assert ckpt_read(&data, &lens, 0, &back, 12).code == CKPT_ERR_SUM; + data[20] = data[20] - 1; + assert ckpt_read(&data, &lens, 0, &back, 12).code == CKPT_OK; + // a flipped checksum byte + data[64] = data[64] + 1; + assert ckpt_read(&data, &lens, 0, &back, 12).code == CKPT_ERR_SUM; + data[64] = data[64] - 1; + // a bad magic + data[1] = 0; + assert ckpt_read(&data, &lens, 0, &back, 12).code == CKPT_ERR_MAGIC; + data[1] = 50; + // a truncated file: one byte short of the checksum + lens[0] = 67; + assert ckpt_read(&data, &lens, 0, &back, 12).code == CKPT_ERR_SHORT; + // truncated inside the header + lens[0] = 10; + assert ckpt_read(&data, &lens, 0, &back, 12).code == CKPT_ERR_SHORT; + lens[0] = 68; + // room for fewer weights than the file holds + assert ckpt_read(&data, &lens, 0, &back, 11).code == CKPT_ERR_SIZE; + // a slot that was never written + assert ckpt_read(&data, &lens, 3, &back, 12).code == IO_ERR_NOENT; + assert ckpt_read(&data, &lens, 0, &back, 12).code == CKPT_OK; +} + +test a_checkpoint_too_big_for_a_slot_is_a_code { + var data : [1024]u8 = undefined; + var lens : [4]i32 = undefined; + zero_volume(&data, &lens); + // a slot holds 256 bytes: 16 header + 4 per weight + 4 checksum + var big : [64]f32 = undefined; + var i : usize = 0; + while (i < 64) { + big[i] = 1.0; + i = i + 1; + } + // 59 weights: 16 + 236 + 4 = 256 exactly fits + assert ckpt_write(&data, &lens, 2, 1, 1, &big, 59) == CKPT_OK; + assert io_size(&lens, 2) == 256; + // 60 weights: 260 bytes, the checksum word does not fit + assert ckpt_write(&data, &lens, 2, 1, 1, &big, 60) == IO_ERR_FULL; + var back : [64]f32 = undefined; + // the half-written slot is not a valid checkpoint + assert ckpt_read(&data, &lens, 2, &back, 64).code != CKPT_OK; +} + +test a_token_stream_round_trips_and_bounds_are_codes { + var data : [1024]u8 = undefined; + var lens : [4]i32 = undefined; + zero_volume(&data, &lens); + const toks : [6]u32 = [0, 1, 50256, 4294967295, 7, 305419896]; + assert tokens_write(&data, &lens, 1, &toks, 6) == TOK_OK; + assert io_size(&lens, 1) == 24; + // token 2 (50256 = 0xC450) is at bytes 8..12, little endian + assert data[256 + 8] == 80; + assert data[256 + 9] == 196; + var back : [8]u32 = undefined; + const info : TokenInfo = tokens_read(&data, &lens, 1, &back, 8); + assert info.code == TOK_OK; + assert info.count == 6; + assert back[0] == 0; + assert back[2] == 50256; + assert back[3] == 4294967295; + assert back[5] == 305419896; + // room for fewer ids than the file holds: a code, and out is untouched + back[0] = 99; + const small : TokenInfo = tokens_read(&data, &lens, 1, &back, 5); + assert small.code == TOK_ERR_SIZE; + assert small.count == 0; + assert back[0] == 99; + // a file whose length is not a whole number of ids + lens[1] = 23; + assert tokens_read(&data, &lens, 1, &back, 8).code == TOK_ERR_ALIGN; + lens[1] = 24; + // a slot that was never written + assert tokens_read(&data, &lens, 3, &back, 8).code == IO_ERR_NOENT; + // an empty stream is valid and has zero ids + assert tokens_write(&data, &lens, 2, &toks, 0) == TOK_OK; + assert tokens_read(&data, &lens, 2, &back, 8).count == 0; + // 64 ids fill a slot exactly; the 65th does not fit + var many : [65]u32 = undefined; + var i : usize = 0; + while (i < 65) { + const id : u32 = @intCast(i); + many[i] = 1000 + id; + i = i + 1; + } + assert tokens_write(&data, &lens, 0, &many, 64) == TOK_OK; + assert io_size(&lens, 0) == 256; + assert tokens_write(&data, &lens, 0, &many, 65) == IO_ERR_FULL; + // the failed write kept the 64 whole ids it had already written + assert io_size(&lens, 0) == 256; +} + +test an_mslh_header_round_trips_and_other_magics_are_refused { + var data : [1024]u8 = undefined; + var lens : [4]i32 = undefined; + zero_volume(&data, &lens); + assert mslh_write_header(&data, &lens, 0, 2, 4000, 1234567) == MSLH_OK; + assert io_size(&lens, 0) == 16; + // the magic is in the file + assert data[0] == 77; + assert data[1] == 83; + assert data[2] == 76; + assert data[3] == 72; + const h : MslhHeader = mslh_read_header(&data, &lens, 0); + assert h.code == MSLH_OK; + assert h.version == 2; + assert h.step == 4000; + assert h.count == 1234567; + // a T27C checkpoint is not an MSLH file, and the reverse + var w : [4]f32 = [1.0, 2.0, 3.0, 4.0]; + assert ckpt_write(&data, &lens, 1, 1, 1, &w, 4) == CKPT_OK; + assert mslh_read_header(&data, &lens, 1).code == MSLH_ERR_MAGIC; + var back : [4]f32 = undefined; + assert ckpt_read(&data, &lens, 0, &back, 4).code == CKPT_ERR_MAGIC; + // an absent slot + assert mslh_read_header(&data, &lens, 3).code == IO_ERR_NOENT; + // a header cut short + lens[0] = 15; + assert mslh_read_header(&data, &lens, 0).code == MSLH_ERR_SHORT; + lens[0] = 16; + // a parameter count no buffer could hold + put_u32(&data, 12, 4294967295); + assert mslh_read_header(&data, &lens, 0).code == MSLH_ERR_COUNT; + put_u32(&data, 12, 1073741823); + assert mslh_read_header(&data, &lens, 0).count == 1073741823; +} + +invariant a_checkpoint_has_a_fixed_overhead { + assert HEADER_BYTES + 4 == 20; + assert NFILES * FILE_CAP == VOLUME_BYTES; +} diff --git a/specs/ml/optimizer/adamw.t27 b/specs/ml/optimizer/adamw.t27 index e37bd73d91..f9c5e39eb1 100644 --- a/specs/ml/optimizer/adamw.t27 +++ b/specs/ml/optimizer/adamw.t27 @@ -1,1111 +1,452 @@ // SPDX-License-Identifier: Apache-2.0 // t27/specs/ml/optimizer/adamw.t27 -// AdamW - Adam with decoupled weight decay and phi-optimizations | phi^2 + 1/phi^2 = 3 | TRINITY +// AdamW - Adam with decoupled weight decay | phi^2 + 1/phi^2 = 3 | TRINITY + +// Why this file was rewritten (epic gHashTag/t27#10024, own_model cell 24): +// +// Before: 1100 lines whose function bodies were comments, a second copy of every +// declaration under a _V1 suffix, tests that were prose (`then elapsed_time_ns < 10`), +// and `init_v1` calling `allocate_array`, which t27c does not have. On the Railway lab +// `t27c test-report` answered BLOCKED: undeclared identifier 'allocate_array'. Removing +// that call alone was not enough: the next error was in the prose tests. +// +// Now: AdamW as arithmetic that runs, in the shape of adam.t27 (#8980), with the three +// moment buffers (first moment m, second moment v, AMSGrad maximum x) living in ONE +// caller-owned f32 array at offsets. Where the offsets come from is the caller's choice: +// a fixed array, or ml/arena_heap.t27, which is the t27 answer to "no heap": a bump +// allocator over that array. `state_alloc` below is the arena path. +// +// What was dropped, and why (nothing true was lost, and nothing here pretends to keep it): +// - the GF16 types: gf16 arithmetic had no working bodies behind them; this is f32. +// - the _V1 duplicate block and its benches: every body was a comment. +// - run_adamw / b1_control_ablation_produces_distinct_bpb: needs a trainer and a data +// set, not an optimizer. The four arms stay as variant_config(0..3); the experiment +// itself belongs to the trainer spec, and its falsification rule (a phi-free control +// matching the phi arm within Monte-Carlo error on equal budget FALSIFIES phi +// specificity) is kept verbatim below. +// +// Non-claims: this proves the update rule on small exact cases and a quadratic. It does not +// say AdamW trains an HSLM; it does not say any phi value is better than a control. module AdamW; - use base::types; - use math::constants; - use numeric::gf16; - - // NOTE (naming): this file carries two versions of the optimizer spec -- - // the pre-B1 spec below and the Phase B1 block appended after the - // Mathematical Notes. Both declared the same constants, types and - // functions, so the module declared those names twice. The pre-B1 - // declarations that the B1 block re-declares carry a _V1 / V1 / _v1 - // suffix (PHI_V1, AdamWConfigV1, init_v1, ...), and the pre-B1 tests, - // invariants and benches call them by those names. The B1 block keeps - // the plain names. Nothing was removed. - - // =========================================================== - // 1. Constants - // =========================================================== - - const PHI_V1 : gf16::GF16 = 1.618033988749895; - const DEFAULT_LEARNING_RATE_V1 : gf16::GF16 = 1e-3; - const DEFAULT_BETA1_V1 : gf16::GF16 = 0.9; - const DEFAULT_BETA2_V1 : gf16::GF16 = 0.999; - const PHI_BETA1 : gf16::GF16 = 0.9 / PHI_V1; // ~ 0.556, phi-damped beta1 - const PHI_BETA2 : gf16::GF16 = 0.999 / PHI_V1; // ~ 0.617, phi-damped beta2 - const DEFAULT_WEIGHT_DECAY_V1 : gf16::GF16 = 0.01; - const DEFAULT_EPSILON_V1 : gf16::GF16 = 1e-8; - const DEFAULT_AMSGRAD_V1 : bool = false; - - // =========================================================== - // 2. Types - // =========================================================== - - pub const AdamWConfigV1 = struct { - learning_rate : gf16::GF16, // eta: step size - beta1 : gf16::GF16, // beta1: first moment decay - beta2 : gf16::GF16, // beta2: second moment decay - weight_decay : gf16::GF16, // lambda: decoupled weight decay - epsilon : gf16::GF16, // eps: numerical stability - amsgrad : bool, // Use AMSGrad variant - use_phi_betas : bool, // Use phi-damped beta1, beta2 - }; - - pub const AdamWStateV1 = struct { - m : []gf16::GF16, // First moment (moving average of gradients) - v : []gf16::GF16, // Second moment (moving average of squared gradients) - v_max : []gf16::GF16, // Max of v (for AMSGrad) - param_count : u32, // Number of parameters - step : u64, // Current optimization step (t) - }; - - pub const OptimizerStepResultV1 = struct { - updated_params : []gf16::GF16, // New parameter values - m : []gf16::GF16, // Updated first moments - v : []gf16::GF16, // Updated second moments - step_norm : gf16::GF16, // L2 norm of the parameter update - }; - - // =========================================================== - // 3. Core Functions - // =========================================================== - - // init_v1(config: AdamWConfigV1, param_count: u32) -> AdamWStateV1 -// Initializes AdamW state with zero moments. - fn init_v1(config: AdamWConfigV1, param_count: u32) -> AdamWStateV1 { - var m = allocate_array(param_count); - var v = allocate_array(param_count); - var v_max = allocate_array(param_count); - var i = 0; - while (i < param_count) { - m[i] = gf16::GF16.from_f32(0.0); - v[i] = gf16::GF16.from_f32(0.0); - v_max[i] = gf16::GF16.from_f32(0.0); - i = i + 1; - } - return AdamWStateV1{ .m = m, .v = v, .v_max = v_max, .param_count = param_count, .step = 0 }; - } - - // step_v1(state: AdamWStateV1, params: []gf16::GF16, grads: []gf16::GF16) -> OptimizerStepResultV1 - // Performs one AdamW optimization step. - fn step_v1(state: AdamWStateV1, params: []gf16::GF16, grads: []gf16::GF16) -> OptimizerStepResultV1 { - // t = step + 1 - // bias_correction1 = 1 - beta1^t - // bias_correction2 = 1 - beta2^t - // lr_t = lr * sqrt(bias_correction2) / bias_correction1 - // - // Apply weight decay (decoupled): - // params[i] = params[i] * (1 - lr * weight_decay) - // - // Update moments: - // m_t = beta1 * m_{t-1} + (1 - beta1) * g_t - // v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2 - // - // If amsgrad: - // v_max_t = max(v_max_{t-1}, v_t) - // denominator = sqrt(v_max_t) + eps - // else: - // denominator = sqrt(v_t) + eps - // - // Update parameters: - // params[i] = params[i] - lr_t * m_t / denominator +use base::types; +use math::constants; +use ml::optimizer::adam; +use ml::arena_heap; + +// ---- constants ---- + +// Standard AdamW defaults (Kingma and Ba 2015; Loshchilov and Hutter 2019) +const DEFAULT_LEARNING_RATE : f32 = 0.001; +const DEFAULT_BETA1 : f32 = 0.9; +const DEFAULT_BETA2 : f32 = 0.999; +const DEFAULT_WEIGHT_DECAY : f32 = 0.01; +const DEFAULT_EPSILON : f32 = 0.00000001; + +// phi anchors. [Open conjecture]: no momentum theory singles out phi. The only verified +// phi fact here is the identity phi^2 + phi^-2 = 3 (Law L5), tested below. +const PHI : f32 = 1.618033988749895; +const PHI_INV : f32 = 0.618033988749895; +const PHI_INV_CUBED : f32 = 0.23606797749979; +const PHI_DAMPED_BETA1 : f32 = 0.5562305898749055; +const PHI_DAMPED_BETA2 : f32 = 0.617415954761145; + +// Control arm: non-Fibonacci rationals of comparable description cost (4 and 7 and 30 are +// outside {1, 2, 3, 5, 8, 13, 21, ...}); 5/8 was rejected as phi-biased. +const RATIONAL_BETA1 : f32 = 0.5714285714285714; +const RATIONAL_WEIGHT_DECAY : f32 = 0.2333333333333333; + +// The four ablation arms of Phase B1 (epic #181). Codes, not an enum, so the tests stay plain. +const VARIANT_CANONICAL : i32 = 0; +const VARIANT_DAMPED : i32 = 1; +const VARIANT_TUNED_STD : i32 = 2; +const VARIANT_RANDOM_RAT : i32 = 3; + +// ---- types ---- + +pub const AdamWConfig = struct { + learning_rate : f32, + beta1 : f32, + beta2 : f32, + weight_decay : f32, + epsilon : f32, + amsgrad : bool, +}; + +/// Offsets of the three moment buffers inside one backing array. code is ARENA_OK (0) or the +/// arena's error; on an error the offsets are ARENA_NO_OFFSET and the arena is the one that +/// went in. +pub const AdamWState = struct { + arena : Arena, + m_off : i32, + v_off : i32, + x_off : i32, + code : i32, +}; + +/// One parameter, its three moments, and whether the step was applied. Test scaffolding. +pub const Probe = struct { + p : f32, + m : f32, + v : f32, + x : f32, + ok : bool, +}; + +// ---- configuration ---- + +/// Variant 0 canonical (beta1 = 1/phi, wd = phi^-3), 1 damped (0.9/phi, 0.01), +/// 2 tuned standard (0.9, 0.01), 3 random rational (4/7, 7/30). Any other code is the +/// standard arm, the one a typo cannot make exotic. beta2 is 0.999 in every arm. +pub fn variant_config(variant: i32, lr: f32) -> AdamWConfig { + var b1 : f32 = DEFAULT_BETA1; + var wd : f32 = DEFAULT_WEIGHT_DECAY; + if (variant == VARIANT_CANONICAL) { + b1 = PHI_INV; + wd = PHI_INV_CUBED; } - - // compute_bias_correction_v1(beta: gf16::GF16, t: u64) -> f32 - // Computes bias correction factor. - fn compute_bias_correction_v1(beta: gf16::GF16, t: u64) -> gf16::GF16 { - // return 1 - pow(beta, t) + if (variant == VARIANT_DAMPED) { + b1 = PHI_DAMPED_BETA1; } - - // update_first_moment_v1(m_prev: gf16::GF16, grad: gf16::GF16, beta1: gf16::GF16) -> f32 - // Updates first moment estimate. - fn update_first_moment_v1(m_prev: gf16::GF16, grad: gf16::GF16, beta1: gf16::GF16) -> gf16::GF16 { - // return beta1 * m_prev + (1 - beta1) * grad + if (variant == VARIANT_RANDOM_RAT) { + b1 = RATIONAL_BETA1; + wd = RATIONAL_WEIGHT_DECAY; } + return AdamWConfig{ .learning_rate = lr, .beta1 = b1, .beta2 = DEFAULT_BETA2, .weight_decay = wd, .epsilon = DEFAULT_EPSILON, .amsgrad = false }; +} - // update_second_moment_v1(v_prev: gf16::GF16, grad: gf16::GF16, beta2: gf16::GF16) -> gf16::GF16 - // Updates second moment estimate. - fn update_second_moment_v1(v_prev: gf16::GF16, grad: gf16::GF16, beta2: gf16::GF16) -> gf16::GF16 { - // return beta2 * v_prev + (1 - beta2) * grad * grad - } +pub fn plain_config(lr: f32, wd: f32, amsgrad: bool) -> AdamWConfig { + return AdamWConfig{ .learning_rate = lr, .beta1 = DEFAULT_BETA1, .beta2 = DEFAULT_BETA2, .weight_decay = wd, .epsilon = DEFAULT_EPSILON, .amsgrad = amsgrad }; +} - // apply_weight_decay_v1(params: []gf16::GF16, lr: gf16::GF16, weight_decay: gf16::GF16) -> []gf16::GF16 - // Applies decoupled weight decay to parameters. - fn apply_weight_decay_v1(params: []gf16::GF16, lr: gf16::GF16, weight_decay: gf16::GF16) -> []gf16::GF16 { - // params[i] = params[i] * (1 - lr * weight_decay) - } +// ---- state in an arena ---- - // compute_update_v1(m: gf16::GF16, v: gf16::GF16, lr_t: gf16::GF16, epsilon: gf16::GF16) -> f32 - // Computes parameter update from moments. - fn compute_update_v1(m: gf16::GF16, v: gf16::GF16, lr_t: gf16::GF16, epsilon: gf16::GF16) -> gf16::GF16 { - // return lr_t * m / (sqrt(v) + epsilon) +/// Take m, v and x (n slots each) from the arena, all three or none, zeroed in the backing +/// array. When the third does not fit, the arena and the backing array are as they came in. +pub fn state_alloc(a: Arena, backing: []f32, n: i32) -> AdamWState { + const r1 : AllocResult = arena_alloc(a, n, 1); + if (r1.code != ARENA_OK) { + return AdamWState{ .arena = a, .m_off = ARENA_NO_OFFSET, .v_off = ARENA_NO_OFFSET, .x_off = ARENA_NO_OFFSET, .code = r1.code }; } - - // get_effective_beta(config: AdamWConfigV1, beta_type: "beta1" | "beta2") -> f32 - // Returns effective beta with phi-damping if enabled. - fn get_effective_beta(config: AdamWConfigV1, beta_type: []const u8) -> gf16::GF16 { - // if not use_phi_betas: return corresponding config beta - // if beta_type == "beta1": return PHI_BETA1 - // else: return PHI_BETA2 + const r2 : AllocResult = arena_alloc(r1.arena, n, 1); + if (r2.code != ARENA_OK) { + return AdamWState{ .arena = a, .m_off = ARENA_NO_OFFSET, .v_off = ARENA_NO_OFFSET, .x_off = ARENA_NO_OFFSET, .code = r2.code }; } - - // amsgrad_update(v: gf16::GF16, v_max: gf16::GF16) -> f32 - // Computes AMSGrad maximum tracking. - fn amsgrad_update(v: gf16::GF16, v_max: gf16::GF16) -> gf16::GF16 { - if (v > v_max) { - return v; - } else { - return v_max; - } + const r3 : AllocResult = arena_alloc(r2.arena, n, 1); + if (r3.code != ARENA_OK) { + return AdamWState{ .arena = a, .m_off = ARENA_NO_OFFSET, .v_off = ARENA_NO_OFFSET, .x_off = ARENA_NO_OFFSET, .code = r3.code }; } - - // =========================================================== - // TDD: Tests - // =========================================================== - - test init_creates_zero_moments - given config = AdamWConfigV1{.learning_rate = 1e-3, .beta1 = 0.9, .beta2 = 0.999, .weight_decay = 0.01, .epsilon = 1e-8, .amsgrad = false, .use_phi_betas = false} - when state = init_v1(config, 10) - then state.m.len == 10 - and state.v.len == 10 - and all(state.m, is_zero) - and all(state.v, is_zero) - and state.step == 0 - - test init_with_amsgrad_creates_v_max - given config = AdamWConfigV1{.amsgrad = true} - when state = init_v1(config, 10) - then state.v_max.len == 10 - and all(state.v_max, is_zero) - - test compute_bias_correction_first_step - given beta = 0.9 - and t = 1 - when result = compute_bias_correction_v1(beta, t) - then approximately_equal(result, 0.1) - - test compute_bias_correction_later_steps - given beta = 0.9 - and t = 10 - when result = compute_bias_correction_v1(beta, t) - then result < 1.0 - and result > 0.0 - - test update_first_momentum_basic - given m_prev = 0.0 - and grad = 1.0 - and beta1 = 0.9 - when result = update_first_moment_v1(m_prev, grad, beta1) - then approximately_equal(result, 0.1) // 0.9 * 0 + 0.1 * 1 - - test update_first_momentum_with_history - given m_prev = 0.5 - and grad = 1.0 - and beta1 = 0.9 - when result = update_first_moment_v1(m_prev, grad, beta1) - then approximately_equal(result, 0.9 * 0.5 + 0.1) - - test update_second_momentum_basic - given v_prev = 0.0 - and grad = 2.0 - and beta2 = 0.999 - when result = update_second_moment_v1(v_prev, grad, beta2) - then approximately_equal(result, 0.001 * 4.0) // 0.999 * 0 + 0.001 * 4 - - test compute_update_basic - given m = 0.1 - and v = 0.01 - and lr_t = 0.001 - and epsilon = 1e-8 - when result = compute_update_v1(m, v, lr_t, epsilon) - then result > 0.0 - - test apply_weight_decay_reduces_parameters - given params = [1.0, 2.0, 3.0] - and lr = 0.01 - and weight_decay = 0.1 - when result = apply_weight_decay_v1(params, lr, weight_decay) - then result[0] < 1.0 - and result[1] < 2.0 - and result[2] < 3.0 - - test step_without_weight_decay_matches_adam - given config = AdamWConfigV1{.learning_rate = 0.001, .beta1 = 0.9, .beta2 = 0.999, .weight_decay = 0.0} - and state = init_v1(config, 1) - and params = [1.0] - and grads = [0.1] - when result = step_v1(state, params, grads) - then result.updated_params.len == 1 - - test step_with_weight_decay_shrinks_params - given config = AdamWConfigV1{.learning_rate = 0.01, .beta1 = 0.9, .beta2 = 0.999, .weight_decay = 0.1} - and state = init_v1(config, 1) - and params = [1.0] - and grads = [0.0] // No gradient, only weight decay - when result = step_v1(state, params, grads) - then result.updated_params[0] < 1.0 - - test amsgrad_update_tracks_maximum - given v = 0.05 - and v_max = 0.03 - when result = amsgrad_update(v, v_max) - then result == 0.05 - - test amsgrad_update_preserves_maximum - given v = 0.02 - and v_max = 0.05 - when result = amsgrad_update(v, v_max) - then result == 0.05 - - test get_effective_beta_phi_damped - given config = AdamWConfigV1{.beta1 = 0.9, .beta2 = 0.999, .use_phi_betas = true} - when beta1 = get_effective_beta(config, "beta1") - and beta2 = get_effective_beta(config, "beta2") - then approximately_equal(beta1, PHI_BETA1) - and approximately_equal(beta2, PHI_BETA2) - - test get_effective_beta_uses_config - given config = AdamWConfigV1{.beta1 = 0.85, .beta2 = 0.995, .use_phi_betas = false} - when beta1 = get_effective_beta(config, "beta1") - and beta2 = get_effective_beta(config, "beta2") - then approximately_equal(beta1, 0.85) - and approximately_equal(beta2, 0.995) - - test step_increases_step_counter - given config = AdamWConfigV1{} - and state = init_v1(config, 10) - and params = random_input(10) - and grads = random_input(10) - when result = step_v1(state, params, grads) - then result.step_norm >= 0.0 // Step was executed - - // =========================================================== - // TDD: Invariants - // =========================================================== - - invariant step_preserves_param_count - given config = any_adamw_config() - and state = init_v1(config, 100) - and params = random_input(100) - and grads = random_input(100) - when result = step_v1(state, params, grads) - then result.updated_params.len == 100 - - invariant moments_have_same_length - given config = any_adamw_config() - and state = init_v1(config, 50) - then state.m.len == 50 - and state.v.len == 50 - - invariant first_moment_remains_bounded - given config = any_adamw_config() - and state = init_v1(config, 10) - and params = finite_input(10) - and grads = bounded_input(10, -1.0, 1.0) - and result = step_v1(state, params, grads) - then all(result.m, is_finite) - - invariant second_moment_non_negative { - const config = any_adamw_config(); - const state = init_v1(config, 10); - const params = finite_input(10); - const grads = finite_input(10); - const result = step_v1(state, params, grads); - // was: then all(result.v, fn(x) x >= 0.0) -- t27 has no closures, so - // the predicate is checked by an explicit loop over every element - var i : usize = 0; - while (i < result.v.len) { - assert(result.v[i] >= 0.0); - i = i + 1; + fill_slots(backing, r1.offset, n, 0.0); + fill_slots(backing, r2.offset, n, 0.0); + fill_slots(backing, r3.offset, n, 0.0); + return AdamWState{ .arena = r3.arena, .m_off = r1.offset, .v_off = r2.offset, .x_off = r3.offset, .code = ARENA_OK }; +} + +// ---- the update ---- + +/// One AdamW step number t (1-based) over params. mom holds m at m_off, v at v_off and the +/// AMSGrad maximum at x_off, each params.len slots. +/// p = p (1 - lr wd) decay, decoupled: never enters m or v +/// m = b1 m + (1 - b1) g v = b2 v + (1 - b2) g^2 +/// x = max(x, v) only when amsgrad +/// m_hat = m / (1 - b1^t) v_hat = (amsgrad ? x : v) / (1 - b2^t) +/// p = p - lr m_hat / (sqrt(v_hat) + eps) +/// Returns false and changes nothing when t is 0 or a gradient is NaN or infinite. +pub fn adamw_step(params: []f32, grads: []const f32, mom: []f32, m_off: i32, v_off: i32, x_off: i32, t: u32, c: AdamWConfig) -> bool { + if (t == 0) { + return false; + } + var i : usize = 0; + while (i < grads.len) { + if (!is_finite(grads[i])) { + return false; } + i = i + 1; } - - invariant bias_correction_in_zero_one_range - given beta = random_gf16_in_range(0.0, 0.999) - and t = random_u64_in_range(1, 1000) - when bc = compute_bias_correction_v1(beta, t) - then bc >= 0.0 - and bc <= 1.0 - - invariant phi_betas_smaller_than_defaults - then PHI_BETA1 < DEFAULT_BETA1_V1 - and PHI_BETA2 < DEFAULT_BETA2_V1 - - invariant weight_decay_preserves_sign - given params = [-1.0, 0.0, 1.0] - and lr = 0.01 - and weight_decay = 0.1 - when result = apply_weight_decay_v1(params, lr, weight_decay) - then result[0] < 0.0 - and result[2] > 0.0 - - invariant amsgrad_v_max_non_decreasing { - const config = AdamWConfigV1{.amsgrad = true}; - const state = init_v1(config, 10); - const params = random_input(10); - const grads = random_input(10); - const result1 = step_v1(state, params, grads); - const result2 = step_v1(state, result1.updated_params, grads); - // was: then all_pairs(result1.v_max, result2.v_max, fn(a, b) b >= a) - // -- t27 has no closures; the pairs are taken element-wise (index i of - // both arrays) by an explicit loop - var i : usize = 0; - while (i < result1.v_max.len) { - assert(result2.v_max[i] >= result1.v_max[i]); - i = i + 1; + const mo : usize = @intCast(m_off); + const vo : usize = @intCast(v_off); + const xo : usize = @intCast(x_off); + const c1 = 1.0 - powi(c.beta1, t); + const c2 = 1.0 - powi(c.beta2, t); + const keep = 1.0 - c.learning_rate * c.weight_decay; + i = 0; + while (i < params.len) { + const g = grads[i]; + const decayed = params[i] * keep; + mom[mo + i] = c.beta1 * mom[mo + i] + (1.0 - c.beta1) * g; + mom[vo + i] = c.beta2 * mom[vo + i] + (1.0 - c.beta2) * g * g; + var second = mom[vo + i]; + if (c.amsgrad) { + if (mom[vo + i] > mom[xo + i]) { + mom[xo + i] = mom[vo + i]; + } + second = mom[xo + i]; } + const m_hat = mom[mo + i] / c1; + const v_hat = second / c2; + params[i] = decayed - c.learning_rate * m_hat / (sqrt(v_hat) + c.epsilon); + i = i + 1; } - - invariant denominator_positive - given v = 0.0 - and epsilon = 1e-8 - when denom = sqrt(v) + epsilon - then denom > 0.0 - - // =========================================================== - // TDD: Benchmarks - // =========================================================== - - bench init_small - given config = any_adamw_config() - when result = init_v1(config, 100) - then elapsed_time_us < 50 - - bench init_large - given config = any_adamw_config() - when result = init_v1(config, 1000000) - then elapsed_time_ms < 20 - - bench step_small - given config = any_adamw_config() - and state = init_v1(config, 100) - and params = random_input(100) - and grads = random_input(100) - when result = step_v1(state, params, grads) - then elapsed_time_us < 200 - - bench step_medium - given config = any_adamw_config() - and state = init_v1(config, 10000) - and params = random_input(10000) - and grads = random_input(10000) - when result = step_v1(state, params, grads) - then elapsed_time_ms < 10 - - bench step_large - given config = any_adamw_config() - and state = init_v1(config, 1000000) - and params = random_input(1000000) - and grads = random_input(1000000) - when result = step_v1(state, params, grads) - then elapsed_time_ms < 200 - - bench compute_bias_correction - given beta = 0.9 - and t = 100 - when result = compute_bias_correction_v1(beta, t) - then elapsed_time_ns < 50 - - bench amsgrad_update - given v = random_positive_gf16() - and v_max = random_positive_gf16() - when result = amsgrad_update(v, v_max) - then elapsed_time_ns < 10 - - // =========================================================== - // Mathematical Notes - // =========================================================== - // - // AdamW Update Rule: - // 1. Apply weight decay (decoupled): - // theta_t = theta_{t-1} * (1 - eta * lambda) - // - // 2. Update moments: - // m_t = beta1 * m_{t-1} + (1 - beta1) * g_t - // v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2 - // - // 3. Bias correction: - // m_hat_t = m_t / (1 - beta1^t) - // v_hat_t = v_t / (1 - beta2^t) - // - // 4. Parameter update: - // theta_t = theta_t - eta * m_hat_t / (sqrt(v_hat_t) + eps) - // - // AMSGrad Variant: - // v_max_t = max(v_max_{t-1}, v_t) - // Use sqrt(v_max_t) instead of sqrt(v_hat_t) in denominator - // - // Phi-Damped Betas: - // beta1_eff = beta1 / phi ~ 0.556 for beta1 = 0.9 - // beta2_eff = beta2 / phi ~ 0.617 for beta2 = 0.999 - // Reduces momentum, provides smoother convergence - // - // Difference from Adam: - // Adam: weight decay added to gradients before moment update - // AdamW: weight decay applied directly to parameters (decoupled) - // - // =========================================================== - - // ===================================================================== - // Phase B1 (epic #181) -- optimizer control ablation blocks (appended) - // phi-canonical (beta1=phi^-1, wd=phi^-3) vs Damped vs TunedStd vs RandomRat. - // All phi claims [Open conjecture]; falsification: non-phi control within MC error. - // ===================================================================== - - // ========================================================= - // 1. Constants - // ========================================================= - - const PHI : gf16::GF16 = 1.618033988749895; - - // Standard AdamW defaults (tuning convention; Kingma & Ba 2015, Loshchilov & Hutter 2019) - const DEFAULT_LEARNING_RATE : gf16::GF16 = 1e-3; - const DEFAULT_BETA1 : gf16::GF16 = 0.9; - const DEFAULT_BETA2 : gf16::GF16 = 0.999; - const DEFAULT_WEIGHT_DECAY : gf16::GF16 = 0.01; - const DEFAULT_EPSILON : gf16::GF16 = 1e-8; - const DEFAULT_AMSGRAD : bool = false; - - // --------------------------------------------------------- - // CANONICAL phi anchor (matches igla-peek/src/optimizer.rs) - // [Open conjecture] -- no momentum theory singles out phi. - // beta1 = phi^(-1) ~ 0.618034 - // weight_decay = phi^(-3) ~ 0.236068 (NOT 0.118034; see NOTES.md) - // lr_default = phi^(-3) ~ 0.236068 - // --------------------------------------------------------- - const PHI_INV : gf16::GF16 = 0.618033988749895; // phi^(-1) - const PHI_INV_SQ : gf16::GF16 = 0.381966011250105; // phi^(-2) - const PHI_INV_CUBED : gf16::GF16 = 0.236067977499790; // phi^(-3) = 1/(phi^3) - // NOTE: igla-peek/src/optimizer.rs test confirms phi^(-3) ~ 0.23607 - // (see: assert!((expected_wd - 0.23607).abs() < 0.001)) - // The brief's "0.118034" for phi^(-3) is a typo; 0.118034 ~ phi^(-4.44), - // which is not a clean phi power. - const PHI_CANONICAL_BETA1 : gf16::GF16 = PHI_INV; // [Open conjecture] - const PHI_CANONICAL_BETA2 : gf16::GF16 = 0.999; // unchanged (standard) - const PHI_CANONICAL_WEIGHT_DECAY : gf16::GF16 = PHI_INV_CUBED; // [Open conjecture] - const PHI_CANONICAL_LR : gf16::GF16 = PHI_INV_CUBED; // [Open conjecture] - - // --------------------------------------------------------- - // PHI-DAMPED variant (old PHI_BETA1 from pre-B1 spec) - // [Open conjecture] -- 0.9/phi is an ad hoc rescaling with no - // theoretical derivation. Kept as named comparison variant. - // beta1 = 0.9 / phi ~ 0.556 - // beta2 = 0.999 / phi ~ 0.617 (for completeness) - // --------------------------------------------------------- - const PHI_DAMPED_BETA1 : gf16::GF16 = 0.9 / PHI; // ~ 0.556 [Open conjecture] - const PHI_DAMPED_BETA2 : gf16::GF16 = 0.999 / PHI; // ~ 0.617 [Open conjecture] - - // --------------------------------------------------------- - // Description-cost anchor (for random-rational control B1d) - // A NON-Fibonacci rational with comparable bit-description cost to PHI_INV. - // Parent decision: 5/8 was rejected (5,8 are Fibonacci numbers, so it is - // a phi-biased control). Use 4/7 and 7/30 -- all of {4,7,30} are outside - // the Fibonacci sequence {1,2,3,5,8,13,21,...}, giving a clean phi-free - // baseline. Magnitudes are matched to the phi anchors (beta1 ~ 0.57 vs - // phi^(-1) ~ 0.62; wd ~ 0.233 vs phi^(-3) ~ 0.236). - // Source: Li & Talwalkar 2019 random-search; MDL: Shaw 2025 (2509.22445) - // --------------------------------------------------------- - const RATIONAL_BETA1 : gf16::GF16 = 0.5714285714285714; // 4/7 -- non-Fibonacci control - const RATIONAL_WEIGHT_DECAY : gf16::GF16 = 0.2333333333333333; // 7/30 -- comparable cost - - // ========================================================= - // 2. Types - // ========================================================= - - // No language-level default is forced: every ablation arm names its variant - // explicitly. The TRAINER runtime default (with_phi_defaults) is Canonical; - // TunedStd is the standard-AdamW CONTROL arm, not a silent default. - pub const PhiVariant = enum { - Canonical, // beta1=phi^(-1)~0.618, wd=phi^(-3)~0.236068 [Open conjecture] (trainer default) - Damped, // beta1=0.9/phi~0.556, wd=0.01 [Open conjecture] - TunedStd, // beta1=0.9, wd=0.01 (AdamW field-standard control) - RandomRat, // beta1=4/7~0.571, wd=7/30~0.233 (non-Fibonacci control) - }; - - pub const AdamWConfig = struct { - learning_rate : gf16::GF16, // eta: step size - beta1 : gf16::GF16, // beta1: first moment decay - beta2 : gf16::GF16, // beta2: second moment decay - weight_decay : gf16::GF16, // lambda: decoupled weight decay - epsilon : gf16::GF16, // eps: numerical stability - amsgrad : bool, // Use AMSGrad variant - phi_variant : PhiVariant, // Which parametrisation to use - }; - - pub const AdamWState = struct { - m : []gf16::GF16, // First moment (moving average of gradients) - v : []gf16::GF16, // Second moment (moving average of squared gradients) - v_max : []gf16::GF16, // Max of v (for AMSGrad) - param_count : u32, // Number of parameters - step : u64, // Current optimization step (t) - }; - - pub const OptimizerStepResult = struct { - updated_params : []gf16::GF16, // New parameter values - m : []gf16::GF16, // Updated first moments - v : []gf16::GF16, // Updated second moments - step_norm : gf16::GF16, // L2 norm of the parameter update - }; - - // ========================================================= - // 3. Core Functions - // ========================================================= - - // init(config: AdamWConfig, param_count: u32) -> AdamWState - fn init(config: AdamWConfig, param_count: u32) -> AdamWState { - // m[i] = 0.0 for all i - // v[i] = 0.0 for all i - // v_max[i] = 0.0 for all i (if amsgrad) - // step = 0 + return true; +} + +// ---- helpers that exist only so the tests go through function calls ---- + +/// One parameter p with moments (m, v, x) taking step number t with gradient g. +fn probe_step(p: f32, g: f32, m: f32, v: f32, x: f32, t: u32, c: AdamWConfig) -> Probe { + var params : [1]f32 = [p]; + const grads : [1]f32 = [g]; + var mom : [3]f32 = [m, v, x]; + const ok : bool = adamw_step(¶ms, &grads, &mom, 0, 1, 2, t, c); + return Probe{ .p = params[0], .m = mom[0], .v = mom[1], .x = mom[2], .ok = ok }; +} + +/// Parameter 2.0 with a zero gradient for `steps` steps at lr 0.1, wd 0.5. which: 0 = p, 1 = m, 2 = v. +fn zero_grad_run(steps: u32, which: i32) -> f32 { + var params : [1]f32 = [2.0]; + const grads : [1]f32 = [0.0]; + var mom : [3]f32 = [0.0, 0.0, 0.0]; + const c : AdamWConfig = plain_config(0.1, 0.5, false); + var t : u32 = 1; + while (t <= steps) { + const ok : bool = adamw_step(¶ms, &grads, &mom, 0, 1, 2, t, c); + t = t + 1; } - - // step(state: AdamWState, params: []gf16::GF16, grads: []gf16::GF16) -> OptimizerStepResult - fn step(state: AdamWState, params: []gf16::GF16, grads: []gf16::GF16) -> OptimizerStepResult { - // t = step + 1 - // bias_correction1 = 1 - beta1^t - // bias_correction2 = 1 - beta2^t - // lr_t = lr * sqrt(bias_correction2) / bias_correction1 - // - // Apply weight decay (decoupled): - // params[i] = params[i] * (1 - lr * weight_decay) - // - // Update moments: - // m_t = beta1 * m_{t-1} + (1 - beta1) * g_t - // v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2 - // - // If amsgrad: - // v_max_t = max(v_max_{t-1}, v_t) - // denominator = sqrt(v_max_t) + eps - // else: - // denominator = sqrt(v_t) + eps - // - // Update parameters: - // params[i] = params[i] - lr_t * m_t / denominator + if (which == 1) { + return mom[0]; } - - fn compute_bias_correction(beta: gf16::GF16, t: u64) -> gf16::GF16 { - // return 1 - pow(beta, t) + if (which == 2) { + return mom[1]; } - - fn update_first_moment(m_prev: gf16::GF16, grad: gf16::GF16, beta1: gf16::GF16) -> gf16::GF16 { - // return beta1 * m_prev + (1 - beta1) * grad + return params[0]; +} + +/// Two steps from p = 1.0, lr 0.1, wd 0: gradient 2.0 then 0.0. which: 0 = p, 1 = v, 2 = x. +fn two_step_run(amsgrad: bool, which: i32) -> f32 { + var params : [1]f32 = [1.0]; + var grads : [1]f32 = [2.0]; + var mom : [3]f32 = [0.0, 0.0, 0.0]; + const c : AdamWConfig = plain_config(0.1, 0.0, amsgrad); + const first : bool = adamw_step(¶ms, &grads, &mom, 0, 1, 2, 1, c); + grads[0] = 0.0; + const second : bool = adamw_step(¶ms, &grads, &mom, 0, 1, 2, 2, c); + if (which == 1) { + return mom[1]; } - - fn update_second_moment(v_prev: gf16::GF16, grad: gf16::GF16, beta2: gf16::GF16) -> gf16::GF16 { - // return beta2 * v_prev + (1 - beta2) * grad * grad + if (which == 2) { + return mom[2]; } - - fn apply_weight_decay(params: []gf16::GF16, lr: gf16::GF16, weight_decay: gf16::GF16) -> []gf16::GF16 { - // params[i] = params[i] * (1 - lr * weight_decay) + return params[0]; +} + +/// a / b at run time, where IEEE gives NaN for 0 / 0 and inf for 1 / 0 +fn ratio(a: f32, b: f32) -> f32 { + return a / b; +} + +/// 1 if the step was applied, 0 if it was refused; a refused step must leave p, m, v alone. +fn applied(g: f32, t: u32) -> i32 { + const r : Probe = probe_step(1.0, g, 0.25, 0.5, 0.75, t, plain_config(0.001, 0.1, true)); + if (r.ok) { + return 1; } - - fn compute_update(m: gf16::GF16, v: gf16::GF16, lr_t: gf16::GF16, epsilon: gf16::GF16) -> gf16::GF16 { - // return lr_t * m / (sqrt(v) + epsilon) + if (r.p != 1.0) { + return 7; } - - // get_effective_betas(variant: PhiVariant) -> (beta1, beta2, weight_decay) - // Returns (beta1, beta2, weight_decay) for the given variant. - fn get_effective_betas(variant: PhiVariant) -> (gf16::GF16, gf16::GF16, gf16::GF16) { - // TODO: Implement proper conditional logic - // (placeholder kept as written: a single value, not the triple) - return DEFAULT_BETA1; + if (r.m != 0.25) { + return 7; } - - test get_effective_betas - // Test Canonical variant - given (b1, b2, wd) = get_effective_betas(PhiVariant.Canonical) - then approximately_equal(b1, PHI_CANONICAL_BETA1) // ~ 0.618 - and approximately_equal(b2, PHI_CANONICAL_BETA2) // 0.999 - and approximately_equal(wd, PHI_CANONICAL_WEIGHT_DECAY) // ~ 0.236068 - - // Test Damped variant - given (b1, b2, wd) = get_effective_betas(PhiVariant.Damped) - then approximately_equal(b1, PHI_DAMPED_BETA1) // ~ 0.556 - and approximately_equal(b2, DEFAULT_BETA2) // 0.999 - and approximately_equal(wd, DEFAULT_WEIGHT_DECAY) // 0.01 - - // Test TunedStd variant - given (b1, b2, wd) = get_effective_betas(PhiVariant.TunedStd) - then approximately_equal(b1, DEFAULT_BETA1) // 0.9 - and approximately_equal(b2, DEFAULT_BETA2) // 0.999 - and approximately_equal(wd, DEFAULT_WEIGHT_DECAY) // 0.01 - - // Test RandomRat variant - given (b1, b2, wd) = get_effective_betas(PhiVariant.RandomRat) - then approximately_equal(b1, RATIONAL_BETA1) // 4/7 ~ 0.5714 - and approximately_equal(b2, DEFAULT_BETA2) // 0.999 - and approximately_equal(wd, RATIONAL_WEIGHT_DECAY) // 7/30 ~ 0.2333 - - - -test amsgrad_update - // Test when v > v_max - given result = amsgrad_update(5.0, 3.0) - then result == 5.0 - - // Test when v <= v_max - given result = amsgrad_update(2.0, 4.0) - then result == 4.0 - - // Test when v == v_max - given result = amsgrad_update(3.0, 3.0) - then result == 3.0 - - // Test with negative values - given result = amsgrad_update(-1.0, -2.0) - then result == -1.0 - - given result = amsgrad_update(-3.0, -1.0) - then result == -1.0 - - // ========================================================= - // EXISTING TDD: Tests (preserved from pre-B1 spec) - // ========================================================= - - test init_creates_zero_moments_v2 - given config = AdamWConfig{.learning_rate = 1e-3, .beta1 = 0.9, .beta2 = 0.999, .weight_decay = 0.01, .epsilon = 1e-8, .amsgrad = false, .phi_variant = PhiVariant.TunedStd} - when state = init(config, 10) - then state.m.len == 10 - and state.v.len == 10 - and all(state.m, is_zero) - and all(state.v, is_zero) - and state.step == 0 - - test init_with_amsgrad_creates_v_max_v2 - given config = AdamWConfig{.amsgrad = true} - when state = init(config, 10) - then state.v_max.len == 10 - and all(state.v_max, is_zero) - - test compute_bias_correction_first_step_v2 - given beta = 0.9 - and t = 1 - when result = compute_bias_correction(beta, t) - then approximately_equal(result, 0.1) - - test compute_bias_correction_later_steps_v2 - given beta = 0.9 - and t = 10 - when result = compute_bias_correction(beta, t) - then result < 1.0 - and result > 0.0 - - test update_first_momentum_basic_v2 - given m_prev = 0.0 - and grad = 1.0 - and beta1 = 0.9 - when result = update_first_moment(m_prev, grad, beta1) - then approximately_equal(result, 0.1) // 0.9 * 0 + 0.1 * 1 - - test update_first_momentum_with_history_v2 - given m_prev = 0.5 - and grad = 1.0 - and beta1 = 0.9 - when result = update_first_moment(m_prev, grad, beta1) - then approximately_equal(result, 0.9 * 0.5 + 0.1) - - test update_second_momentum_basic_v2 - given v_prev = 0.0 - and grad = 2.0 - and beta2 = 0.999 - when result = update_second_moment(v_prev, grad, beta2) - then approximately_equal(result, 0.001 * 4.0) - - test compute_update_basic_v2 - given m = 0.1 - and v = 0.01 - and lr_t = 0.001 - and epsilon = 1e-8 - when result = compute_update(m, v, lr_t, epsilon) - then result > 0.0 - - test apply_weight_decay_reduces_parameters_v2 - given params = [1.0, 2.0, 3.0] - and lr = 0.01 - and weight_decay = 0.1 - when result = apply_weight_decay(params, lr, weight_decay) - then result[0] < 1.0 - and result[1] < 2.0 - and result[2] < 3.0 - - test step_without_weight_decay_matches_adam_v2 - given config = AdamWConfig{.learning_rate = 0.001, .beta1 = 0.9, .beta2 = 0.999, .weight_decay = 0.0} - and state = init(config, 1) - and params = [1.0] - and grads = [0.1] - when result = step(state, params, grads) - then result.updated_params.len == 1 - - test step_with_weight_decay_shrinks_params_v2 - given config = AdamWConfig{.learning_rate = 0.01, .beta1 = 0.9, .beta2 = 0.999, .weight_decay = 0.1} - and state = init(config, 1) - and params = [1.0] - and grads = [0.0] - when result = step(state, params, grads) - then result.updated_params[0] < 1.0 - - test amsgrad_update_tracks_maximum_v2 - given v = 0.05 - and v_max = 0.03 - when result = amsgrad_update(v, v_max) - then result == 0.05 - - test amsgrad_update_preserves_maximum_v2 - given v = 0.02 - and v_max = 0.05 - when result = amsgrad_update(v, v_max) - then result == 0.05 - - test step_increases_step_counter_v2 - given config = AdamWConfig{} - and state = init(config, 10) - and params = random_input(10) - and grads = random_input(10) - when result = step(state, params, grads) - then result.step_norm >= 0.0 - - // ========================================================= - // PHASE B1: Constant-identity tests (reconciliation) - // ========================================================= - - // B1-CONST-1: Canonical phi anchor matches igla-peek/src/optimizer.rs - // [Verified arithmetic] phi^(-1) = 1/phi ~ 0.618034 - test b1_canonical_beta1_matches_trainer - given expected = 1.0 / 1.618033988749895 - when val = PHI_CANONICAL_BETA1 - then approximately_equal(val, expected) // tolerance 1e-9 - - // B1-CONST-2: Canonical weight_decay matches trainer phi^(-3) - // [Verified arithmetic] phi^(-3) = 1/(phi^3) ~ 0.236068 - // igla-peek test: assert!((expected_wd - 0.23607).abs() < 0.001) - test b1_canonical_wd_matches_trainer - given expected = 1.0 / (1.618033988749895 * 1.618033988749895 * 1.618033988749895) - when val = PHI_CANONICAL_WEIGHT_DECAY - then approximately_equal(val, expected) // ~ 0.23607, tolerance 1e-9 - - // B1-CONST-3: Damped beta1 (legacy) equals 0.9/phi - test b1_damped_beta1_is_0p9_over_phi - given expected = 0.9 / 1.618033988749895 - when val = PHI_DAMPED_BETA1 - then approximately_equal(val, expected) // ~ 0.556 - - // B1-CONST-4: Canonical and damped are numerically distinct - test b1_canonical_and_damped_are_distinct - then PHI_CANONICAL_BETA1 > PHI_DAMPED_BETA1 - and approximately_equal(PHI_CANONICAL_BETA1, 0.618034) // 3 sig figs - and approximately_equal(PHI_DAMPED_BETA1, 0.55557) // 3 sig figs - - // B1-CONST-5: Identity law L5 -- [Verified] - // phi^2 + phi^(-2) = 3 (Law L5, the only Verified phi-fact) - test b1_trinity_identity_l5 - given phi = 1.618033988749895 - when lhs = phi * phi + 1.0 / (phi * phi) - then approximately_equal(lhs, 3.0) // tolerance 1e-12 - - // B1-CONST-6: Random-rational is distinct from all phi values - test b1_random_rational_distinct_from_phi - then RATIONAL_BETA1 != PHI_CANONICAL_BETA1 - and RATIONAL_BETA1 != PHI_DAMPED_BETA1 - and RATIONAL_BETA1 != DEFAULT_BETA1 - - // ========================================================= - // PHASE B1: Control-test design (ablation) - // - // FALSIFICATION CRITERION (mandatory, non-negotiable): - // If ANY of variants (b), (c), or (d) achieves BPB within - // Monte-Carlo error of variant (a) on an equal tuning budget, - // phi-specificity of AdamW hyperparameters is FALSIFIED. - // The experiment reports the null first. - // - // Experimental protocol: - // - Metric: BPB (bits-per-byte) on held-out validation set - // - Champion baseline: BPB = 2.1919 (seed=43, ~81K steps, - // AdamW, hidden=828, sha=cd91c45), Hoffmann 2022 (2203.15556) regime - // - Budget: equal total gradient steps AND equal HP-tuning budget - // (Li & Talwalkar 2019 random-search, arXiv:1902.07638; - // Melis/Dyer/Blunsom 2017 equal-budget, arXiv:1707.05589) - // - Seeds: minimum 3 seeds per variant; report mean +/- std - // - Monte-Carlo error: estimate via bootstrap on validation set - // - AdamW implicit L-inf bias: Xie & Li 2024 (arXiv:2404.04454) - // - Weight-decay scaling: Wang & Aitchison 2025 (arXiv:2405.13698) - // - Cautious weight decay: Chen 2025 (arXiv:2510.12402) - // ========================================================= - - // B1a: phi-canonical (beta1=phi^(-1)~0.618, wd=phi^(-3)~0.236068) - // [Open conjecture] -- this is the trainer default - test b1_variant_a_phi_canonical_config - given (b1, b2, wd) = get_effective_betas(PhiVariant.Canonical) - then approximately_equal(b1, PHI_CANONICAL_BETA1) // ~ 0.618 - and approximately_equal(b2, 0.999) - and approximately_equal(wd, PHI_CANONICAL_WEIGHT_DECAY) // ~ 0.236068 - - // B1b: phi-damped (beta1=0.9/phi~0.556, wd=0.01) - // [Open conjecture] -- pre-B1 spec default - test b1_variant_b_phi_damped_config - given (b1, b2, wd) = get_effective_betas(PhiVariant.Damped) - then approximately_equal(b1, PHI_DAMPED_BETA1) // ~ 0.556 - and approximately_equal(b2, DEFAULT_BETA2) // 0.999 - and approximately_equal(wd, DEFAULT_WEIGHT_DECAY) // 0.01 - - // B1c: tuned-standard AdamW (beta1=0.9, wd=0.01) - // Tuning convention from Kingma & Ba 2015 / Loshchilov & Hutter 2019 - test b1_variant_c_tuned_standard_config - given (b1, b2, wd) = get_effective_betas(PhiVariant.TunedStd) - then approximately_equal(b1, 0.9) - and approximately_equal(b2, 0.999) - and approximately_equal(wd, 0.01) - - // B1d: random-rational (beta1=4/7, wd=7/30 -- non-Fibonacci control) - // Equal description-cost control; MDL Shaw 2025 (2509.22445) - test b1_variant_d_random_rational_config - given (b1, b2, wd) = get_effective_betas(PhiVariant.RandomRat) - then approximately_equal(b1, RATIONAL_BETA1) // 4/7 ~ 0.5714 - and approximately_equal(b2, DEFAULT_BETA2) // 0.999 - and approximately_equal(wd, RATIONAL_WEIGHT_DECAY) // 7/30 ~ 0.2333 - - // B1e: Falsification integration test (runs the ablation scaffold) - // - // HEADLINE: phi-specificity is FALSIFIED if any non-phi control - // matches variant (a) within Monte-Carlo error on equal budget. - // - // This test verifies the scaffold produces distinct BPB measurements - // (not that phi wins -- that is the experimental outcome to determine). - test b1_control_ablation_produces_distinct_bpb - given n_steps = 1000 // mini-run for CI; full run uses 81000 - and seed = 43 - and bpb_a = run_adamw(PhiVariant.Canonical, n_steps, seed) - and bpb_b = run_adamw(PhiVariant.Damped, n_steps, seed) - and bpb_c = run_adamw(PhiVariant.TunedStd, n_steps, seed) - and bpb_d = run_adamw(PhiVariant.RandomRat, n_steps, seed) - then all_finite([bpb_a, bpb_b, bpb_c, bpb_d]) - and all_positive([bpb_a, bpb_b, bpb_c, bpb_d]) - // Note: NOT asserting bpb_a < bpb_b, bpb_c, bpb_d -- that would - // beg the question. The experiment records which variant wins. - - // ========================================================= - // EXISTING TDD: Invariants (preserved) - // ========================================================= - - invariant step_preserves_param_count_v2 - given config = any_adamw_config() - and state = init(config, 100) - and params = random_input(100) - and grads = random_input(100) - when result = step(state, params, grads) - then result.updated_params.len == 100 - - invariant moments_have_same_length_v2 - given config = any_adamw_config() - and state = init(config, 50) - then state.m.len == 50 - and state.v.len == 50 - - invariant first_moment_remains_bounded_v2 - given config = any_adamw_config() - and state = init(config, 10) - and params = finite_input(10) - and grads = bounded_input(10, -1.0, 1.0) - and result = step(state, params, grads) - then all(result.m, is_finite) - - invariant second_moment_non_negative_v2 { - const config = any_adamw_config(); - const state = init(config, 10); - const params = finite_input(10); - const grads = finite_input(10); - const result = step(state, params, grads); - // was: then all(result.v, fn(x) x >= 0.0) -- t27 has no closures, so - // the predicate is checked by an explicit loop over every element - var i : usize = 0; - while (i < result.v.len) { - assert(result.v[i] >= 0.0); - i = i + 1; - } + if (r.v != 0.5) { + return 7; } - - invariant bias_correction_in_zero_one_range_v2 - given beta = random_gf16_in_range(0.0, 0.999) - and t = random_u64_in_range(1, 1000) - when bc = compute_bias_correction(beta, t) - then bc >= 0.0 - and bc <= 1.0 - - invariant weight_decay_preserves_sign_v2 - given params = [-1.0, 0.0, 1.0] - and lr = 0.01 - and weight_decay = 0.1 - when result = apply_weight_decay(params, lr, weight_decay) - then result[0] < 0.0 - and result[2] > 0.0 - - invariant amsgrad_v_max_non_decreasing_v2 { - const config = AdamWConfig{.amsgrad = true}; - const state = init(config, 10); - const params = random_input(10); - const grads = random_input(10); - const result1 = step(state, params, grads); - const result2 = step(state, result1.updated_params, grads); - // was: then all_pairs(result1.v_max, result2.v_max, fn(a, b) b >= a) - // -- t27 has no closures; the pairs are taken element-wise (index i of - // both arrays) by an explicit loop - var i : usize = 0; - while (i < result1.v_max.len) { - assert(result2.v_max[i] >= result1.v_max[i]); - i = i + 1; - } + if (r.x != 0.75) { + return 7; } - - invariant denominator_positive_v2 - given v = 0.0 - and epsilon = 1e-8 - when denom = sqrt(v) + epsilon - then denom > 0.0 - - // ========================================================= - // PHASE B1: New invariants for reconciliation - // ========================================================= - - // B1-INV-1: Canonical beta1 is in valid range for convergence theory - // Heavy-ball / EMA window 1/(1-beta1); valid range (0, 1) open. - // [Open conjecture] that phi^(-1) is "optimal" -- no theorem backs this. - invariant b1_canonical_beta1_valid - then PHI_CANONICAL_BETA1 > 0.0 - and PHI_CANONICAL_BETA1 < 1.0 - - // B1-INV-2: Canonical weight_decay is in sensible range - invariant b1_canonical_wd_valid - then PHI_CANONICAL_WEIGHT_DECAY > 0.0 - and PHI_CANONICAL_WEIGHT_DECAY < 1.0 - - // B1-INV-3: Canonical beta1 > damped beta1 (ordering fact, not a claim) - invariant b1_canonical_greater_than_damped - then PHI_CANONICAL_BETA1 > PHI_DAMPED_BETA1 - - // B1-INV-4: All four control variants have distinct beta1 values - // (ensures the ablation is not degenerate) - invariant b1_variants_distinct - then PHI_CANONICAL_BETA1 != PHI_DAMPED_BETA1 - and PHI_CANONICAL_BETA1 != DEFAULT_BETA1 - and PHI_CANONICAL_BETA1 != RATIONAL_BETA1 - and PHI_DAMPED_BETA1 != DEFAULT_BETA1 - and PHI_DAMPED_BETA1 != RATIONAL_BETA1 - and DEFAULT_BETA1 != RATIONAL_BETA1 - - // B1-INV-5: Trinity identity L5 -- [Verified] -- must hold to fp tolerance - invariant b1_trinity_l5_verified - given phi = 1.618033988749895 - when lhs = phi * phi + 1.0 / (phi * phi) - then approximately_equal(lhs, 3.0) // tolerance 1e-10 - - // B1-INV-6: Weight decay values are distinct across variants that differ - invariant b1_canonical_wd_differs_from_default - then PHI_CANONICAL_WEIGHT_DECAY != DEFAULT_WEIGHT_DECAY - - // B1-INV-7: No phi anchor is presented as Verified (enforces honesty) - // This invariant is meta-level: code review must confirm no .t27 comment - // or doc labels PHI_CANONICAL_BETA1 or PHI_DAMPED_BETA1 as [Verified]. - // [Verified] is reserved for phi^2 + phi^(-2) = 3 (Law L5 only). - invariant b1_phi_anchors_labeled_open_conjecture - // Verified in code review: PHI_CANONICAL_BETA1, PHI_CANONICAL_WEIGHT_DECAY, - // PHI_DAMPED_BETA1, PHI_DAMPED_BETA2 all carry [Open conjecture] in comments. - assert true // placeholder; enforcement is by CI lint on claim labels - - // ========================================================= - // EXISTING TDD: Benchmarks (preserved) - // ========================================================= - - bench init_small_v2 - given config = any_adamw_config() - when result = init(config, 100) - then elapsed_time_us < 50 - - bench init_large_v2 - given config = any_adamw_config() - when result = init(config, 1000000) - then elapsed_time_ms < 20 - - bench step_small_v2 - given config = any_adamw_config() - and state = init(config, 100) - and params = random_input(100) - and grads = random_input(100) - when result = step(state, params, grads) - then elapsed_time_us < 200 - - bench step_medium_v2 - given config = any_adamw_config() - and state = init(config, 10000) - and params = random_input(10000) - and grads = random_input(10000) - when result = step(state, params, grads) - then elapsed_time_ms < 10 - - bench step_large_v2 - given config = any_adamw_config() - and state = init(config, 1000000) - and params = random_input(1000000) - and grads = random_input(1000000) - when result = step(state, params, grads) - then elapsed_time_ms < 200 - - bench compute_bias_correction_v2 - given beta = 0.9 - and t = 100 - when result = compute_bias_correction(beta, t) - then elapsed_time_ns < 50 - - bench amsgrad_update_v2 - given v = random_positive_gf16() - and v_max = random_positive_gf16() - when result = amsgrad_update(v, v_max) - then elapsed_time_ns < 10 - - // ========================================================= - // PHASE B1: Ablation benchmark - // ========================================================= - - // B1-BENCH: Compare throughput of all four variants at equal step count. - // Purpose: confirm no variant has a hidden compute advantage. - bench b1_ablation_equal_budget_throughput - given n = 10000 - and steps = 1000 - and config_a = AdamWConfig{.phi_variant = PhiVariant.Canonical} - and config_b = AdamWConfig{.phi_variant = PhiVariant.Damped} - and config_c = AdamWConfig{.phi_variant = PhiVariant.TunedStd} - and config_d = AdamWConfig{.phi_variant = PhiVariant.RandomRat} - when throughput_a = steps_per_second(config_a, n, steps) - and throughput_b = steps_per_second(config_b, n, steps) - and throughput_c = steps_per_second(config_c, n, steps) - and throughput_d = steps_per_second(config_d, n, steps) - // All variants must be within 5% of each other (same algo, different scalars) - then ratio(throughput_a, throughput_c) > 0.95 - and ratio(throughput_b, throughput_c) > 0.95 - and ratio(throughput_d, throughput_c) > 0.95 - - // ========================================================= - // Mathematical Notes - // ========================================================= - // - // AdamW Update Rule (Loshchilov & Hutter 2019): - // 1. Apply weight decay (decoupled): - // theta_t = theta_{t-1} * (1 - eta * lambda) - // - // 2. Update moments: - // m_t = beta1 * m_{t-1} + (1 - beta1) * g_t - // v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2 - // - // 3. Bias correction: - // m_hat_t = m_t / (1 - beta1^t) - // v_hat_t = v_t / (1 - beta2^t) - // - // 4. Parameter update: - // theta_t = theta_t - eta * m_hat_t / (sqrt(v_hat_t) + eps) - // - // EMA window: 1/(1-beta1) steps. For phi^(-1): window = phi^2 ~ 2.618. - // For 0.9: window = 10. For 0.9/phi: window ~ 2.24. - // Shorter window means faster response to gradient changes but - // noisier estimates -- no theory singles out phi as optimal. - // - // Literature context (used as BASELINES, not as phi validation): - // Xie & Li 2024 (2404.04454): AdamW has implicit L-inf gradient bias; - // relevant to comparing variants with different beta1 values. - // Wang & Aitchison 2025 (2405.13698): weight-decay scaling matters at scale. - // Chen 2025 (2510.12402): Cautious Weight Decay; reference for wd ablation. - // Li & Talwalkar 2019 (1902.07638): random-search is a strong HP baseline. - // Melis/Dyer/Blunsom 2017 (1707.05589): equal-budget comparison protocol. - // Hoffmann 2022 (2203.15556): BPB / Chinchilla scaling; our metric anchor. - // - // ========================================================= + return 0; +} + +/// Minimise (x - 3)^2 from 0 at lr 0.1. +fn descend(steps: u32, wd: f32) -> f32 { + var params : [1]f32 = [0.0]; + var mom : [3]f32 = [0.0, 0.0, 0.0]; + const c : AdamWConfig = plain_config(0.1, wd, false); + var t : u32 = 1; + while (t <= steps) { + const grads : [1]f32 = [2.0 * (params[0] - 3.0)]; + const ok : bool = adamw_step(¶ms, &grads, &mom, 0, 1, 2, t, c); + t = t + 1; + } + return params[0]; +} + +/// The error code of allocating state for 3 parameters from an arena of cap slots. +fn arena_code(cap: i32) -> i32 { + var backing : [24]f32 = undefined; + fill_slots(&backing, 0, 24, 9.0); + return state_alloc(arena_new(cap), &backing, 3).code; +} + +/// The arena's top after that allocation. +fn arena_top(cap: i32) -> i32 { + var backing : [24]f32 = undefined; + fill_slots(&backing, 0, 24, 9.0); + return state_alloc(arena_new(cap), &backing, 3).arena.top; +} + +/// Allocate state for 3 parameters from a backing array of 24 slots, step the three parameters +/// once through the arena offsets, and return parameter `index` (which 0) or the first-moment +/// slot of that parameter (which 1). +fn arena_driven(index: usize, which: i32, amsgrad: bool) -> f32 { + var backing : [24]f32 = undefined; + fill_slots(&backing, 0, 24, 9.0); + var params : [3]f32 = [1.0, 2.0, 3.0]; + const grads : [3]f32 = [0.1, 0.2, 0.3]; + const s : AdamWState = state_alloc(arena_new(24), &backing, 3); + const ok : bool = adamw_step(¶ms, &grads, &backing, s.m_off, s.v_off, s.x_off, 1, plain_config(0.001, 0.0, amsgrad)); + if (which == 1) { + const base : usize = @intCast(s.m_off); + return backing[base + index]; + } + return params[index]; +} + +/// The same three parameters stepped through three separate caller-owned arrays. +fn separate_driven(index: usize, amsgrad: bool) -> f32 { + var params : [3]f32 = [1.0, 2.0, 3.0]; + const grads : [3]f32 = [0.1, 0.2, 0.3]; + var mom : [9]f32 = [0.0] ** 9; + const ok : bool = adamw_step(¶ms, &grads, &mom, 0, 3, 6, 1, plain_config(0.001, 0.0, amsgrad)); + return params[index]; +} + +fn phi_identity() -> f32 { + return PHI * PHI + 1.0 / (PHI * PHI); +} + +// ---- tests ---- + +// Bias correction makes the first step lr * sign(g), whatever |g| is. +test the_first_step_is_lr_times_the_sign_of_the_gradient { + const c : AdamWConfig = plain_config(0.001, 0.0, false); + const small : Probe = probe_step(1.0, 0.1, 0.0, 0.0, 0.0, 1, c); + const large : Probe = probe_step(1.0, 100.0, 0.0, 0.0, 0.0, 1, c); + const neg : Probe = probe_step(1.0, -5.0, 0.0, 0.0, 0.0, 1, c); + assert small.ok; + assert @abs(small.p - 0.999) < 0.000001; + assert @abs(large.p - 0.999) < 0.000001; + assert @abs(neg.p - 1.001) < 0.000001; +} + +// m = (1 - b1) g and v = (1 - b2) g^2 after the first step, and x stays 0 without amsgrad. +test the_moments_after_one_step_are_exact { + const c : AdamWConfig = plain_config(0.001, 0.0, false); + const r : Probe = probe_step(1.0, 0.3, 0.0, 0.0, 0.0, 1, c); + assert @abs(r.m - 0.03) < 0.0000001; + assert @abs(r.v - 0.00009) < 0.0000001; + assert r.x == 0.0; +} + +// Decoupled: with a zero gradient the parameter shrinks by (1 - lr wd) per step and the +// moments never see the decay. Adam with L2 would put wd p into m and v. +test the_decay_is_decoupled_from_the_moments { + assert @abs(zero_grad_run(1, 0) - 1.9) < 0.00001; + assert @abs(zero_grad_run(3, 0) - 1.71475) < 0.00001; + assert zero_grad_run(3, 1) == 0.0; + assert zero_grad_run(3, 2) == 0.0; +} + +// With amsgrad the denominator uses the largest v seen: after gradients 2.0 then 0.0, v has +// decayed to 0.003996 but x still holds 0.004, so the second step is a little shorter. +test amsgrad_keeps_the_largest_second_moment { + assert @abs(two_step_run(true, 1) - 0.003996) < 0.000001; + assert @abs(two_step_run(true, 2) - 0.004) < 0.000001; + assert two_step_run(false, 2) == 0.0; + const gap : f32 = two_step_run(true, 0) - two_step_run(false, 0); + assert gap > 0.00002; + assert gap < 0.00005; +} + +// NaN, infinity and step 0 are refused, and a refused step changes nothing. +test a_bad_gradient_or_step_zero_is_refused { + assert applied(1.5, 1) == 1; + assert applied(-2.3, 4) == 1; + assert applied(ratio(0.0, 0.0), 1) == 0; + assert applied(ratio(1.0, 0.0), 1) == 0; + assert applied(ratio(-1.0, 0.0), 1) == 0; + assert applied(1.5, 0) == 0; +} + +test it_descends_a_quadratic { + assert @abs(descend(500, 0.0) - 3.0) < 0.01; +} + +// Decoupled decay pulls the minimiser toward 0, without removing it from the basin. +test decay_pulls_the_minimiser_toward_zero { + const free : f32 = descend(500, 0.0); + const decayed : f32 = descend(500, 0.1); + assert decayed < free - 0.05; + assert decayed > 0.0; +} + +test the_four_ablation_arms_are_what_the_comment_says { + const canon : AdamWConfig = variant_config(VARIANT_CANONICAL, 0.001); + assert @abs(canon.beta1 - 0.618034) < 0.000001; + assert @abs(canon.weight_decay - 0.236068) < 0.000001; + const damped : AdamWConfig = variant_config(VARIANT_DAMPED, 0.001); + assert @abs(damped.beta1 - 0.556231) < 0.000001; + assert damped.weight_decay == DEFAULT_WEIGHT_DECAY; + const tuned : AdamWConfig = variant_config(VARIANT_TUNED_STD, 0.001); + assert tuned.beta1 == 0.9; + assert tuned.weight_decay == DEFAULT_WEIGHT_DECAY; + const rat : AdamWConfig = variant_config(VARIANT_RANDOM_RAT, 0.001); + assert @abs(rat.beta1 - 0.571429) < 0.000001; + assert @abs(rat.weight_decay - 0.233333) < 0.000001; + assert canon.beta2 == 0.999; + assert rat.beta2 == 0.999; + assert canon.learning_rate == 0.001; +} + +test an_unknown_variant_is_the_standard_arm { + const odd : AdamWConfig = variant_config(99, 0.002); + assert odd.beta1 == 0.9; + assert odd.weight_decay == DEFAULT_WEIGHT_DECAY; + assert odd.learning_rate == 0.002; +} + +test the_phi_identity_is_the_only_verified_phi_fact { + assert @abs(phi_identity() - 3.0) < 0.000001; + assert @abs(PHI * PHI_INV - 1.0) < 0.000001; + assert @abs(PHI_INV * PHI_INV * PHI_INV - PHI_INV_CUBED) < 0.000001; + assert PHI_INV > PHI_DAMPED_BETA1; + assert RATIONAL_BETA1 != PHI_INV; + assert RATIONAL_BETA1 != PHI_DAMPED_BETA1; +} + +// The arena path: three buffers of 3 slots at offsets 0, 3, 6, all zeroed from a backing array +// that held 9.0, and stepping through them gives exactly the separate-arrays result. +test the_state_comes_from_the_arena { + assert arena_code(24) == 0; + assert arena_top(24) == 9; + assert arena_driven(1, 0, false) == separate_driven(1, false); + assert arena_driven(2, 0, false) == separate_driven(2, false); + assert @abs(arena_driven(2, 1, false) - 0.03) < 0.0000001; + assert arena_driven(1, 0, true) == separate_driven(1, true); + assert arena_driven(2, 0, true) == separate_driven(2, true); +} + +// 3 buffers of 3 need 9 slots; with 8 the third fails and the arena is as it came in. +test an_arena_too_small_is_a_code_not_a_trap { + assert arena_code(8) == 1; + assert arena_top(8) == 0; + assert arena_code(9) == 0; + assert arena_top(9) == 9; +} + +invariant adamw_defaults_are_valid { + assert DEFAULT_BETA1 < 1.0; + assert DEFAULT_BETA2 < 1.0; + assert DEFAULT_BETA1 < DEFAULT_BETA2; + assert DEFAULT_EPSILON > 0.0; + assert DEFAULT_WEIGHT_DECAY >= 0.0; + assert DEFAULT_WEIGHT_DECAY < 1.0; +}