code wiki / _hdl_build / nx_nofloat_muon_ns.nx
`nx_nofloat_muon_ns.nx` source

↩ module page · 108 lines · 4567 B
1// nx_nofloat_muon_ns.nx -- SOVEREIGN no-float Newton-Schulz orthogonalization = the core of the Muon
2// optimizer (K3 F405 gap; operator 2026-07-19 "beyond SOTA, novel"). Muon = MomentUm Orthogonalized by
3// Newton-Schulz (Keller Jordan). NS quintic iteration drives a matrix toward its orthogonal polar factor:
4//   X0 = G / ||G||_F  (Frobenius-normalize so spectral norm <= 1)
5//   X_{k+1} = a*X + b*(A*X) + c*(A*A*X),  A = X*X^T,  (a,b,c) = (3.4445, -4.7750, 2.0315)
6// After ~5 iters X*X^T -> I. Done in PURE INTEGER Q16 (matmul accumulate-then-shift + integer isqrt) =
7// BIT-EXACT DETERMINISTIC orthogonalization -- a float training stack CANNOT reproduce this bit-for-bit
8// (non-associative accumulation); that is the novel exceed. Grounded: kellerjordan.github.io/posts/muon.
9// license_tier: ORIGINAL   No hw writes (Rule 26).
10import "nx_syscalls.nx"
11import "nx_vecmath.nx"
12
13const MN_Q: i64 = 65536           // Q16 one
14const MN_QBITS: i64 = 16
15const MN_NMAX: i64 = 8            // max matrix dim supported
16// NS quintic coeffs in Q16 (Muon): a=3.4445, b=-4.7750, c=2.0315
17const MN_A: i64 = 225755
18const MN_B: i64 = 0 - 312934
19const MN_C: i64 = 133136
20const MN_ITERS: i64 = 5
21
22// integer sqrt of a non-negative i64 (floor). Newton iteration.
23func mn_isqrt(v: i64) -> i64 { return vm_isqrt(v) }
24// C[n x n] = A[n x n] * B[n x n] in Q16 (accumulate i64, >>16). A,B,C distinct buffers.
25func mn_mm(a: *i64, b: *i64, c: *i64, n: i64) -> i64 {
26    var i: i64 = 0
27    while i < n {
28        var j: i64 = 0
29        while j < n {
30            var s: i64 = 0
31            var k: i64 = 0
32            while k < n { s = s + a[i*n+k] * b[k*n+j]; k = k + 1 }
33            c[i*n+j] = s >> MN_QBITS
34            j = j + 1
35        }
36        i = i + 1
37    }
38    return 0
39}
40// C = A^T (n x n)
41func mn_transpose(a: *i64, c: *i64, n: i64) -> i64 {
42    var i: i64 = 0
43    while i < n { var j: i64 = 0; while j < n { c[j*n+i] = a[i*n+j]; j = j + 1 } i = i + 1 }
44    return 0
45}
46// out = pa*A + pb*B + pc*C  (scalars in Q16, matrices in Q16) elementwise, n x n
47func mn_comb(a: *i64, b: *i64, c: *i64, out: *i64, pa: i64, pb: i64, pc: i64, n: i64) -> i64 {
48    var i: i64 = 0
49    let nn: i64 = n * n
50    while i < nn {
51        out[i] = ((pa * a[i]) >> MN_QBITS) + ((pb * b[i]) >> MN_QBITS) + ((pc * c[i]) >> MN_QBITS)
52        i = i + 1
53    }
54    return 0
55}
56// Newton-Schulz orthogonalize G (n x n, Q16) -> X (n x n, Q16), `iters` NS steps.
57// Muon NS = APPROXIMATE orthogonalization: singular values -> ~1 (fixed points 0.87/1.26), not exactly 1.
58func mn_orthogonalize_it(g: *i64, x: *i64, n: i64, iters: i64) -> i64 {
59    let nn: i64 = n * n
60    // Frobenius norm ||G||_F = sqrt(sum g^2).  g in Q16 -> g^2 in Q32; sum in i64; sqrt -> Q16 norm.
61    var ss: i64 = 0
62    var i: i64 = 0
63    while i < nn { let v: i64 = g[i] >> 8; ss = ss + v * v; i = i + 1 }   // (g>>8)^2 = g^2 >> 16 = Q16-scaled sumsq
64    let fro: i64 = mn_isqrt(ss)                                          // ~ ||G||_F in Q8-ish; normalize by ratio
65    if fro <= 0 { return 0 }
66    // X0 = G / ||G||_F  : x = g * Q / (fro<<8)   (fro was sqrt of Q16-scaled -> scale back)
67    let denom: i64 = fro << 8
68    i = 0
69    while i < nn { x[i] = (g[i] << MN_QBITS) / denom; i = i + 1 }
70    // iteration buffers
71    let xt: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64
72    let aM: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64
73    let ax: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64
74    let aax: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64
75    var it: i64 = 0
76    while it < iters {
77        mn_transpose(x, xt, n)          // X^T
78        mn_mm(x, xt, aM, n)             // A = X X^T
79        mn_mm(aM, x, ax, n)             // A X
80        mn_mm(aM, ax, aax, n)           // A A X
81        mn_comb(x, ax, aax, x, MN_A, MN_B, MN_C, n)   // X = aX + b(AX) + c(AAX)
82        it = it + 1
83    }
84    return 0
85}
86// default 5-iteration Muon NS
87func mn_orthogonalize(g: *i64, x: *i64, n: i64) -> i64 { return mn_orthogonalize_it(g, x, n, MN_ITERS) }
88// orthogonality error metric: sum |(X X^T)[i][j] - I[i][j]| in Q16 (0 = perfectly orthogonal)
89func mn_ortho_err(x: *i64, n: i64) -> i64 {
90    let xt: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64
91    let g: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64
92    mn_transpose(x, xt, n)
93    mn_mm(x, xt, g, n)
94    var e: i64 = 0
95    var i: i64 = 0
96    while i < n {
97        var j: i64 = 0
98        while j < n {
99            var t: i64 = g[i*n+j]
100            if i == j { t = t - MN_Q }
101            if t < 0 { t = 0 - t }
102            e = e + t
103            j = j + 1
104        }
105        i = i + 1
106    }
107    return e
108}
nx_nofloat_muon_ns.nx source

`nx_nofloat_muon_ns.nx` source
