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

1 files changed, 333 insertions, 0 deletions

diff --git a/src/model.rs b/src/model.rs
new file mode 100644
index 0000000..acac139
--- /dev/null
+++ b/src/model.rs
@@ -0,0 +1,333 @@
+use burn::module::Module;
+use burn::module::ParamId;
+use burn::nn::{Embedding, EmbeddingConfig, RmsNorm, RmsNormConfig};
+use burn::prelude::*;
+use burn::tensor::activation::{silu, softmax};
+use burn::tensor::Bool;
+
+use crate::config::{HfConfig, LoraConfig};
+use crate::lora::LoraLinear;
+
+/// RoPE with the Hugging Face Llama "half-split" convention: head_dim is
+/// split in half and rotated as (x1, x2) -> (-x2, x1).
+/// Note: `burn::nn::RotaryEncoding` uses the interleaved convention (it
+/// rotates adjacent pairs), which is incompatible with these weights.
+#[derive(Module, Debug)]
+pub struct Rope<B: Backend> {
+    cos: Tensor<B, 2>,
+    sin: Tensor<B, 2>,
+}
+
+impl<B: Backend> Rope<B> {
+    pub fn new(head_dim: usize, max_seq: usize, theta: f32, device: &B::Device) -> Self {
+        let half = head_dim / 2;
+        let inv_freq: Vec<f32> = (0..half)
+            .map(|i| 1.0 / theta.powf(2.0 * i as f32 / head_dim as f32))
+            .collect();
+        let inv_freq = Tensor::<B, 1>::from_floats(inv_freq.as_slice(), device);
+        let pos = Tensor::<B, 1, Int>::arange(0..max_seq as i64, device).float();
+
+        // [max_seq, half] = outer(pos, inv_freq)
+        let freqs = pos
+            .unsqueeze_dim::<2>(1)
+            .matmul(inv_freq.unsqueeze_dim::<2>(0));
+        // duplicate to cover the whole head_dim: [f0..fh, f0..fh]
+        let emb = Tensor::cat(vec![freqs.clone(), freqs], 1);
+
+        Self {
+            cos: emb.clone().cos(),
+            sin: emb.sin(),
+        }
+    }
+
+    /// x: [batch, heads, seq, head_dim]
+    pub fn apply(&self, x: Tensor<B, 4>) -> Tensor<B, 4> {
+        let [b, h, t, hd] = x.dims();
+        let half = hd / 2;
+
+        let cos = self.cos.clone().slice([0..t, 0..hd]).unsqueeze::<4>();
+        let sin = self.sin.clone().slice([0..t, 0..hd]).unsqueeze::<4>();
+
+        let x1 = x.clone().slice([0..b, 0..h, 0..t, 0..half]);
+        let x2 = x.clone().slice([0..b, 0..h, 0..t, half..hd]);
+        let rotated = Tensor::cat(vec![-x2, x1], 3);
+
+        x * cos + rotated * sin
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct Attention<B: Backend> {
+    pub q: LoraLinear<B>,
+    pub k: LoraLinear<B>,
+    pub v: LoraLinear<B>,
+    pub o: LoraLinear<B>,
+    pub n_heads: usize,
+    pub n_kv_heads: usize,
+    pub head_dim: usize,
+}
+
+impl<B: Backend> Attention<B> {
+    fn new(cfg: &HfConfig, lora: &LoraConfig, device: &B::Device) -> Self {
+        let hd = cfg.head_dim();
+        let d = cfg.hidden_size;
+        let kv = cfg.num_key_value_heads * hd;
+        let r = |name: &str| {
+            if lora.targets(name) {
+                Some((lora.rank, lora.scaling()))
+            } else {
+                None
+            }
+        };
+        Self {
+            q: LoraLinear::new(d, d, r("q"), device),
+            k: LoraLinear::new(d, kv, r("k"), device),
+            v: LoraLinear::new(d, kv, r("v"), device),
+            o: LoraLinear::new(d, d, r("o"), device),
+            n_heads: cfg.num_attention_heads,
+            n_kv_heads: cfg.num_key_value_heads,
+            head_dim: hd,
+        }
+    }
+
+    /// Repeats each KV head `groups` times, interleaved, like Hugging Face
+    /// `repeat_kv`: [k0,k0,k1,k1,...].
+    fn repeat_kv(&self, x: Tensor<B, 4>, groups: usize) -> Tensor<B, 4> {
+        if groups == 1 {
+            return x;
+        }
+        let [b, n_kv, t, hd] = x.dims();
+        x.reshape([b, n_kv, 1, t, hd])
+            .repeat_dim(2, groups)
+            .reshape([b, n_kv * groups, t, hd])
+    }
+
+    pub fn forward(&self, x: Tensor<B, 3>, rope: &Rope<B>) -> Tensor<B, 3> {
+        let [b, t, d] = x.dims();
+        let hd = self.head_dim;
+        let groups = self.n_heads / self.n_kv_heads;
+        let device = x.device();
+
+        let q = self
+            .q
+            .forward(x.clone())
+            .reshape([b, t, self.n_heads, hd])
+            .swap_dims(1, 2);
+        let k = self
+            .k
+            .forward(x.clone())
+            .reshape([b, t, self.n_kv_heads, hd])
+            .swap_dims(1, 2);
+        let v = self
+            .v
+            .forward(x)
+            .reshape([b, t, self.n_kv_heads, hd])
+            .swap_dims(1, 2);
+
+        let q = rope.apply(q);
+        let k = rope.apply(k);
+
+        let k = self.repeat_kv(k, groups);
+        let v = self.repeat_kv(v, groups);
+
+        let scale = 1.0 / (hd as f64).sqrt();
+        let scores = q.matmul(k.swap_dims(2, 3)) * scale;
+
+        // Causal mask. Mind burn's semantics: `tril_mask` is `false` on the
+        // triangle that is KEPT (lower + diagonal) and `true` elsewhere,
+        // which is what gets filled with -inf.
+        let mask = Tensor::<B, 2, Bool>::tril_mask([t, t], 0, &device).unsqueeze::<4>();
+        let scores = scores.mask_fill(mask, f32::NEG_INFINITY);
+        let attn = softmax(scores, 3);
+
+        let out = attn.matmul(v).swap_dims(1, 2).reshape([b, t, d]);
+        self.o.forward(out)
+    }
+
+    fn freeze_base(self) -> Self {
+        Self {
+            q: self.q.freeze_base(),
+            k: self.k.freeze_base(),
+            v: self.v.freeze_base(),
+            o: self.o.freeze_base(),
+            ..self
+        }
+    }
+
+    fn trainable_ids(&self, out: &mut Vec<ParamId>) {
+        for l in [&self.q, &self.k, &self.v, &self.o] {
+            l.trainable_ids(out);
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct Mlp<B: Backend> {
+    pub gate: LoraLinear<B>,
+    pub up: LoraLinear<B>,
+    pub down: LoraLinear<B>,
+}
+
+impl<B: Backend> Mlp<B> {
+    fn new(cfg: &HfConfig, lora: &LoraConfig, device: &B::Device) -> Self {
+        let (d, i) = (cfg.hidden_size, cfg.intermediate_size);
+        let r = |name: &str| {
+            if lora.targets(name) {
+                Some((lora.rank, lora.scaling()))
+            } else {
+                None
+            }
+        };
+        Self {
+            gate: LoraLinear::new(d, i, r("gate"), device),
+            up: LoraLinear::new(d, i, r("up"), device),
+            down: LoraLinear::new(i, d, r("down"), device),
+        }
+    }
+
+    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
+        let gate = silu(self.gate.forward(x.clone()));
+        self.down.forward(gate * self.up.forward(x))
+    }
+
+    fn freeze_base(self) -> Self {
+        Self {
+            gate: self.gate.freeze_base(),
+            up: self.up.freeze_base(),
+            down: self.down.freeze_base(),
+        }
+    }
+
+    fn trainable_ids(&self, out: &mut Vec<ParamId>) {
+        for l in [&self.gate, &self.up, &self.down] {
+            l.trainable_ids(out);
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct Block<B: Backend> {
+    pub attn_norm: RmsNorm<B>,
+    pub attn: Attention<B>,
+    pub ffn_norm: RmsNorm<B>,
+    pub ffn: Mlp<B>,
+}
+
+impl<B: Backend> Block<B> {
+    fn new(cfg: &HfConfig, lora: &LoraConfig, device: &B::Device) -> Self {
+        let norm = || {
+            RmsNormConfig::new(cfg.hidden_size)
+                .with_epsilon(cfg.rms_norm_eps)
+                .init(device)
+        };
+        Self {
+            attn_norm: norm(),
+            attn: Attention::new(cfg, lora, device),
+            ffn_norm: norm(),
+            ffn: Mlp::new(cfg, lora, device),
+        }
+    }
+
+    pub fn forward(&self, x: Tensor<B, 3>, rope: &Rope<B>) -> Tensor<B, 3> {
+        let h = x.clone() + self.attn.forward(self.attn_norm.forward(x), rope);
+        h.clone() + self.ffn.forward(self.ffn_norm.forward(h))
+    }
+
+    fn freeze_base(self) -> Self {
+        Self {
+            attn: self.attn.freeze_base(),
+            ffn: self.ffn.freeze_base(),
+            ..self
+        }
+    }
+
+    fn trainable_ids(&self, out: &mut Vec<ParamId>) {
+        self.attn.trainable_ids(out);
+        self.ffn.trainable_ids(out);
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct Llama<B: Backend> {
+    pub embed: Embedding<B>,
+    pub layers: Vec<Block<B>>,
+    pub norm: RmsNorm<B>,
+    pub rope: Rope<B>,
+}
+
+impl<B: Backend> Llama<B> {
+    pub fn new(cfg: &HfConfig, lora: &LoraConfig, max_seq: usize, device: &B::Device) -> Self {
+        Self {
+            embed: EmbeddingConfig::new(cfg.vocab_size, cfg.hidden_size).init(device),
+            layers: (0..cfg.num_hidden_layers)
+                .map(|_| Block::new(cfg, lora, device))
+                .collect(),
+            norm: RmsNormConfig::new(cfg.hidden_size)
+                .with_epsilon(cfg.rms_norm_eps)
+                .init(device),
+            rope: Rope::new(cfg.head_dim(), max_seq, cfg.rope_theta, device),
+        }
+    }
+
+    /// tokens: [batch, seq] -> logits [batch, seq, vocab]
+    pub fn forward(&self, tokens: Tensor<B, 2, Int>) -> Tensor<B, 3> {
+        let mut h = self.embed.forward(tokens);
+        for layer in &self.layers {
+            h = layer.forward(h, &self.rope);
+        }
+        let h = self.norm.forward(h);
+        // tied embeddings: lm_head = embed.weight^T
+        h.matmul(self.embed.weight.val().transpose().unsqueeze())
+    }
+
+    /// Freezes all base weights; only the LoRA A/B stay trainable.
+    pub fn freeze_base(self) -> Self {
+        let Self {
+            embed,
+            layers,
+            norm,
+            rope,
+        } = self;
+        Self {
+            embed: Embedding {
+                weight: embed.weight.set_require_grad(false),
+            },
+            layers: layers.into_iter().map(|l| l.freeze_base()).collect(),
+            norm: RmsNorm {
+                gamma: norm.gamma.set_require_grad(false),
+                epsilon: norm.epsilon,
+            },
+            rope,
+        }
+    }
+
+    /// L2 norm of all B matrices together. They start at zero, so this is a
+    /// fingerprint to check that an adapter was trained and that it loads
+    /// back identical.
+    pub fn lora_norm(&self) -> f32 {
+        let mut acc = 0.0f32;
+        for l in &self.layers {
+            let mut push = |lin: &LoraLinear<B>| {
+                if let Some(ad) = &lin.lora {
+                    let v: f32 = ad.b.val().powf_scalar(2.0).sum().into_scalar().elem();
+                    acc += v;
+                }
+            };
+            for lin in [&l.attn.q, &l.attn.k, &l.attn.v, &l.attn.o] {
+                push(lin);
+            }
+            for lin in [&l.ffn.gate, &l.ffn.up, &l.ffn.down] {
+                push(lin);
+            }
+        }
+        acc.sqrt()
+    }
+
+    pub fn trainable_ids(&self) -> Vec<ParamId> {
+        let mut out = Vec::new();
+        for l in &self.layers {
+            l.trainable_ids(&mut out);
+        }
+        out
+    }
+}