1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
|
use burn::module::{Module, Param, ParamId};
use burn::nn::{Linear, LinearConfig};
use burn::prelude::*;
use burn::tensor::Distribution;
/// Low-rank adapter: delta_W = (A @ B) * scaling.
/// A starts as Gaussian noise and B as zeros, so at the start of training
/// delta_W = 0 and the model is identical to the base one.
#[derive(Module, Debug)]
pub struct LoraAdapter<B: Backend> {
/// [d_in, rank]
pub a: Param<Tensor<B, 2>>,
/// [rank, d_out]
pub b: Param<Tensor<B, 2>>,
pub scaling: f64,
}
impl<B: Backend> LoraAdapter<B> {
pub fn new(d_in: usize, d_out: usize, rank: usize, scaling: f64, device: &B::Device) -> Self {
let std = 1.0 / (rank as f64).sqrt();
let a = Tensor::random([d_in, rank], Distribution::Normal(0.0, std), device);
let b = Tensor::zeros([rank, d_out], device);
Self {
a: Param::from_tensor(a),
b: Param::from_tensor(b),
scaling,
}
}
}
/// Frozen Linear + optional LoRA adapter.
#[derive(Module, Debug)]
pub struct LoraLinear<B: Backend> {
pub base: Linear<B>,
pub lora: Option<LoraAdapter<B>>,
}
impl<B: Backend> LoraLinear<B> {
pub fn new(d_in: usize, d_out: usize, rank: Option<(usize, f64)>, device: &B::Device) -> Self {
let base = LinearConfig::new(d_in, d_out).with_bias(false).init(device);
let lora = rank.map(|(r, scaling)| LoraAdapter::new(d_in, d_out, r, scaling, device));
Self { base, lora }
}
pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
let out = self.base.forward(x.clone());
match &self.lora {
None => out,
Some(lora) => {
let delta = x
.matmul(lora.a.val().unsqueeze())
.matmul(lora.b.val().unsqueeze());
out + delta * lora.scaling
}
}
}
/// Freezes the base weight: it stops accumulating gradients in the autodiff graph.
pub fn freeze_base(self) -> Self {
let Self { base, lora } = self;
let base = Linear {
weight: base.weight.set_require_grad(false),
bias: base.bias.map(|b| b.set_require_grad(false)),
};
Self { base, lora }
}
/// IDs of the trainable parameters (only A and B).
pub fn trainable_ids(&self, out: &mut Vec<ParamId>) {
if let Some(lora) = &self.lora {
out.push(lora.a.id);
out.push(lora.b.id);
}
}
}
|
