1 files changed, 111 insertions, 0 deletions
diff --git a/src/train.rs b/src/train.rs
new file mode 100644
index 0000000..87cc629
--- /dev/null
+++ b/src/train.rs
@@ -0,0 +1,111 @@
+use anyhow::Result;
+use burn::optim::{AdamWConfig, GradientsParams, Optimizer};
+use burn::prelude::*;
+use burn::tensor::activation::log_softmax;
+use burn::tensor::backend::AutodiffBackend;
+use burn::tensor::TensorData;
+
+use crate::data::{collate, Example};
+use crate::model::Llama;
+
+/// Causal cross-entropy with a loss mask.
+///
+/// The shift is the key part: position `i` predicts token `i+1`. And `mask`
+/// makes only the assistant answer count, so the prompt adds no gradient.
+/// This is exactly what llama.cpp's `llama-finetune` does not do.
+pub fn masked_ce_loss<B: Backend>(
+    logits: Tensor<B, 3>,
+    ids: Tensor<B, 2, Int>,
+    mask: Tensor<B, 2>,
+) -> Tensor<B, 1> {
+    let [b, t, v] = logits.dims();
+    let n = b * (t - 1);
+
+    let logits = logits.slice([0..b, 0..t - 1, 0..v]).reshape([n, v]);
+    let targets = ids.slice([0..b, 1..t]).reshape([n, 1]);
+    let mask = mask.slice([0..b, 1..t]).reshape([n, 1]);
+
+    let logp = log_softmax(logits, 1).gather(1, targets);
+    let denom = mask.clone().sum().clamp_min(1.0);
+
+    -(logp * mask).sum() / denom
+}
+
+pub struct TrainCfg {
+    pub steps: usize,
+    pub batch_size: usize,
+    pub lr: f64,
+}
+
+pub fn train<B: AutodiffBackend>(
+    mut model: Llama<B>,
+    examples: &[Example],
+    cfg: TrainCfg,
+    device: &B::Device,
+) -> Result<Llama<B>> {
+    let trainable = model.trainable_ids();
+    println!(
+        "training {} LoRA tensors on {} examples\n",
+        trainable.len(),
+        examples.len()
+    );
+
+    // B starts at zero. If it is still zero at the end, the gradients never
+    // reached LoRA; if W_base changes, freezing failed.
+    let norm_base = |m: &Llama<B>| -> f32 {
+        m.layers[0]
+            .attn
+            .q
+            .base
+            .weight
+            .val()
+            .powf_scalar(2.0)
+            .sum()
+            .sqrt()
+            .into_scalar()
+            .elem()
+    };
+    println!(
+        "before -> |B| = {:.6}   |W_base| = {:.4}",
+        model.lora_norm(),
+        norm_base(&model)
+    );
+
+    let mut optim = AdamWConfig::new().init::<B, Llama<B>>();
+    let mut cursor = 0usize;
+
+    for step in 1..=cfg.steps {
+        if cursor + cfg.batch_size > examples.len() {
+            cursor = 0;
+        }
+        let batch = collate(&examples[cursor..cursor + cfg.batch_size]);
+        cursor += cfg.batch_size;
+
+        let dims = [batch.batch, batch.seq];
+        let ids = Tensor::<B, 2, Int>::from_data(TensorData::new(batch.ids, dims), device);
+        let mask = Tensor::<B, 2>::from_data(TensorData::new(batch.mask, dims), device);
+
+        let logits = model.forward(ids.clone());
+        let loss = masked_ce_loss(logits, ids, mask);
+        let value: f32 = loss.clone().into_scalar().elem();
+
+        let mut grads = loss.backward();
+        // The heart of LoRA: only the gradients of the A and B tensors are
+        // collected. The rest of the model is never touched.
+        let grads = GradientsParams::from_params(&mut grads, &model, &trainable);
+        model = optim.step(cfg.lr, model, grads);
+
+        println!(
+            "step {step:4}/{}  loss {value:.4}  seq {}",
+            cfg.steps, batch.seq
+        );
+    }
+
+    println!(
+        "after  -> |B| = {:.6}   |W_base| = {:.4}",
+        model.lora_norm(),
+        norm_base(&model)
+    );
+
+    Ok(model)
+}