nx_genlora.nx source
↩ module page · 242 lines · 10609 B
1// nx_genlora.nx -- apply LoRA adapters to a base weight. Sovereign, model-agnostic LoRA support.
2//
3// W' = W + (alpha / rank) * multiplier * (up @ down)
4//
5// read out of sd.cpp's lora.hpp, not guessed: `scale_value = alpha / rank; scale_value *= multiplier`.
6// With no `.alpha` tensor, alpha defaults to rank, so the scale is the multiplier alone.
7//
8// ⚠TWO THINGS VARY INDEPENDENTLY, AND BOTH MUST BE HANDLED:
9//
10// 1. THE NAMING CONVENTION.
11// base model.diffusion_model.layers.0.attention.qkv.weight
12// kohya lora_unet_layers_0_attention_qkv.lora_down.weight (flattened, underscores)
13// peft diffusion_model.layers.0.attention.to_q.lora_A.weight (dotted, A=down B=up)
14//
15// 2. WHETHER THE ADAPTER TRAINED THE **FUSED** MODULE OR THE UNFUSED ONES. This is the part that
16// bites. The Z-Image checkpoint stores ONE fused `attention.qkv.weight` [11520, 3840]. The
17// kohya adapter here matches it directly. But all NINE peft adapters in the same folder train
18// `to_q` / `to_k` / `to_v` SEPARATELY -- there is no `attention.qkv` key in any of them.
19// A name-only lookup therefore finds nothing, folds nothing, and returns "this adapter does
20// not cover this layer" for EVERY layer of EVERY peft adapter, with no error anywhere.
21// ★★★★★ AN ADAPTER THAT MATCHES NO KEYS IS INDISTINGUISHABLE FROM ONE THAT IS WORKING WEAKLY.
22// ⇒ a fused base weight must be folded SLICE-WISE from the unfused adapter modules:
23// to_q -> rows [0, out/3) to_k -> [out/3, 2out/3) to_v -> [2out/3, out)
24//
25// ★ MEASURED CAVEAT worth more than the code: on Z-Image layer 0 the kohya adapter moves the
26// activation by only 0.86% while the oracle's own Q8_0 error is 0.6% median. Fitting the scale by
27// least squares against the oracle recovered 0.1725 where alpha/rank is 0.1667 -- consistent, but
28// INSIDE THE NOISE. ★★★★★ YOU CANNOT VALIDATE AN ADAPTER ON A LAYER WHERE THE ADAPTER BARELY ACTS;
29// grade against an EXACT reference (nx_gen_lora_verify does).
30//
31// WHY FOLD RATHER THAN APPLY AT RUNTIME: sd.cpp keeps the adapter separate and pays
32// 2*rank*(in+out)*tokens every step. Folding pays 2*rank*in*out ONCE and leaves the hot matmul
33// byte-identical to the un-adapted path -- so an adapter costs nothing per step and cannot slow
34// down or perturb the kernel.
35// license_tier: ORIGINAL
36
37import "nx_syscalls.nx"
38import "nx_le.nx"
39import "nx_f32.nx"
40import "nx_f32_div.nx"
41import "nx_f32_cvt.nx"
42import "nx_f16.nx"
43import "nx_strconv.nx"
44import "nx_genweights.nx"
45const LORA_MAGIC_1024: i64 = 1024
46
47// verdicts. "absent" and "malformed" are DIFFERENT answers: an adapter that covers only some
48// layers is normal and expected; one whose own tensors disagree on rank is a defect.
49const LORA_ABSENT: i64 = 0
50const LORA_APPLIED: i64 = 1
51const LORA_BAD_SHAPE: i64 = 0 - 10
52const LORA_BAD_PAIR: i64 = 0 - 11
53const LORA_BAD_TYPE: i64 = 0 - 12
54
55func _lo_len(s: *u8) -> i64 {
56 var n: i64 = 0
57 while s[n] != (0 as u8) { n = n + 1 }
58 return n
59}
60func _lo_app(dst: *u8, at: i64, src: *u8) -> i64 {
61 var o: i64 = at
62 var i: i64 = 0
63 while src[i] != (0 as u8) { dst[o] = src[i]; o = o + 1; i = i + 1 }
64 dst[o] = 0
65 return o
66}
67func _lo_ends(s: *u8, suf: *u8) -> i64 {
68 let sl: i64 = _lo_len(s)
69 let fl: i64 = _lo_len(suf)
70 if sl < fl { return 0 }
71 var i: i64 = 0
72 while i < fl { if s[sl - fl + i] != suf[i] { return 0 } i = i + 1 }
73 return 1
74}
75// Replace the trailing `nold` chars of a stem -- turns a FUSED base module name into the UNFUSED
76// module an adapter actually trained.
77func _lo_retail(stem: *u8, nold: i64, rep: *u8) -> i64 {
78 return _lo_app(stem, _lo_len(stem) - nold, rep)
79}
80
81// base "model.diffusion_model.<path>.weight" -> kohya stem "lora_unet_<path with _>"
82func nx_lora_key_kohya(out: *u8, base: *u8) -> i64 {
83 let pre: *u8 = "model.diffusion_model." as *u8
84 let pl: i64 = _lo_len(pre)
85 var i: i64 = 0
86 while i < pl { if base[i] != pre[i] { return 0 - 1 } i = i + 1 }
87 let bl: i64 = _lo_len(base)
88 let sl: i64 = 7 // ".weight"
89 if bl <= pl + sl { return 0 - 2 }
90 var o: i64 = _lo_app(out, 0, "lora_unet_" as *u8)
91 var p: i64 = pl
92 while p < bl - sl {
93 var c: i64 = base[p] as i64
94 if c == 0x2E { c = 0x5F } // '.' -> '_'
95 out[o] = c as u8
96 o = o + 1
97 p = p + 1
98 }
99 out[o] = 0
100 return o
101}
102
103// base "model.diffusion_model.<path>.weight" -> peft stem "diffusion_model.<path>"
104func nx_lora_key_peft(out: *u8, base: *u8) -> i64 {
105 let pre: *u8 = "model." as *u8
106 let pl: i64 = _lo_len(pre)
107 var i: i64 = 0
108 while i < pl { if base[i] != pre[i] { return 0 - 1 } i = i + 1 }
109 let bl: i64 = _lo_len(base)
110 let sl: i64 = 7
111 if bl <= pl + sl { return 0 - 2 }
112 var o: i64 = 0
113 var p: i64 = pl
114 while p < bl - sl { out[o] = base[p]; o = o + 1; p = p + 1 }
115 out[o] = 0
116 return o
117}
118
119func _lo_probe(lw: *i64, stem: *u8, suffix: *u8, nm: *u8, inf: *i64) -> i64 {
120 let o: i64 = _lo_app(nm, _lo_app(nm, 0, stem), suffix)
121 let idx: i64 = nx_gw_find(lw, nm, o)
122 if idx < 0 { return 0 }
123 inf[0] = idx
124 return 1
125}
126
127// Fold ONE adapter pair into rows [row_off, row_off + n_rows) of packed-f32 W [out_dim, in_dim].
128// Returns LORA_APPLIED / LORA_ABSENT / negative.
129func nx_lora_fold_slice(lw: *i64, stem: *u8, dn_suf: *u8, up_suf: *u8, W: *u8,
130 in_dim: i64, row_off: i64, n_rows: i64, mult_f32: i64) -> i64 {
131 let nm: *u8 = sys_mmap(LORA_MAGIC_1024)
132 let inf: *i64 = sys_mmap(64) as *i64
133
134 if _lo_probe(lw, stem, dn_suf, nm, inf) == 0 { return LORA_ABSENT }
135
136 // down [rank, in] in torch -> dim0=in dim1=rank after the safetensors [out,in] reversal
137 let rank: i64 = nx_gw_dim1(lw, inf[0])
138 if nx_gw_dim0(lw, inf[0]) != in_dim { return LORA_BAD_SHAPE }
139 if rank <= 0 { return LORA_BAD_SHAPE }
140 let down: *u8 = sys_mmap(rank * in_dim * 4 + 64)
141 if nx_gw_to_f32_packed(lw, inf[0], down, rank * in_dim) != 0 { return LORA_BAD_TYPE }
142
143 if _lo_probe(lw, stem, up_suf, nm, inf) == 0 { return LORA_BAD_PAIR }
144 if nx_gw_dim1(lw, inf[0]) != n_rows { return LORA_BAD_SHAPE }
145 if nx_gw_dim0(lw, inf[0]) != rank { return LORA_BAD_SHAPE }
146 let up: *u8 = sys_mmap(n_rows * rank * 4 + 64)
147 if nx_gw_to_f32_packed(lw, inf[0], up, n_rows * rank) != 0 { return LORA_BAD_TYPE }
148
149 var scale: i64 = mult_f32
150 if _lo_probe(lw, stem, ".alpha" as *u8, nm, inf) == 1 {
151 let ab: *u8 = sys_mmap(64)
152 if nx_gw_to_f32_packed(lw, inf[0], ab, 1) == 0 {
153 scale = __f32_mul(nx_f32_div(nx_le_read_u32(ab, 0), nx_i32_to_f32(rank)), mult_f32)
154 }
155 }
156
157 // W[row_off+o][:] += (scale * up[o][r]) * down[r][:] -- an AXPY, not a dot.
158 // The dot form (sum over r) strides down[] by in_dim and vectorises badly; the AXPY form walks
159 // W and down contiguously and lets __f32x8_fma accumulate straight into the weight in place.
160 let cv: *u8 = sys_mmap(64)
161 let nch: i64 = in_dim / 8
162 var o: i64 = 0
163 while o < n_rows {
164 let wrow: i64 = (W as i64) + (row_off + o) * in_dim * 4
165 var r: i64 = 0
166 while r < rank {
167 let c: i64 = __f32_mul(scale, nx_le_read_u32(up, (o * rank + r) * 4))
168 if (c & 0x7FFFFFFF) != 0 { // a zero row contributes nothing, exactly
169 var k: i64 = 0
170 while k < 8 { nx_le_write_u32(cv, k * 4, c); k = k + 1 }
171 let drow: i64 = (down as i64) + r * in_dim * 4
172 var ch: i64 = 0
173 while ch < nch {
174 __f32x8_fma((wrow + ch * 32) as *u8, cv, (drow + ch * 32) as *u8)
175 ch = ch + 1
176 }
177 var i: i64 = nch * 8
178 while i < in_dim {
179 let f: i64 = (row_off + o) * in_dim + i
180 nx_le_write_u32(W, f * 4, __f32_add(nx_le_read_u32(W, f * 4),
181 __f32_mul(c, nx_le_read_u32(down, r * in_dim * 4 + i * 4))))
182 i = i + 1
183 }
184 }
185 r = r + 1
186 }
187 o = o + 1
188 }
189 return LORA_APPLIED
190}
191
192// Fold whatever this adapter has for `base` into W. Tries, in order:
193// kohya fused -> peft fused -> peft UNFUSED q/k/v thirds -> peft renamed out projection.
194// Returns LORA_APPLIED if ANY of them applied, LORA_ABSENT if none did, negative if one was
195// present but malformed.
196func nx_lora_fold(lw: *i64, base: *u8, W: *u8, in_dim: i64, out_dim: i64, mult_f32: i64) -> i64 {
197 let stem: *u8 = sys_mmap(LORA_MAGIC_1024)
198 let dn_k: *u8 = ".lora_down.weight" as *u8
199 let up_k: *u8 = ".lora_up.weight" as *u8
200 let dn_p: *u8 = ".lora_A.weight" as *u8
201 let up_p: *u8 = ".lora_B.weight" as *u8
202
203 if nx_lora_key_kohya(stem, base) > 0 {
204 let rc: i64 = nx_lora_fold_slice(lw, stem, dn_k, up_k, W, in_dim, 0, out_dim, mult_f32)
205 if rc != LORA_ABSENT { return rc }
206 }
207 if nx_lora_key_peft(stem, base) <= 0 { return LORA_ABSENT }
208
209 let rc2: i64 = nx_lora_fold_slice(lw, stem, dn_p, up_p, W, in_dim, 0, out_dim, mult_f32)
210 if rc2 != LORA_ABSENT { return rc2 }
211
212 // ---- fused qkv <- three unfused projections ----
213 if _lo_ends(base, "attention.qkv.weight" as *u8) == 1 {
214 let third: i64 = out_dim / 3
215 if third * 3 != out_dim { return LORA_BAD_SHAPE }
216 var applied: i64 = 0
217 var i: i64 = 0
218 while i < 3 {
219 nx_lora_key_peft(stem, base)
220 if i == 0 { _lo_retail(stem, 3, "to_q" as *u8) }
221 if i == 1 { _lo_retail(stem, 3, "to_k" as *u8) }
222 if i == 2 { _lo_retail(stem, 3, "to_v" as *u8) }
223 let r: i64 = nx_lora_fold_slice(lw, stem, dn_p, up_p, W, in_dim,
224 i * third, third, mult_f32)
225 if r < 0 { return r }
226 if r == LORA_APPLIED { applied = applied + 1 }
227 i = i + 1
228 }
229 // A PARTIAL q/k/v SET IS A DEFECT, NOT A PARTIAL SUCCESS: folding q and v but not k
230 // produces a subtly wrong attention that still renders a picture.
231 if applied == 3 { return LORA_APPLIED }
232 if applied == 0 { return LORA_ABSENT }
233 return LORA_BAD_PAIR
234 }
235 // ---- peft renames the attention output projection ----
236 if _lo_ends(base, "attention.out.weight" as *u8) == 1 {
237 nx_lora_key_peft(stem, base)
238 _lo_retail(stem, 3, "to_out.0" as *u8)
239 return nx_lora_fold_slice(lw, stem, dn_p, up_p, W, in_dim, 0, out_dim, mult_f32)
240 }
241 return LORA_ABSENT
242}