nx_llm_diff_probe.nx source
↩ module page · 112 lines · 5344 B
1// nx_llm_diff_probe.nx -- NUMERICAL cross-check vs the transformers oracle (same GGUF). Dumps our forward's
2// last-token ("is") hidden-state L2-norm after embed + each of the 24 layers, as milli-units (x1000) for direct
3// compare to the reference norms [390, 4468, 7320, ... 279079]. The FIRST layer whose norm diverges localizes the
4// residual bug. Also prints embed[0..6] + after-L0[0..6] as f32 hex (ref embed ~[-0.005,0.015,-0.0067,...],
5// ref afterL0 ~[0.2343,-0.093,0.1314,...]). Fully sovereign forward; oracle is a throwaway debug reference. exp_exit:0
6import "nx_syscalls.nx"
7import "nx_tier.nx"
8import "nx_gguf.nx"
9import "nx_gguf_load.nx"
10import "nx_gguf_meta.nx"
11import "nx_f32.nx"
12import "nx_f32_rmsnorm.nx"
13import "nx_f32_kv_cache.nx"
14import "nx_f32_lazy_weight.nx"
15import "nx_f32_llama_block.nx"
16import "nx_f32_llama_block_v4.nx"
17import "nx_f32_llama_stack_v4.nx"
18import "nx_f32_llama_layer_lazy_load.nx"
19import "nx_f32_llm.nx"
20import "nx_f32_llm_v4.nx"
21import "nx_f32_llm_read_dims.nx"
22
23func pw(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
24func pn(v: i64) -> i64 { let bb: *u8=sys_mmap(28); var m: i64=v; if m<0{sys_write(1,"-" as *u8,1);m=0-m} let t: *u8=sys_mmap(28); var k: i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i: i64=0; while i<k{bb[i]=t[k-1-i];i=i+1} sys_write(1,bb,k); return 0 }
25func ph(v: i64) -> i64 { let b: *u8=sys_mmap(8); var i: i64=0; while i<8 { let nib: i64=(v>>((7-i)*4))&0xF; if nib<10 {b[i]=(48+nib) as u8} else {b[i]=(87+nib) as u8} i=i+1 } sys_write(1,b,8); return 0 }
26
27// int(v * 1000), truncating; v is f32 raw bits
28func milli(v: i64) -> i64 {
29 let m: i64 = nx_f32_mul(v, 0x447A0000) // *1000.0
30 let bits: i64 = m & 0xFFFFFFFF
31 let sign: i64 = (bits >> 31) & 1
32 let exp: i64 = ((bits >> 23) & 0xFF) - 127
33 if exp < 0 { return 0 }
34 let mant: i64 = (bits & 0x7FFFFF) | 0x800000
35 var iv: i64 = 0
36 if exp <= 23 { iv = mant >> (23 - exp) } else { iv = mant << (exp - 23) }
37 if sign == 1 { iv = 0 - iv }
38 return iv
39}
40func l2norm(h: *i64, n: nx_int) -> i64 {
41 var ss: i64 = 0
42 var i: nx_int = 0
43 while i < n { ss = nx_f32_add(ss, nx_f32_mul(h[i], h[i])); i = i + 1 }
44 return nx_f32_sqrt(ss)
45}
46func dump8(label: *u8, h: *i64) -> i64 {
47 pw(label); pw(" hex[0..6]:")
48 var i: nx_int = 0
49 while i < 6 { pw(" "); ph(h[i] & 0xFFFFFFFF); i = i + 1 }
50 pw("\n"); return 0
51}
52
53func main() -> i64 {
54 let path: *u8 = "/tmp/nx_real_model.gguf" as *u8
55 let len_out: *i64 = sys_mmap(8) as *i64
56 let buf: *u8 = sys_read_file(path, len_out)
57 if buf == (0 as *u8) { return 10 }
58 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
59 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
60 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
61 let oe: *i64 = sys_mmap(8) as *i64
62 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, oe) != NX_FLD_OK { return 30 }
63 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, oe) != NX_FLV4_OK { return 40 }
64
65 let eps: i64 = 0x358637BD
66 let attn_scale: i64 = 0x3E000000
67 let rope_base: i64 = 0x415D0EAB
68 let H: nx_int = model.hidden_dim
69 let NL: nx_int = model.n_layers
70
71 // Q6_K ffn_down direct dequant check vs REF down_proj[0,:8] + row0 norm.
72 let l0: *NxF32LlamaLayerLazy = model.layers[2] as *NxF32LlamaLayerLazy
73 let wd: *NxF32LazyWeight = l0.W_down
74 pw(" W_down dtype="); pn(wd.dtype_tag as i64); pw(" rows="); pn(wd.rows as i64); pw(" cols="); pn(wd.cols as i64); pw("\n")
75 dump8(" W_down.storage[0..6]" as *u8, wd.storage)
76 pw(" W_down row0(4864) norm*1000="); pn(milli(l2norm(wd.storage, model.ffn_dim))); pw("\n")
77
78 let toks: *i64 = sys_mmap(16 * 8) as *i64
79 toks[0]=785; toks[1]=6722; toks[2]=315; toks[3]=9625; toks[4]=374
80 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(NL, model.n_kv_heads, 64, model.head_dim)
81 let logits: *i64 = sys_mmap(model.vocab_size * 8) as *i64
82 var p: nx_int = 0
83 while p < 4 {
84 let one: *i64 = (((toks as i64) + p * 8)) as *i64
85 if nx_f32_llm_forward_v4(model, one, 1, cache, eps, attn_scale, rope_base, 1, logits) != NX_FLV4_OK { return 60 }
86 p = p + 1
87 }
88
89 pw("=== OUR per-layer last-token norm x1000 (compare REF [390,4468,7320,...,279079]) ===\n" as *u8)
90 let bufA: *i64 = sys_mmap(H * 8) as *i64
91 let bufB: *i64 = sys_mmap(H * 8) as *i64
92 var d: nx_int = 0
93 while d < H { bufA[d] = model.embed_weights[374 * H + d]; d = d + 1 }
94 pw(" embed norm*1000="); pn(milli(l2norm(bufA, H))); pw(" (REF 390)\n")
95 dump8(" embed" as *u8, bufA)
96 var cur: *i64 = bufA
97 var use_a: nx_int = 1
98 var L: nx_int = 0
99 while L < NL {
100 let layer: *NxF32LlamaLayerLazy = model.layers[L] as *NxF32LlamaLayerLazy
101 var dst: *i64 = bufB
102 if use_a == 0 { dst = bufA }
103 nx_f32_llama_block_forward_v4(cur, 1, H, model.n_heads, model.n_kv_heads, model.head_dim, model.ffn_dim,
104 layer, cache, L, eps, attn_scale, rope_base, 1, dst)
105 cur = dst
106 use_a = 1 - use_a
107 pw(" L"); pn(L as i64); pw(" norm*1000="); pn(milli(l2norm(cur, H))); pw("\n")
108 if L == 0 { dump8(" afterL0" as *u8, cur) }
109 L = L + 1
110 }
111 return 0
112}