nx_f32_llm_load.nx source
↩ module page · 129 lines · 5121 B
1// nx_f32_llm_load.nx -- model-level GGUF binder.
2//
3// Populates an NxF32LlamaModel from a parsed GGUF:
4// - 3 top-level tensors:
5// token_embd.weight -> embed_weights
6// output_norm.weight -> gamma_out
7// output.weight -> lm_head (tied to embed if missing)
8// - N per-layer NxF32LlamaLayer structs via nx_f32_llama_layer_load_from_gguf
9//
10// Caller responsibility: pre-fill model dimension fields (n_layers,
11// hidden_dim, n_heads, n_kv_heads, head_dim, ffn_dim, vocab_size).
12// v1 does not parse GGUF metadata to derive them -- a separate
13// metadata-walker brick is queued for that. This keeps the binder
14// focused on tensor binding only.
15//
16// Tied-embed detection: if "output.weight" is not found, lm_head
17// points to the same storage as embed_weights (the canonical Llama
18// memory-saving trick used by Qwen, Llama-3, etc.). Caller must
19// not free either pointer independently.
20//
21// genealogy_id: llama_gguf_naming_gerganov_2024 + standard_model_binding
22// + tied_embed_qwen_llama3
23// lineage_id: substrate_f32_llm_load_v1
24
25import "nx_syscalls.nx"
26import "nx_tier.nx"
27import "nx_gguf.nx"
28import "nx_gguf_load_f32.nx"
29import "nx_f32_llama_block.nx"
30import "nx_f32_llama_layer_load.nx"
31import "nx_f32_llm.nx"
32
33const NX_FLM_OK: nx_int = 0
34const NX_FLM_ERR_NULL: nx_int = 1
35const NX_FLM_ERR_BAD_DIM: nx_int = 2
36const NX_FLM_ERR_NOT_FOUND: nx_int = 3
37const NX_FLM_ERR_LAYER: nx_int = 4
38const NX_FLM_N_VERDICTS: nx_int = 5
39
40func nx_flm_verdict_is_valid(v: nx_int) -> nx_int {
41 if v < 0 { return 0 }
42 if v >= NX_FLM_N_VERDICTS { return 0 }
43 return 1
44}
45
46// Load weights into a model struct whose dim fields are already set.
47//
48// Returns NX_FLM_OK on success. If "output.weight" is missing,
49// lm_head is tied to embed_weights (assumed transposed-compatible
50// per Llama convention -- caller can detect via pointer equality).
51
52func nx_f32_llm_load_weights_from_gguf(buf: *u8, hdr: *NxGgufHeader,
53 model: *NxF32LlamaModel,
54 out_err: *i64) -> nx_int {
55 if model == (0 as *NxF32LlamaModel) {
56 out_err[0] = NX_FLM_ERR_NULL
57 return NX_FLM_ERR_NULL
58 }
59 if model.n_layers <= 0 {
60 out_err[0] = NX_FLM_ERR_BAD_DIM
61 return NX_FLM_ERR_BAD_DIM
62 }
63
64 let nv_out: *i64 = sys_mmap(8) as *i64
65 let inner_err: *i64 = sys_mmap(8) as *i64
66
67 // ===== token_embd.weight =====
68 let n_te: *u8 = sys_mmap(17)
69 n_te[0]=0x74 as u8; n_te[1]=0x6f as u8; n_te[2]=0x6b as u8; n_te[3]=0x65 as u8
70 n_te[4]=0x6e as u8; n_te[5]=0x5f as u8; n_te[6]=0x65 as u8; n_te[7]=0x6d as u8
71 n_te[8]=0x62 as u8; n_te[9]=0x64 as u8; n_te[10]=0x2e as u8; n_te[11]=0x77 as u8
72 n_te[12]=0x65 as u8; n_te[13]=0x69 as u8; n_te[14]=0x67 as u8; n_te[15]=0x68 as u8
73 n_te[16]=0x74 as u8
74 let embed: *i64 = nx_gguf_load_tensor_to_f32(buf, hdr, n_te, 17,
75 nv_out, inner_err)
76 if inner_err[0] != NX_GLF_OK {
77 out_err[0] = NX_FLM_ERR_NOT_FOUND
78 return NX_FLM_ERR_NOT_FOUND
79 }
80 model.embed_weights = embed
81
82 // ===== output_norm.weight =====
83 let n_on: *u8 = sys_mmap(18)
84 n_on[0]=0x6f as u8; n_on[1]=0x75 as u8; n_on[2]=0x74 as u8; n_on[3]=0x70 as u8
85 n_on[4]=0x75 as u8; n_on[5]=0x74 as u8; n_on[6]=0x5f as u8; n_on[7]=0x6e as u8
86 n_on[8]=0x6f as u8; n_on[9]=0x72 as u8; n_on[10]=0x6d as u8; n_on[11]=0x2e as u8
87 n_on[12]=0x77 as u8; n_on[13]=0x65 as u8; n_on[14]=0x69 as u8; n_on[15]=0x67 as u8
88 n_on[16]=0x68 as u8; n_on[17]=0x74 as u8
89 let gout: *i64 = nx_gguf_load_tensor_to_f32(buf, hdr, n_on, 18,
90 nv_out, inner_err)
91 if inner_err[0] != NX_GLF_OK {
92 out_err[0] = NX_FLM_ERR_NOT_FOUND
93 return NX_FLM_ERR_NOT_FOUND
94 }
95 model.gamma_out = gout
96
97 // ===== output.weight (tied-embed fallback) =====
98 let n_ow: *u8 = sys_mmap(13)
99 n_ow[0]=0x6f as u8; n_ow[1]=0x75 as u8; n_ow[2]=0x74 as u8; n_ow[3]=0x70 as u8
100 n_ow[4]=0x75 as u8; n_ow[5]=0x74 as u8; n_ow[6]=0x2e as u8; n_ow[7]=0x77 as u8
101 n_ow[8]=0x65 as u8; n_ow[9]=0x69 as u8; n_ow[10]=0x67 as u8; n_ow[11]=0x68 as u8
102 n_ow[12]=0x74 as u8
103 let lm: *i64 = nx_gguf_load_tensor_to_f32(buf, hdr, n_ow, 13,
104 nv_out, inner_err)
105 if inner_err[0] == NX_GLF_OK {
106 model.lm_head = lm
107 } else {
108 // Tied embed: lm_head shares storage with embed_weights.
109 model.lm_head = embed
110 }
111
112 // ===== Per-layer binding =====
113 model.layers = sys_mmap(model.n_layers * 8) as *i64
114 var L: nx_int = 0
115 while L < model.n_layers {
116 let layer: *NxF32LlamaLayer = nx_f32_llama_layer_alloc()
117 let v: nx_int = nx_f32_llama_layer_load_from_gguf(
118 buf, hdr, L, layer, inner_err)
119 if v != NX_FLL_OK {
120 out_err[0] = NX_FLM_ERR_LAYER
121 return NX_FLM_ERR_LAYER
122 }
123 model.layers[L] = layer as i64
124 L = L + 1
125 }
126
127 out_err[0] = NX_FLM_OK
128 return NX_FLM_OK
129}