LoRA.xtl

1⍝│ LoRA: low-rank adapters: the adapted forward pass in factored form, the merge, the counts, the gradients and steps, LoRA-FA
2⍝│
3⍝│ Import it with an alias of your choice: "lo:" u_se< "LoRA".
4⍝│ Put libs/LoRA/src on XETAL_PATH ("just path"); the reference is libs/LoRA/docs.
5⍝│ Names with l: are exported; h: names are private to this file.
6⍝│
7⍝│ A frozen weight matrix W (m x n: m outputs, n inputs) is adapted by
8⍝│ a low-rank update s B A: A is r x n, B is m x r, the rank r much
9⍝│ smaller than m and n, s the scale alpha / r. An adapter is the tuple
10⍝│ (W, A, B, s). The inputs are columns: x is one input of n numbers or
11⍝│ an n x N matrix of N inputs, and W x is W '+ '× i̲nner x. The
12⍝│ forward pass keeps the factors apart, B (A x), as training does:
13⍝│ the merged matrix W + s B A is made only to deploy it.
14⍝│
15⍝│ Training moves A and B by the gradients of the squared error
16⍝│ (g_rads, s_gd, a_dam through the Optim library); LoRA-FA freezes A
17⍝│ (regenerated from a seed, s_eeded) and trains B alone.
18
19"op:" u̲se< "Optim"
20
21⍝━━ The forward pass
22
23⍝│ alpha s̲caling r: the adapter's scale s, alpha / r (so that a
24⍝│ change of rank keeps the update's size).
25⍝│ >> "lo:" u_se< "LoRA"
26⍝│ >> 16.0 lo:s_caling 8
27⍝│ 2.0
28ˡs̲caling ← { alpha r → alpha ÷ f̲loat r }
29
30⍝│ alpha s̲calingRs r: rsLoRA's rank-stabilized scale, alpha / sqrt r
31⍝│ (Kalajdzievski, "A Rank Stabilization Scaling Factor for Fine-Tuning
32⍝│ with LoRA", 2023): with alpha / r the update's learning slows as r
33⍝│ grows; over sqrt r it does not. An adapter takes either, (W, A, B, s).
34⍝│ >> "lo:" u_se< "LoRA"
35⍝│ >> 16.0 lo:s_calingRs 4
36⍝│ 8.0
37ˡs̲calingRs ← { alpha r → alpha ÷ s̲qrt f̲loat r }
38
39⍝│ (W, A, B, s) f̲orward x: the adapted layer, W x + s B (A x), the
40⍝│ update in its factored form: A x (r numbers per input), then B of
41⍝│ that, never the m x n matrix B A.
42⍝│ >> "lo:" u_se< "LoRA"
43⍝│ >> (2 2 r_eshape 1.0 0.0 0.0 1.0, 1 2 r_eshape 1.0 1.0, 2 1 r_eshape 1.0 -1.0, 0.5) lo:f_orward 2.0 4.0
44⍝│ 5.0 1.0
45ˡf̲orward ← { (W, A, B, s) x → (W '+ '× i̲nner x) + s × B '+ '× i̲nner A '+ '× i̲nner x }
46
47⍝━━ The merge
48
49⍝│ d̲elta (A, B, s): the update itself, s B A, an m x n matrix of
50⍝│ rank r at most.
51⍝│ >> "lo:" u_se< "LoRA"
52⍝│ >> lo:d_elta (1 2 r_eshape 1.0 1.0, 2 1 r_eshape 1.0 -1.0, 0.5)
53⍝│  0.5  0.5
54⍝│ -0.5 -0.5
55ˡd̲elta ← { (A, B, s) → s × B '+ '× i̲nner A }
56
57⍝│ m̲erge (W, A, B, s): the adapter folded into the weights, W + s B
58⍝│ A: one matrix again, as fast as the base layer, for one task.
59⍝│ >> "lo:" u_se< "LoRA"
60⍝│ >> lo:m_erge (2 2 r_eshape 1.0 0.0 0.0 1.0, 1 2 r_eshape 1.0 1.0, 2 1 r_eshape 1.0 -1.0, 0.5)
61⍝│  1.5 0.5
62⍝│ -0.5 0.5
63ˡm̲erge ← { (W, A, B, s) → W + ˡd̲elta (A, B, s) }
64
65⍝━━ Gradients
66
67⍝│ (W, A, B, s) l̲oss (x, y): the squared error of the adapted layer
68⍝│ on the inputs x (columns) against the wanted outputs y, half the sum
69⍝│ of the squared differences, so that its gradient at the output, g,
70⍝│ is the difference itself.
71⍝│ >> "lo:" u_se< "LoRA"
72⍝│ >> (2 2 r_eshape 1.0 0.0 0.0 1.0, 1 2 r_eshape 1.0 1.0, 2 1 r_eshape 1.0 -1.0, 0.5) lo:l_oss (2 1 r_eshape 2.0 4.0, 2 1 r_eshape 4.0 2.0)
73⍝│ 1.0
74ˡl̲oss ← { a (x, y) →
75  e ← (a ˡf̲orward x) − y
76  0.5 × '+ r̲/ r̲avel e × e
77}
78
79⍝│ (W, A, B, s) g̲rads (x, g): the gradients for the inputs x (an n x
80⍝│ N matrix, one input per column) and g, the loss's gradient at the
81⍝│ output (m x N): with z = A x, dL/dA = s (B' g) x', dL/dB = s g z' and
82⍝│ dL/dx = W' g + s A' (B' g), as the tuple (dA, dB, dx). B' g is made
83⍝│ once and used twice; no m x n matrix but W's own is made.
84⍝│ >> "lo:" u_se< "LoRA"
85⍝│ >> (dA, dB, dx) := (2 2 r_eshape 1.0 0.0 0.0 1.0, 1 2 r_eshape 1.0 1.0, 2 1 r_eshape 1.0 -1.0, 0.5) lo:g_rads (2 1 r_eshape 2.0 4.0, 2 1 r_eshape 1.0 -1.0)
86⍝│ >> dB
87⍝│  3.0
88⍝│ -3.0
89ˡg̲rads ← { (W, A, B, s) (x, g) →
90  b ← (o̲\ B) '+ '× i̲nner g
91  (s × b '+ '× i̲nner o̲\ x, s × g '+ '× i̲nner o̲\ A '+ '× i̲nner x, ((o̲\ W) '+ '× i̲nner g) + s × (o̲\ A) '+ '× i̲nner b)
92}
93
94⍝━━ Training steps
95
96⍝│ lr s̲gd ((W, A, B, s), (x, y)): one step of gradient descent on
97⍝│ the squared error (l_oss) at the rate lr, A and B moved by Optim's
98⍝│ s_gd, W and s left as they are: the adapter after the step. p_ower
99⍝│ repeats it.
100⍝│ >> "lo:" u_se< "LoRA"
101⍝│ >> (_, A, B, _) := 0.1 lo:s_gd ((2 2 r_eshape 1.0 0.0 0.0 1.0, 1 2 r_eshape 1.0 1.0, 2 1 r_eshape 1.0 -1.0, 0.5), (2 1 r_eshape 2.0 4.0, 2 1 r_eshape 4.0 2.0))
102⍝│ >> B
103⍝│  0.7
104⍝│ -0.7
105ˡs̲gd ← { lr ((W, A, B, s), (x, y)) →
106  (dA, dB, _) ← (W, A, B, s) ˡg̲rads (x, ((W, A, B, s) ˡf̲orward x) − y)
107  (W, lr ᵒᵖs̲gd (A, dA), lr ᵒᵖs̲gd (B, dB), s)
108}
109
110⍝│ a̲damStart (W, A, B, s): what Adam starts from: the adapter, its
111⍝│ four running averages at 0 (A's two, then B's) and the step count 1.
112ˡa̲damStart ← { (W, A, B, s) → ((W, A, B, s), (0.0 × A, 0.0 × A, 0.0 × B, 0.0 × B), 1) }
113
114⍝│ lr a̲dam (((W, A, B, s), (mA, vA, mB, vB), t), (x, y)): Adam's step
115⍝│ t on the squared error, A and B moved by Optim's a_dam at the rate
116⍝│ lr; the next state, the step count one more.
117⍝│ >> "lo:" u_se< "LoRA"
118⍝│ >> ((_, _, B, _), _, _) := 0.1 lo:a_dam ((lo:a_damStart (2 2 r_eshape 1.0 0.0 0.0 1.0, 1 2 r_eshape 1.0 1.0, 2 1 r_eshape 1.0 -1.0, 0.5)), (2 1 r_eshape 2.0 4.0, 2 1 r_eshape 4.0 2.0))
119⍝│ >> B
120⍝│  0.9000000003333333
121⍝│ -0.9000000003333333
122ˡa̲dam ← { lr (((W, A, B, s), (mA, vA, mB, vB), t), (x, y)) →
123  (dA, dB, _) ← (W, A, B, s) ˡg̲rads (x, ((W, A, B, s) ˡf̲orward x) − y)
124  (A2, mA2, vA2) ← (lr, t) ᵒᵖa̲dam (A, mA, vA, dA)
125  (B2, mB2, vB2) ← (lr, t) ᵒᵖa̲dam (B, mB, vB, dB)
126  ((W, A2, B2, s), (mA2, vA2, mB2, vB2), t + 1)
127}
128
129⍝━━ LoRA-FA: A frozen
130
131⍝│ seed s̲eeded r n: an r x n matrix regenerated from a whole-number
132⍝│ seed (LoRA-FA's frozen A, which a checkpoint then need not hold):
133⍝│ sin(seed k k) for k = 1 .. r n, over the square root of n (the
134⍝│ inputs' count, so A x keeps x's size). No chance: the same seed
135⍝│ gives the same matrix everywhere.
136⍝│ >> "lo:" u_se< "LoRA"
137⍝│ >> 7 lo:s_eeded 2 4
138⍝│ 0.32849329935939453 0.13545289415393452  0.08367785015140346 -0.44499780218341667
139⍝│ -0.4005672975890204  0.3115061055018265 -0.26829917759428185  0.47413628607807934
140ˡs̲eeded ← { seed rn →
141  k ← f̲loat r̲ange '× r̲/ rn
142  rn r̲eshape (s̲in (f̲loat seed) × k × k) ÷ s̲qrt f̲loat 2 s̲elect rn
143}
144
145⍝│ (W, A, B, s) g̲radsFA (x, g): LoRA-FA's gradients, (dB, dx): A is
146⍝│ frozen, so its gradient is skipped, but x's still flows back through
147⍝│ A, dx = W' g + s A' (B' g).
148ˡg̲radsFA ← { (W, A, B, s) (x, g) →
149  (s × g '+ '× i̲nner o̲\ A '+ '× i̲nner x, ((o̲\ W) '+ '× i̲nner g) + s × (o̲\ A) '+ '× i̲nner (o̲\ B) '+ '× i̲nner g)
150}
151
152⍝│ lr s̲gdFA ((W, A, B, s), (x, y)): LoRA-FA's step of gradient
153⍝│ descent: B alone moves. Every update s B A it reaches has its rows
154⍝│ in A's row space.
155ˡs̲gdFA ← { lr ((W, A, B, s), (x, y)) →
156  (dB, _) ← (W, A, B, s) ˡg̲radsFA (x, ((W, A, B, s) ˡf̲orward x) − y)
157  (W, A, lr ᵒᵖs̲gd (B, dB), s)
158}
159
160⍝│ a̲damStartFA (W, A, B, s): LoRA-FA's Adam state: the adapter, B's
161⍝│ two running averages at 0 (A has none: half of LoRA's optimizer
162⍝│ memory on a square layer), the step 1.
163ˡa̲damStartFA ← { (W, A, B, s) → ((W, A, B, s), (0.0 × B, 0.0 × B), 1) }
164
165⍝│ lr a̲damFA (((W, A, B, s), (mB, vB), t), (x, y)): LoRA-FA's Adam
166⍝│ step: B alone moves; the next state.
167ˡa̲damFA ← { lr (((W, A, B, s), (mB, vB), t), (x, y)) →
168  (dB, _) ← (W, A, B, s) ˡg̲radsFA (x, ((W, A, B, s) ˡf̲orward x) − y)
169  (B2, mB2, vB2) ← (lr, t) ᵒᵖa̲dam (B, mB, vB, dB)
170  ((W, A, B2, s), (mB2, vB2), t + 1)
171}
172
173⍝━━ Rank
174
175⍝│ One step of Gram-Schmidt on the rows: the k-th row's remainder R[k],
176⍝│ if larger than tol, becomes a direction and is taken out of every
177⍝│ row, and the count c goes up.
178ʰs̲tep ← { tol (R, k, c) →
179  v ← k s̲elect R
180  n ← s̲qrt '+ r̲/ v × v
181  n ≤ tol ? (R, k + 1, c)
182  (R − (R '+ '× i̲nner v ÷ n) '× t̲able v ÷ n, k + 1, c + 1)
183}
184
185⍝│ r̲ank M: the number of independent rows of the matrix M (its
186⍝│ rank), by Gram-Schmidt: a row counts when what is left of it, after
187⍝│ the directions of the rows before it are taken out, is larger than
188⍝│ 1e-9 times the largest row.
189⍝│ >> "lo:" u_se< "LoRA"
190⍝│ >> lo:r_ank 3 3 r_eshape 1.0 2.0 3.0 2.0 4.0 6.0 0.0 1.0 1.0
191⍝│ 2
192ˡr̲ank ← { M →
193  tol ← 0.000000001 × s̲qrt 'm̲ax r̲/ '+ r̲/₂ M × M
194  (_, _, c) ← (f̲irst s̲hape M) '{ R → tol ʰs̲tep R } p̲ower (M, 1, 0)
195  c
196}
197
198⍝━━ Counts
199
200⍝│ r p̲arams m n: the numbers an adapter of rank r on an m x n matrix
201⍝│ trains, r (m + n) (B's m r and A's r n), against the m n of the
202⍝│ matrix itself.
203⍝│ >> "lo:" u_se< "LoRA"
204⍝│ >> 16 lo:p_arams 4096 14336
205⍝│ 294912
206ˡp̲arams ← { r mn → r × '+ r̲/ mn }
207
208⍝│ r e̲xpert d f: the numbers rank-r adapters train on one SwiGLU
209⍝│ expert of hidden size d and inner size f: its three matrices (gate
210⍝│ and up, f x d; down, d x f), 3 r (d + f).
211⍝│ >> "lo:" u_se< "LoRA"
212⍝│ >> 16 lo:e_xpert 4096 14336
213⍝│ 884736
214ˡe̲xpert ← { r df → 3 × r ˡp̲arams df }
215
216⍝│ r e̲xpertFA d f: the same expert under LoRA-FA: the numbers it
217⍝│ trains (the B's, r (2f + d)), then the frozen A's it regenerates
218⍝│ (r (2d + f)); together r e_xpert d f.
219⍝│ >> "lo:" u_se< "LoRA"
220⍝│ >> 16 lo:e_xpertFA 4096 14336
221⍝│ 524288 360448
222ˡe̲xpertFA ← { r df →
223  (d, f) ← (f̲irst df, 2 s̲elect df)
224  (r × f + f + d) c̲at r × d + d + f
225}
226
227⍝│ b b̲ytes n: the bytes n numbers take at b bits each (16 for BF16
228⍝│ and FP16, 32 for FP32, 4 for NF4).
229⍝│ >> "lo:" u_se< "LoRA"
230⍝│ >> 16 lo:b_ytes 884736
231⍝│ 1769472.0
232ˡb̲ytes ← { b n → (f̲loat b × n) ÷ 8.0 }
233
234⍝│ r c̲osts m n: the multiply-adds for one input through rank-r
235⍝│ factors on an m x n matrix, B (A x) as r (m + n), then what making
236⍝│ the matrix B A first costs, r m n.
237⍝│ >> "lo:" u_se< "LoRA"
238⍝│ >> 8 lo:c_osts 64 64
239⍝│ 1024 32768
240ˡc̲osts ← { r mn → (r ˡp̲arams mn) c̲at r × '× r̲/ mn }