llm-tune: LoRA fine-tuning for Llama in Rust with burnHEADmain

1 files changed, 75 insertions, 0 deletions

diff --git a/src/lora.rs b/src/lora.rs
new file mode 100644
index 0000000..169289f
--- /dev/null
+++ b/src/lora.rs
@@ -0,0 +1,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);
+        }
+    }
+}