nx_gguf_load_model.nx source
↩ module page · 147 lines · 6212 B
1// nx_gguf_load_model.nx -- top-level model-tensor loader.
2//
3// L4 brick: the missing companion to nx_gguf_load_block.nx. Where
4// block_load reads the 9 per-layer tensors, this brick reads the 3
5// top-level (non-per-layer) tensors a Llama-class GGUF carries:
6//
7// token_embd.weight -- vocab x hidden (embedding table)
8// output.weight -- hidden x vocab (output projection /
9// "lm head"; often
10// tied to token_embd)
11// output_norm.weight -- hidden (final RMSNorm scale)
12//
13// Together with nx_gguf_load_block, this brick covers the FULL set of
14// tensor names a Llama-2 / Llama-3 / Mistral / Qwen GGUF contains
15// (modulo extra K/V biases that don't appear in standard files).
16//
17// After this brick, the path from text prompt -> first emitted token
18// is loader-complete; the next workstream is `nx_llm_run` v2 which
19// composes load_model + load_block + transformer_block_forward +
20// final_norm + output_projection + softmax + sample.
21//
22// Bits-up composition:
23// nx_gguf.nx, nx_gguf_load.nx -- header + per-tensor load
24// nx_tensor.nx -- NxTensor output
25//
26// genealogy_id: llama_gguf_naming_gerganov_2024 + standard_lm_head
27// lineage_id: substrate_gguf_model_loader_v1
28
29// nx_safety_envelope:
30// intended_use: "Load the 3 top-level Llama-class GGUF
31// tensors (token_embd / output / output_norm)
32// into NxTensor + raw-i64 pointer outputs"
33// sil_target: SIL2
34// asil_target: QM
35// dal_target: DAL C
36// evidence: [gguf_naming_spec_public, clean_room,
37// composes_only_shipped_canonical_bricks]
38// hazard_register: [bug-tape-missing-output-weight-tied-embed,
39// bug-tape-vocab-mismatch-vs-tokenizer]
40// residual_risk: "Llama files where output.weight is TIED
41// to token_embd.weight may omit output.weight
42// entirely; caller must handle the NOT_FOUND
43// verdict by reusing token_embd"
44// verdict: NOT_YET_EVALUATED
45
46import "nx_syscalls.nx"
47import "nx_tier.nx"
48import "nx_tensor.nx"
49import "nx_gguf.nx"
50import "nx_gguf_load.nx"
51
52// ===== Sealed-enum: ModelLoadVerdict ==============================
53
54const NX_GML_OK: nx_int = 0
55const NX_GML_ERR_EMBED_NOT_FOUND: nx_int = 1
56const NX_GML_ERR_OUTPUT_NOT_FOUND: nx_int = 2
57const NX_GML_ERR_NORM_NOT_FOUND: nx_int = 3
58const NX_GML_ERR_LOAD_FAIL: nx_int = 4
59const NX_GML_N_VERDICTS: nx_int = 5
60
61func nx_gml_verdict_is_valid(v: nx_int) -> nx_int {
62 if v < 0 { return 0 }
63 if v >= NX_GML_N_VERDICTS { return 0 }
64 return 1
65}
66
67// ===== Bundle struct ==============================================
68//
69// Top-level model weights. Caller pre-allocates this struct;
70// loader populates the 3 fields.
71//
72// output_weight is *NxTensor (full tensor pointer) so the matmul to
73// vocab logits can dispatch normally. output_norm_gamma is *i64
74// (1-D raw vector) so it can plug directly into nx_rmsnorm_forward
75// per its existing signature.
76//
77// token_embd is *NxTensor (used by nx_embedding_lookup which expects
78// a tensor pointer).
79
80struct NxGgufModelWeights {
81 token_embd: *NxTensor, // [vocab_size, hidden_dim]
82 output_weight: *NxTensor, // [hidden_dim, vocab_size] (may be 0 if tied)
83 output_norm_gamma: *i64, // [hidden_dim]
84 is_output_tied: nx_int // 1 if output.weight was missing
85 // and caller should reuse token_embd
86}
87
88const NX_GML_BUNDLE_BYTES: nx_int = 32 // 4 fields * 8
89
90// ===== Public: load top-level model tensors =======================
91
92func nx_gguf_load_model_weights(buf: *u8, hdr: *NxGgufHeader,
93 w_out: *NxGgufModelWeights,
94 out_err: *i64) -> nx_int {
95 // ----- token_embd.weight (17 bytes) -----
96 let n_te: *u8 = sys_mmap(17)
97 n_te[0]=0x74; n_te[1]=0x6f; n_te[2]=0x6b; n_te[3]=0x65
98 n_te[4]=0x6e; n_te[5]=0x5f; n_te[6]=0x65; n_te[7]=0x6d
99 n_te[8]=0x62; n_te[9]=0x64; n_te[10]=0x2e; n_te[11]=0x77
100 n_te[12]=0x65; n_te[13]=0x69; n_te[14]=0x67; n_te[15]=0x68
101 n_te[16]=0x74
102 let t_embd: *NxTensor = nx_gguf_load_tensor(buf, hdr, n_te, 17, out_err)
103 if out_err[0] != NX_GL_OK {
104 out_err[0] = NX_GML_ERR_EMBED_NOT_FOUND
105 return NX_GML_ERR_EMBED_NOT_FOUND
106 }
107 w_out.token_embd = t_embd
108
109 // ----- output_norm.weight (18 bytes) -----
110 let n_on: *u8 = sys_mmap(18)
111 n_on[0]=0x6f; n_on[1]=0x75; n_on[2]=0x74; n_on[3]=0x70
112 n_on[4]=0x75; n_on[5]=0x74; n_on[6]=0x5f; n_on[7]=0x6e
113 n_on[8]=0x6f; n_on[9]=0x72; n_on[10]=0x6d; n_on[11]=0x2e
114 n_on[12]=0x77; n_on[13]=0x65; n_on[14]=0x69; n_on[15]=0x67
115 n_on[16]=0x68; n_on[17]=0x74
116 let t_norm: *NxTensor = nx_gguf_load_tensor(buf, hdr, n_on, 18, out_err)
117 if out_err[0] != NX_GL_OK {
118 out_err[0] = NX_GML_ERR_NORM_NOT_FOUND
119 return NX_GML_ERR_NORM_NOT_FOUND
120 }
121 w_out.output_norm_gamma = t_norm.storage as *i64
122
123 // ----- output.weight (13 bytes) -- OPTIONAL if tied -----
124 //
125 // Llama family: when "output.weight" is missing, the lm head is
126 // tied to token_embd.weight (transposed at use time). We detect
127 // this by trying the load and falling back gracefully if NOT_FOUND.
128 let n_ow: *u8 = sys_mmap(13)
129 n_ow[0]=0x6f; n_ow[1]=0x75; n_ow[2]=0x74; n_ow[3]=0x70
130 n_ow[4]=0x75; n_ow[5]=0x74; n_ow[6]=0x2e; n_ow[7]=0x77
131 n_ow[8]=0x65; n_ow[9]=0x69; n_ow[10]=0x67; n_ow[11]=0x68
132 n_ow[12]=0x74
133 let t_out: *NxTensor = nx_gguf_load_tensor(buf, hdr, n_ow, 13, out_err)
134 if out_err[0] == NX_GL_OK {
135 w_out.output_weight = t_out
136 w_out.is_output_tied = 0
137 } else {
138 // Tied-embedding fallback: leave output_weight null, mark tied.
139 w_out.output_weight = 0 as *NxTensor
140 w_out.is_output_tied = 1
141 // Reset out_err so the overall verdict is OK.
142 out_err[0] = NX_GML_OK
143 }
144
145 out_err[0] = NX_GML_OK
146 return NX_GML_OK
147}