nx_nofloat_muon_gate.nx source

1// nx_nofloat_muon_gate.nx -- gate: the FULL Per-Head Muon optimizer loop (momentum + NS-orthogonalize + 2// apply), deterministic, per-head. Honest scope: proves the OPTIMIZER STEP mechanics + the novel exceed 3// (bit-exact), NOT a full training convergence run (that rides the GPU seat F101). 4// T1 MOMENTUM heavy-ball: repeated aligned grads grow the buffer beyond a single grad (geometric sum) 5// T2 ORTHOGONALIZED update: U = NS(momentum) has far lower orthogonality error than the raw momentum 6// T3 UPDATE APPLIED EXACTLY: W_after == W_before - lr*U elementwise (mechanical correctness) 7// T4 DETERMINISTIC: two identical per-head runs -> bit-identical weights (the exceed: float Muon drifts) 8// T5 PER-HEAD INDEPENDENCE: perturbing head-0's grad changes head-0's W but NOT head-1's W 9// T6 STABLE: 12 steps stay bounded (no blowup) 10// license_tier: ORIGINAL No hw writes (Rule 26). expect_exit: 0 11import "nx_nofloat_muon.nx" 12import "nx_nofloat_muon_ns.nx" 13import "nx_gate_verdict.nx" 14import "nx_syscalls.nx" 15 16const UG_Q: i64 = 65536 17const UG_MU: i64 = 58982 // 0.9 momentum 18const UG_LR: i64 = 6553 // 0.1 learning rate 19 20func ug_abs(x: i64) -> i64 { if x < 0 { return 0 - x } return x } 21func ug_copy(src: *i64, dst: *i64, nn: i64) -> i64 { var i: i64 = 0; while i < nn { dst[i] = src[i]; i = i + 1 } return 0 } 22// |X|_1 sum of abs 23func ug_l1(x: *i64, nn: i64) -> i64 { var s: i64 = 0; var i: i64 = 0; while i < nn { s = s + ug_abs(x[i]); i = i + 1 } return s } 24 25func main() -> i64 { 26 let ctr: *i64 = gv_ctr() 27 gv_head("nx_nofloat_muon gate -- full Per-Head Muon optimizer loop (momentum + NS + apply), deterministic" as *u8) 28 let n: i64 = 4 29 let nn: i64 = n * n 30 31 // a well-conditioned-ish gradient matrix (Q16) 32 let grad: *i64 = sys_mmap(64*8) as *i64 33 var i: i64 = 0 34 while i < nn { grad[i] = (((i % 4) + 1) * (UG_Q / 5)) - (UG_Q/3); i = i + 1 } 35 let u: *i64 = sys_mmap(64*8) as *i64 36 37 // ---- T1: momentum heavy-ball accumulation ---- 38 let buf1: *i64 = sys_mmap(64*8) as *i64 39 ug_copy(grad, buf1, nn) // buf = grad (first step, no prior momentum) -- start from grad-sized 40 i = 0 41 while i < nn { buf1[i] = 0; i = i + 1 } 42 mu_momentum(buf1, grad, UG_MU, nn); let l1a: i64 = ug_l1(buf1, nn) // after 1 43 mu_momentum(buf1, grad, UG_MU, nn); mu_momentum(buf1, grad, UG_MU, nn); let l1c: i64 = ug_l1(buf1, nn) // after 3 44 var t1: i64 = 0 45 if l1c > (l1a + (l1a/2)) { t1 = 1 } // 3-step momentum > 1.5x single-step (geometric growth) 46 gv_check("T1 MOMENTUM heavy-ball: 3 aligned grads accumulate > 1.5x a single grad" as *u8, t1, ctr) 47 48 // ---- T2: NS-orthogonalized update vs raw momentum ---- 49 let bufo: *i64 = sys_mmap(64*8) as *i64 50 ug_copy(grad, bufo, nn) 51 mu_momentum(bufo, grad, UG_MU, nn) // some momentum matrix 52 let errbuf: i64 = mn_ortho_err(bufo, n) // orthogonality error of the RAW momentum 53 mn_orthogonalize(bufo, u, n) // U = NS(momentum) 54 let erru: i64 = mn_ortho_err(u, n) // error of the orthogonalized update 55 var t2: i64 = 0 56 if erru < errbuf { if erru < (nn * UG_Q) { t2 = 1 } } // NS improved + U approx-orthogonal (avg<1/entry) 57 gv_check("T2 ORTHOGONALIZED update: NS(momentum) far lower ortho-error than raw momentum" as *u8, t2, ctr) 58 59 // ---- T3: update applied exactly (W == W - lr*U) ---- 60 let w: *i64 = sys_mmap(64*8) as *i64 61 let buf: *i64 = sys_mmap(64*8) as *i64 62 let wsave: *i64 = sys_mmap(64*8) as *i64 63 i = 0 64 while i < nn { w[i] = ((i*7) % 11) * (UG_Q/9); buf[i] = 0; i = i + 1 } 65 ug_copy(w, wsave, nn) 66 // expected: replicate momentum->NS->apply on copies 67 let bufe: *i64 = sys_mmap(64*8) as *i64 68 let ue: *i64 = sys_mmap(64*8) as *i64 69 let we: *i64 = sys_mmap(64*8) as *i64 70 ug_copy(w, we, nn); ug_copy(buf, bufe, nn) 71 mu_momentum(bufe, grad, UG_MU, nn) 72 mn_orthogonalize(bufe, ue, n) 73 i = 0 74 while i < nn { we[i] = we[i] - ((UG_LR * ue[i]) >> 16); i = i + 1 } 75 // actual step 76 mu_step(w, buf, grad, UG_MU, UG_LR, n, u) 77 var t3: i64 = 1 78 i = 0 79 while i < nn { if w[i] != we[i] { t3 = 0 } i = i + 1 } 80 // also confirm it actually moved 81 var moved: i64 = 0 82 i = 0 83 while i < nn { if w[i] != wsave[i] { moved = 1 } i = i + 1 } 84 if moved == 0 { t3 = 0 } 85 gv_check("T3 UPDATE APPLIED EXACTLY: W_after == W_before - lr*NS(momentum), and W moved" as *u8, t3, ctr) 86 87 // ---- T4: deterministic per-head step ---- 88 let nh: i64 = 2 89 let hd: i64 = 4 90 let phn: i64 = nh*hd*hd 91 let wA: *i64 = sys_mmap(128*8) as *i64 92 let bA: *i64 = sys_mmap(128*8) as *i64 93 let gA: *i64 = sys_mmap(128*8) as *i64 94 let wB: *i64 = sys_mmap(128*8) as *i64 95 let bB: *i64 = sys_mmap(128*8) as *i64 96 let gB: *i64 = sys_mmap(128*8) as *i64 97 i = 0 98 while i < phn { let gv: i64 = (((i%5)+1)*(UG_Q/6)) - (UG_Q/4); wA[i]=(i%9)*(UG_Q/10); bA[i]=0; gA[i]=gv; wB[i]=(i%9)*(UG_Q/10); bB[i]=0; gB[i]=gv; i=i+1 } 99 mu_perhead_step(wA, bA, gA, UG_MU, UG_LR, nh, hd, u) 100 mu_perhead_step(wB, bB, gB, UG_MU, UG_LR, nh, hd, u) 101 var t4: i64 = 1 102 i = 0 103 while i < phn { if wA[i] != wB[i] { t4 = 0 } i = i + 1 } 104 gv_check("T4 DETERMINISTIC per-head step: two runs bit-identical (float Muon drifts)" as *u8, t4, ctr) 105 106 // ---- T5: per-head independence ---- 107 let wC: *i64 = sys_mmap(128*8) as *i64 108 let bC: *i64 = sys_mmap(128*8) as *i64 109 let gC: *i64 = sys_mmap(128*8) as *i64 110 i = 0 111 while i < phn { wC[i]=(i%9)*(UG_Q/10); bC[i]=0; gC[i]=gA[i]; i=i+1 } // gA already consumed bA; rebuild grads 112 i = 0 113 while i < phn { let gv: i64 = (((i%5)+1)*(UG_Q/6)) - (UG_Q/4); gC[i]=gv; i=i+1 } 114 gC[0] = gC[0] + (UG_Q/2) // perturb ONLY head-0 (block 0 = indices 0..hd*hd) 115 // fresh baseline for compare 116 let wD: *i64 = sys_mmap(128*8) as *i64 117 let bD: *i64 = sys_mmap(128*8) as *i64 118 let gD: *i64 = sys_mmap(128*8) as *i64 119 i = 0 120 while i < phn { wD[i]=(i%9)*(UG_Q/10); bD[i]=0; let gv: i64 = (((i%5)+1)*(UG_Q/6)) - (UG_Q/4); gD[i]=gv; i=i+1 } 121 i = 0 122 while i < phn { wC[i]=(i%9)*(UG_Q/10); bC[i]=0; i=i+1 } 123 mu_perhead_step(wC, bC, gC, UG_MU, UG_LR, nh, hd, u) // perturbed head-0 124 mu_perhead_step(wD, bD, gD, UG_MU, UG_LR, nh, hd, u) // baseline 125 let blk: i64 = hd*hd 126 var head0_changed: i64 = 0 127 var head1_same: i64 = 1 128 i = 0 129 while i < blk { if wC[i] != wD[i] { head0_changed = 1 } i = i + 1 } 130 i = blk 131 while i < 2*blk { if wC[i] != wD[i] { head1_same = 0 } i = i + 1 } 132 var t5: i64 = 0 133 if head0_changed == 1 { if head1_same == 1 { t5 = 1 } } 134 gv_check("T5 PER-HEAD INDEPENDENCE: perturbing head-0 grad moves head-0 W only, head-1 untouched" as *u8, t5, ctr) 135 136 // ---- T6: stable over 12 steps ---- 137 let wS: *i64 = sys_mmap(64*8) as *i64 138 let bS: *i64 = sys_mmap(64*8) as *i64 139 i = 0 140 while i < nn { wS[i]=(i%7)*(UG_Q/8); bS[i]=0; i=i+1 } 141 var step: i64 = 0 142 while step < 12 { mu_step(wS, bS, grad, UG_MU, UG_LR, n, u); step = step + 1 } 143 var t6: i64 = 1 144 i = 0 145 while i < nn { if ug_abs(wS[i]) > 1000000000 { t6 = 0 } i = i + 1 } 146 gv_check("T6 STABLE over 12 optimizer steps (bounded, no blowup)" as *u8, t6, ctr) 147 148 let rc: i64 = gv_verdict("NOFLOAT-MUON-GATE" as *u8, ctr, "full Per-Head Muon optimizer: momentum + NS-orthogonalized update, applied exactly, per-head independent, bit-exact deterministic no-float" as *u8) 149 sys_exit(rc) 150 return rc 151}