nx_lowrank_weight_spectrum.nx source
↩ module page · 185 lines · 8419 B
1// nx_lowrank_weight_spectrum.nx -- the FORWARD-FREE VRAM lever: are the real model's WEIGHT matrices low-rank?
2// (FlashSVD / over-parameterization: arXiv 2508.01506). Unlike the KV-cache lever, this needs NO forward --
3// just load a weight tensor + measure its singular spectrum -- so it iterates without the slow software-f32
4// forward. Measures attn_q (square 896x896) AND ffn_gate (4864x896, the big MLP weight where most params live)
5// for Qwen2.5-0.5B: cumulative covariance energy in the top-r of the smaller (cols) dimension. High energy at
6// small r => low-rank => SVD-compressible. Sovereign (nx_cc->nxasm, no gcc). NOTE: SVD funcs copied from
7// nx_lowrank_kv_real (proven f32+reorth) -- TODO extract both to nx_lowrank_svd.nx (DRY #15). license_tier: ORIGINAL
8import "nx_syscalls.nx"
9import "nx_gguf.nx"
10import "nx_gguf_load.nx"
11import "nx_gguf_meta.nx"
12import "nx_gguf_load_f32.nx"
13import "nx_f32.nx"
14import "nx_f32_div.nx"
15import "nx_f32_hw.nx"
16
17func w_puts(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
18func w_putn(v: i64) -> i64 {
19 if v == 0 { sys_write(1, "0" as *u8, 1); return 0 }
20 var m: i64 = v
21 if m < 0 { sys_write(1, "-" as *u8, 1); m = 0 - m }
22 let d: *u8 = sys_mmap(24); var k: i64 = 0
23 while m > 0 { d[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
24 var j: i64 = k - 1
25 while j >= 0 { sys_write(1, ((d as i64)+j) as *u8, 1); j = j - 1 }
26 return 0
27}
28
29// ---- f32 eigen-spectrum machinery (copied verbatim from nx_lowrank_kv_real; f32 + reorthogonalization) ----
30func r_is_zero(x: i64) -> i64 { if (x & 0x7fffffff) == 0 { return 1 } return 0 }
31func r_dot(a: *i64, b: *i64, n: i64) -> i64 {
32 var s: i64 = 0
33 var i: i64 = 0
34 while i < n { s = nx_f32_add(s, nx_f32_mul(a[i], b[i])); i = i + 1 }
35 return s
36}
37func r_normalize(v: *i64, n: i64) -> i64 {
38 let ns: i64 = r_dot(v, v, n)
39 if r_is_zero(ns) == 1 { return -1 }
40 let nrm: i64 = nx_f32_sqrt(ns)
41 if r_is_zero(nrm) == 1 { return -1 }
42 let inv: i64 = nx_f32_div(f32_of(1), nrm)
43 var i: i64 = 0
44 while i < n { v[i] = nx_f32_mul(v[i], inv); i = i + 1 }
45 return 0
46}
47func r_gmatvec(G: *i64, v: *i64, d: i64, out: *i64) -> i64 {
48 var a: i64 = 0
49 while a < d {
50 var s: i64 = 0
51 var b: i64 = 0
52 while b < d { s = nx_f32_add(s, nx_f32_mul(G[a * d + b], v[b])); b = b + 1 }
53 out[a] = s
54 a = a + 1
55 }
56 return 0
57}
58func r_deflate1(v: *i64, basis: *i64, found: i64, d: i64) -> i64 {
59 var j: i64 = 0
60 while j < found {
61 var s: i64 = 0
62 var c: i64 = 0
63 while c < d { s = nx_f32_add(s, nx_f32_mul(v[c], basis[j * d + c])); c = c + 1 }
64 c = 0
65 while c < d { v[c] = nx_f32_sub(v[c], nx_f32_mul(s, basis[j * d + c])); c = c + 1 }
66 j = j + 1
67 }
68 return 0
69}
70func r_reorth(v: *i64, basis: *i64, found: i64, d: i64) -> i64 {
71 r_deflate1(v, basis, found, d); r_deflate1(v, basis, found, d); return 0
72}
73func r_power_one(G: *i64, d: i64, basis: *i64, found: i64, v: *i64, w: *i64) -> i64 {
74 var c: i64 = 0
75 while c < d { v[c] = 0; c = c + 1 }
76 v[found % d] = f32_of(1)
77 r_reorth(v, basis, found, d)
78 if r_normalize(v, d) < 0 { return 0 }
79 var it: i64 = 0
80 while it < 64 {
81 r_gmatvec(G, v, d, w)
82 r_reorth(w, basis, found, d)
83 if r_is_zero(r_dot(w, w, d)) == 1 { it = 64 } else {
84 if r_normalize(w, d) < 0 { it = 64 } else {
85 c = 0
86 while c < d { v[c] = w[c]; c = c + 1 }
87 it = it + 1
88 }
89 }
90 }
91 c = 0
92 while c < d { basis[found * d + c] = v[c]; c = c + 1 }
93 r_gmatvec(G, v, d, w)
94 return r_dot(v, w, d)
95}
96// gram of W[rows x cols] (row-major) over the 'rows' axis -> cols x cols covariance.
97// Caller passes the LARGER dim as rows so cov = min(rows,cols)^2 (cheap + captures the full rank).
98func r_gram(W: *i64, rows: i64, cols: i64, G: *i64) -> i64 {
99 var a: i64 = 0
100 while a < cols {
101 var b: i64 = 0
102 while b < cols {
103 var s: i64 = 0
104 var t: i64 = 0
105 while t < rows { s = nx_f32_add(s, nx_f32_mul(W[t * cols + a], W[t * cols + b])); t = t + 1 }
106 G[a * cols + b] = s
107 b = b + 1
108 }
109 a = a + 1
110 }
111 return 0
112}
113func r_trace(G: *i64, d: i64) -> i64 {
114 var s: i64 = 0
115 var i: i64 = 0
116 while i < d { s = nx_f32_add(s, G[i * d + i]); i = i + 1 }
117 return s
118}
119
120// top-32 cumulative covariance energy (permille) of W[rows x cols]; gram over rows -> cols x cols spectrum.
121func w_spectrum(W: *i64, rows: i64, cols: i64, label: *u8) -> i64 {
122 let G: *i64 = sys_mmap(cols * cols * 8) as *i64
123 r_gram(W, rows, cols, G)
124 let total: i64 = r_trace(G, cols)
125 w_puts(label); w_puts(" ("); w_putn(rows); w_puts("x"); w_putn(cols); w_puts(", cov "); w_putn(cols); w_puts("^2) top-r: ")
126 if r_is_zero(total) == 1 { w_puts("(zero)\n"); return 0 }
127 let basis: *i64 = sys_mmap(32 * cols * 8) as *i64
128 let v: *i64 = sys_mmap(cols * 8) as *i64
129 let ww: *i64 = sys_mmap(cols * 8) as *i64
130 var cum: i64 = 0
131 var f: i64 = 0
132 while f < 32 {
133 let lam: i64 = r_power_one(G, cols, basis, f, v, ww)
134 cum = nx_f32_add(cum, lam)
135 if f == 0 { w_puts("r1="); w_putn(f32_int(nx_f32_mul(nx_f32_div(cum, total), f32_of(1000)))); w_puts(" ") }
136 if f == 3 { w_puts("r4="); w_putn(f32_int(nx_f32_mul(nx_f32_div(cum, total), f32_of(1000)))); w_puts(" ") }
137 if f == 7 { w_puts("r8="); w_putn(f32_int(nx_f32_mul(nx_f32_div(cum, total), f32_of(1000)))); w_puts(" ") }
138 if f == 15 { w_puts("r16="); w_putn(f32_int(nx_f32_mul(nx_f32_div(cum, total), f32_of(1000)))); w_puts(" ") }
139 if f == 31 { w_puts("r32="); w_putn(f32_int(nx_f32_mul(nx_f32_div(cum, total), f32_of(1000)))); w_puts("permille") }
140 f = f + 1
141 }
142 w_puts("\n")
143 return 0
144}
145
146// load + measure a named weight tensor (rows = larger dim_1/out, cols = smaller dim_0/in for these projections)
147func w_measure(buf: *u8, hdr: *NxGgufHeader, name: *u8, name_len: nx_int, label: *u8) -> i64 {
148 let idx: nx_int = nx_gguf_find_tensor(hdr, name, name_len)
149 if idx < 0 { w_puts(label); w_puts(": NOT FOUND\n"); return -1 }
150 let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, idx)
151 let nv_out: *i64 = sys_mmap(8) as *i64
152 let err: *i64 = sys_mmap(8) as *i64
153 let W: *i64 = nx_gguf_load_tensor_to_f32(buf, hdr, name, name_len, nv_out, err)
154 if W == (0 as *i64) { w_puts(label); w_puts(": load failed err="); w_putn(err[0]); w_puts("\n"); return -1 }
155 // dim_0 = in (cols, fast), dim_1 = out (rows). gram over the LARGER dim.
156 var rows: i64 = ti.dim_1
157 var cols: i64 = ti.dim_0
158 if cols > rows {
159 // transpose-view not supported by r_gram; for these weights dim_1(out) >= dim_0(in) holds, so just note
160 w_puts(label); w_puts(": cols>rows (skip; needs transposed gram)\n"); return -1
161 }
162 w_spectrum(W, rows, cols, label)
163 return 0
164}
165
166func main() -> i64 {
167 let path: *u8 = "/home/elderwesto/nx_stage/nx_real_model.gguf\x00"
168 let len_out: *i64 = sys_mmap(8) as *i64
169 let buf: *u8 = sys_read_file(path, len_out)
170 if buf == (0 as *u8) { w_puts("weight: model read failed\n"); return 10 }
171 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
172 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { w_puts("weight: gguf parse failed\n"); return 20 }
173 w_puts("=== nx_lowrank_weight_spectrum: are real Qwen2.5-0.5B WEIGHT matrices low-rank? (FORWARD-FREE) ===\n")
174 w_puts(" top-r energy in the SMALLER (cols) dim; high permille at small r => low-rank => SVD-compressible (FlashSVD)\n\n")
175
176 w_measure(buf, hdr, "blk.0.attn_q.weight\x00", 19, " attn_q L0" as *u8) // attention (small, 896x896)
177 w_measure(buf, hdr, "blk.0.ffn_gate.weight\x00", 21, " ffn_gate L0" as *u8) // BIG MLP weight (most params)
178 w_measure(buf, hdr, "blk.0.ffn_up.weight\x00", 19, " ffn_up L0" as *u8) // BIG MLP weight
179
180 w_puts("\n COMPRESSION: a rank-r approx of [rows x cols] = r*(rows+cols) vs rows*cols. For ffn (4864x896) at\n")
181 w_puts(" r=32: 32*(4864+896)=184320 vs 4358144 = ~24x; attn (896x896) at r=32: 32*1792=57344 vs 802816 = ~14x.\n")
182 w_puts(" HONEST: top-32 only (full 90%-rank slow if not low-rank); layer 0; software gram (ffn ~195s). Forward-FREE.\n")
183 sys_exit(0)
184 return 0
185}