include/brotensor/ops/lora.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | // brotensor ops/lora.h — Low-rank (LoRA) adapter delta on a frozen base | ||
| 4 | // linear, optionally gated on the bottleneck. | ||
| 5 | // | ||
| 6 | // Header-only: pure orchestration of the device-neutral linear + elementwise | ||
| 7 | // free functions, so it runs on whichever device the input tensors live on. | ||
| 8 | // There is no new vtable op and no backend code — a LoRA "layer" is just a | ||
| 9 | // composition of linear_forward / linear_backward with an elementwise gate. | ||
| 10 | // | ||
| 11 | // For a frozen base weight W:(out,in) and bias b:(out,1), input x:(in,1), | ||
| 12 | // down-projection A:(r,in), up-projection B:(out,r), scalar `scale`, and an | ||
| 13 | // optional per-rank gate g:(r,1): | ||
| 14 | // | ||
| 15 | // h = A x (bottleneck, r-dim) | ||
| 16 | // hg = g (.) h (h unchanged if g is null) | ||
| 17 | // y = W x + b + scale * (B hg) | ||
| 18 | // | ||
| 19 | // The gate is how conditioning rides on top of a single shared adapter: with | ||
| 20 | // g = g(condition) and g(0) = 0 the delta vanishes and y reproduces the base | ||
| 21 | // model exactly. The base weight/bias never receive a gradient (frozen); the | ||
| 22 | // trainable parameters are A, B and whatever produces the gate. | ||
| 23 | |||
| 24 | #include "../tensor.h" | ||
| 25 | #include "linear.h" | ||
| 26 | #include "elementwise.h" | ||
| 27 | |||
| 28 | namespace brotensor { | ||
| 29 | |||
| 30 | // LoRA forward. `y` is resized to (out,1). `h_out` / `hg_out` are the | ||
| 31 | // bottleneck activations cached for lora_backward — pass them straight back | ||
| 32 | // in. `g` may be null (plain ungated LoRA), else shape (r,1). | ||
| 33 | 270 | inline void lora_forward(const Tensor& W, const Tensor& b, const Tensor& x, | |
| 34 | const Tensor& A, const Tensor& B, float scale, | ||
| 35 | const Tensor* g, | ||
| 36 | Tensor& y, Tensor& h_out, Tensor& hg_out) { | ||
| 37 | 270 | const Device dev = x.device; | |
| 38 | // base: y = W x + b | ||
| 39 | 270 | linear_forward(W, b, x, y); | |
| 40 | // bottleneck h = A x (no bias) | ||
| 41 | 270 | Tensor zr = Tensor::zeros_on(dev, A.rows, 1); | |
| 42 |
1/2✓ Branch 0 taken 270 times.
✗ Branch 1 not taken.
|
270 | linear_forward(A, zr, x, h_out); |
| 43 | // gate: hg = g (.) h | ||
| 44 |
1/2✓ Branch 0 taken 270 times.
✗ Branch 1 not taken.
|
270 | hg_out = h_out.clone(); |
| 45 |
3/4✓ Branch 0 taken 139 times.
✓ Branch 1 taken 131 times.
✓ Branch 2 taken 139 times.
✗ Branch 3 not taken.
|
270 | if (g) mul_inplace(hg_out, *g); |
| 46 | // delta = scale * (B hg); y += delta | ||
| 47 |
1/2✓ Branch 0 taken 270 times.
✗ Branch 1 not taken.
|
270 | Tensor zo = Tensor::zeros_on(dev, B.rows, 1); |
| 48 | 270 | Tensor delta; | |
| 49 |
1/2✓ Branch 0 taken 270 times.
✗ Branch 1 not taken.
|
270 | linear_forward(B, zo, hg_out, delta); |
| 50 |
1/2✓ Branch 0 taken 270 times.
✗ Branch 1 not taken.
|
270 | scale_inplace(delta, scale); |
| 51 |
1/2✓ Branch 0 taken 270 times.
✗ Branch 1 not taken.
|
270 | add_inplace(y, delta); |
| 52 | 270 | } | |
| 53 | |||
| 54 | // LoRA backward. `dY`:(out,1) upstream. Base W,b are frozen (no grad). | ||
| 55 | // dA:(r,in), dB:(out,r) are ACCUMULATED — caller zeros before the first call. | ||
| 56 | // dG:(r,1) accumulated if non-null AND `g` was non-null in the forward. | ||
| 57 | // dX:(in,1) overwritten if non-null (base + lora input grad). Pass null when | ||
| 58 | // the input is frozen (e.g. the AdaIN style vector) to skip the base | ||
| 59 | // backward entirely. | ||
| 60 | // `h` / `hg` are the cached bottleneck activations from lora_forward. | ||
| 61 | 6 | inline void lora_backward(const Tensor& W, const Tensor& x, | |
| 62 | const Tensor& A, const Tensor& B, float scale, | ||
| 63 | const Tensor* g, | ||
| 64 | const Tensor& h, const Tensor& hg, | ||
| 65 | const Tensor& dY, | ||
| 66 | Tensor& dA, Tensor& dB, | ||
| 67 | Tensor* dG, Tensor* dX) { | ||
| 68 | 6 | const Device dev = x.device; | |
| 69 | 6 | const int out = B.rows, r = A.rows, in = A.cols; | |
| 70 | // dp = scale * dY (grad into p = B hg, since y = base + scale*p) | ||
| 71 | 6 | Tensor dp = dY.clone(); | |
| 72 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | scale_inplace(dp, scale); |
| 73 | // through B: dB += dp hg^T ; dHg = B^T dp | ||
| 74 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | Tensor dHg, dBbias = Tensor::zeros_on(dev, out, 1); |
| 75 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | linear_backward(B, hg, dp, dHg, dB, dBbias); |
| 76 | // gate split: dH = dHg (.) g ; dG += dHg (.) h (hg = g (.) h) | ||
| 77 | 6 | Tensor dH; | |
| 78 |
2/2✓ Branch 0 taken 3 times.
✓ Branch 1 taken 3 times.
|
6 | if (g) { |
| 79 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | dH = dHg.clone(); |
| 80 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | mul_inplace(dH, *g); |
| 81 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | if (dG) { |
| 82 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | Tensor t = dHg.clone(); |
| 83 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | mul_inplace(t, h); |
| 84 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | add_inplace(*dG, t); |
| 85 | 3 | } | |
| 86 | 3 | } else { | |
| 87 |
1/2✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
|
3 | dH = dHg; |
| 88 | } | ||
| 89 | // through A: dA += dH x^T ; dX_lora = A^T dH | ||
| 90 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | Tensor dXl, dAbias = Tensor::zeros_on(dev, r, 1); |
| 91 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | linear_backward(A, x, dH, dXl, dA, dAbias); |
| 92 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (dX) { |
| 93 | // base contributes W^T dY to dX; W is frozen so dW/dB are discarded. | ||
| 94 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | Tensor dWbase = Tensor::zeros_on(dev, out, in); |
| 95 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | Tensor dBbase = Tensor::zeros_on(dev, out, 1); |
| 96 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | linear_backward(W, x, dY, *dX, dWbase, dBbase); |
| 97 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | add_inplace(*dX, dXl); |
| 98 | 6 | } | |
| 99 | 6 | } | |
| 100 | |||
| 101 | } // namespace brotensor | ||
| 102 |