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

8 files changed, 1096 insertions, 0 deletions

diff --git a/src/adapter.rs b/src/adapter.rs
new file mode 100644
index 0000000..efc94ff
--- /dev/null
+++ b/src/adapter.rs
@@ -0,0 +1,67 @@
+use anyhow::{bail, Result};
+use burn::prelude::*;
+use burn::store::{ModuleSnapshot, SafetensorsStore};
+use std::path::Path;
+
+use crate::config::LoraConfig;
+use crate::model::Llama;
+
+const WEIGHTS: &str = "adapter.safetensors";
+const CONFIG: &str = "adapter_config.json";
+
+/// Only the adapter tensors: `....lora.a` and `....lora.b`.
+/// Everything else (base weights, norms, embeddings) stays out of the file.
+const LORA_RE: &str = r"\.lora\.[ab]$";
+
+/// Saves only the adapter. No `ModuleAdapter` is applied: these tensors are
+/// created and loaded in burn layout, so transposing them like the HF
+/// weights would corrupt them.
+pub fn save<B: Backend>(model: &Llama<B>, lora: &LoraConfig, dir: impl AsRef<Path>) -> Result<()> {
+    let dir = dir.as_ref();
+    std::fs::create_dir_all(dir)?;
+
+    let mut store = SafetensorsStore::from_file(dir.join(WEIGHTS)).with_regex(LORA_RE);
+    model.save_into(&mut store)?;
+
+    std::fs::write(dir.join(CONFIG), serde_json::to_string_pretty(lora)?)?;
+    Ok(())
+}
+
+/// Reads the adapter config. It is needed BEFORE building the model: the
+/// rank defines the shape of A and B, and the targets which projections get them.
+pub fn load_config(dir: impl AsRef<Path>) -> Result<LoraConfig> {
+    let path = dir.as_ref().join(CONFIG);
+    if !path.exists() {
+        bail!("{} does not exist", path.display());
+    }
+    Ok(serde_json::from_str(&std::fs::read_to_string(&path)?)?)
+}
+
+/// Applies an adapter to a model already loaded with the base weights.
+pub fn load<B: Backend>(model: &mut Llama<B>, dir: impl AsRef<Path>) -> Result<usize> {
+    let path = dir.as_ref().join(WEIGHTS);
+    if !path.exists() {
+        bail!("{} does not exist", path.display());
+    }
+
+    // `allow_partial`: the file only holds LoRA; the rest of the model keeps
+    // the base weights that were already loaded.
+    let mut store = SafetensorsStore::from_file(&path).allow_partial(true);
+    let result = model.load_from(&mut store)?;
+
+    if !result.errors.is_empty() {
+        bail!("errors applying the adapter: {:?}", result.errors);
+    }
+    // `allow_partial` also lets an incomplete file through: tensors that are
+    // not applied would keep their random init and silently corrupt the
+    // output. Require all of them.
+    let expected = model.trainable_ids().len();
+    if result.applied.len() != expected {
+        bail!(
+            "the adapter applied {} of {} expected tensors; different rank or targets?",
+            result.applied.len(),
+            expected
+        );
+    }
+    Ok(result.applied.len())
+}
diff --git a/src/config.rs b/src/config.rs
new file mode 100644
index 0000000..71789c5
--- /dev/null
+++ b/src/config.rs
@@ -0,0 +1,73 @@
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
+use std::path::Path;
+
+/// The subset of Hugging Face `config.json` needed for Llama.
+#[derive(Debug, Clone, Deserialize)]
+pub struct HfConfig {
+    pub hidden_size: usize,
+    pub intermediate_size: usize,
+    pub num_hidden_layers: usize,
+    pub num_attention_heads: usize,
+    pub num_key_value_heads: usize,
+    pub vocab_size: usize,
+    pub rms_norm_eps: f64,
+    pub rope_theta: f32,
+    pub max_position_embeddings: usize,
+    #[serde(default)]
+    pub tie_word_embeddings: bool,
+}
+
+impl HfConfig {
+    pub fn load(dir: impl AsRef<Path>) -> Result<Self> {
+        let path = dir.as_ref().join("config.json");
+        let raw = std::fs::read_to_string(&path)
+            .map_err(|e| anyhow::anyhow!("{}: {e}", path.display()))?;
+        let cfg: Self = serde_json::from_str(&raw)?;
+        // the model has no separate lm_head: it reuses the embedding matrix
+        if !cfg.tie_word_embeddings {
+            anyhow::bail!("only models with tie_word_embeddings = true are supported");
+        }
+        Ok(cfg)
+    }
+
+    pub fn head_dim(&self) -> usize {
+        self.hidden_size / self.num_attention_heads
+    }
+
+    /// Longest sequence the RoPE table covers. Capped so the table stays
+    /// small on CPU; SmolLM2 declares 8192.
+    pub fn max_seq(&self) -> usize {
+        self.max_position_embeddings.min(2048)
+    }
+}
+
+/// LoRA hyperparameters. Saved next to the adapter because the rank
+/// defines the shape of A and B.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LoraConfig {
+    pub rank: usize,
+    pub alpha: f64,
+    /// Projections that get an adapter.
+    pub targets: Vec<String>,
+}
+
+impl Default for LoraConfig {
+    fn default() -> Self {
+        Self {
+            rank: 8,
+            alpha: 16.0,
+            targets: vec!["q".into(), "v".into()],
+        }
+    }
+}
+
+impl LoraConfig {
+    pub fn scaling(&self) -> f64 {
+        self.alpha / self.rank as f64
+    }
+
+    pub fn targets(&self, name: &str) -> bool {
+        self.targets.iter().any(|t| t == name)
+    }
+}
diff --git a/src/data.rs b/src/data.rs
new file mode 100644
index 0000000..2224583
--- /dev/null
+++ b/src/data.rs
@@ -0,0 +1,104 @@
+use anyhow::Result;
+use serde::Deserialize;
+use std::path::Path;
+use tokenizers::Tokenizer;
+
+#[derive(Debug, Deserialize)]
+struct DollyRow {
+    instruction: String,
+    #[serde(default)]
+    context: String,
+    response: String,
+}
+
+/// A tokenized example. `mask[i] = 1.0` marks the positions that count
+/// towards the loss: only the answer, never the prompt.
+#[derive(Debug, Clone)]
+pub struct Example {
+    pub ids: Vec<i64>,
+    pub mask: Vec<f32>,
+}
+
+fn encode(tok: &Tokenizer, text: &str) -> Result<Vec<i64>> {
+    Ok(tok
+        .encode(text, false)
+        .map_err(|e| anyhow::anyhow!("{e}"))?
+        .get_ids()
+        .iter()
+        .map(|&i| i as i64)
+        .collect())
+}
+
+pub fn load_dolly(
+    path: impl AsRef<Path>,
+    tok: &Tokenizer,
+    max_len: usize,
+    limit: usize,
+) -> Result<Vec<Example>> {
+    let raw = std::fs::read_to_string(path)?;
+    let mut out = Vec::new();
+
+    for line in raw.lines() {
+        if out.len() >= limit {
+            break;
+        }
+        let row: DollyRow = match serde_json::from_str(line) {
+            Ok(r) => r,
+            Err(_) => continue,
+        };
+
+        // ChatML, the SmolLM2 chat template
+        let user = if row.context.trim().is_empty() {
+            row.instruction.clone()
+        } else {
+            format!("{}\n\n{}", row.instruction, row.context)
+        };
+        let prompt = format!("<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n");
+        let answer = format!("{}<|im_end|>", row.response);
+
+        let p = encode(tok, &prompt)?;
+        let a = encode(tok, &answer)?;
+        if p.len() + a.len() > max_len || a.is_empty() {
+            continue;
+        }
+
+        let mut ids = p.clone();
+        ids.extend_from_slice(&a);
+        let mut mask = vec![0.0; p.len()];
+        mask.extend(std::iter::repeat(1.0).take(a.len()));
+
+        out.push(Example { ids, mask });
+    }
+    Ok(out)
+}
+
+/// Pads a batch to its longest example. Padding has mask 0, so it adds
+/// nothing to the loss.
+pub struct Batch {
+    pub ids: Vec<i64>,
+    pub mask: Vec<f32>,
+    pub batch: usize,
+    pub seq: usize,
+}
+
+pub fn collate(examples: &[Example]) -> Batch {
+    let seq = examples.iter().map(|e| e.ids.len()).max().unwrap_or(1);
+    let batch = examples.len();
+    let mut ids = Vec::with_capacity(batch * seq);
+    let mut mask = Vec::with_capacity(batch * seq);
+
+    for e in examples {
+        ids.extend_from_slice(&e.ids);
+        mask.extend_from_slice(&e.mask);
+        for _ in e.ids.len()..seq {
+            ids.push(0);
+            mask.push(0.0);
+        }
+    }
+    Batch {
+        ids,
+        mask,
+        batch,
+        seq,
+    }
+}
diff --git a/src/loader.rs b/src/loader.rs
new file mode 100644
index 0000000..8d9c76f
--- /dev/null
+++ b/src/loader.rs
@@ -0,0 +1,116 @@
+use anyhow::{bail, Result};
+use burn::prelude::*;
+use burn::store::{
+    ModuleAdapter, ModuleSnapshot, PyTorchToBurnAdapter, SafetensorsStore, TensorSnapshot,
+};
+use burn::tensor::DType;
+use std::path::Path;
+use std::rc::Rc;
+
+/// The weights ship as BF16, which the ndarray backend does not support.
+/// burn's `HalfPrecisionAdapter` only covers F32<->F16, so BF16 -> F32 is
+/// converted here, lazily.
+#[derive(Debug, Clone, Default)]
+struct Bf16ToF32;
+
+impl ModuleAdapter for Bf16ToF32 {
+    fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot {
+        if snapshot.dtype != DType::BF16 {
+            return snapshot.clone();
+        }
+        let original = snapshot.clone_data_fn();
+        let cast = Rc::new(move || Ok(original()?.convert_dtype(DType::F32)));
+        TensorSnapshot::from_closure(
+            cast,
+            DType::F32,
+            snapshot.shape.clone(),
+            snapshot.path_stack.clone().unwrap_or_default(),
+            snapshot.container_stack.clone().unwrap_or_default(),
+            snapshot.tensor_id.unwrap_or_default(),
+        )
+    }
+
+    fn clone_box(&self) -> Box<dyn ModuleAdapter> {
+        Box::new(self.clone())
+    }
+}
+
+use crate::model::Llama;
+
+/// Renaming rules: Hugging Face names -> burn module paths.
+/// Base Linear layers live under `.base.weight` because they are wrapped in LoraLinear.
+const REMAP: &[(&str, &str)] = &[
+    (r"^model\.embed_tokens\.weight$", "embed.weight"),
+    (r"^model\.norm\.weight$", "norm.gamma"),
+    (
+        r"^model\.layers\.(\d+)\.input_layernorm\.weight$",
+        "layers.${1}.attn_norm.gamma",
+    ),
+    (
+        r"^model\.layers\.(\d+)\.post_attention_layernorm\.weight$",
+        "layers.${1}.ffn_norm.gamma",
+    ),
+    (
+        r"^model\.layers\.(\d+)\.self_attn\.q_proj\.weight$",
+        "layers.${1}.attn.q.base.weight",
+    ),
+    (
+        r"^model\.layers\.(\d+)\.self_attn\.k_proj\.weight$",
+        "layers.${1}.attn.k.base.weight",
+    ),
+    (
+        r"^model\.layers\.(\d+)\.self_attn\.v_proj\.weight$",
+        "layers.${1}.attn.v.base.weight",
+    ),
+    (
+        r"^model\.layers\.(\d+)\.self_attn\.o_proj\.weight$",
+        "layers.${1}.attn.o.base.weight",
+    ),
+    (
+        r"^model\.layers\.(\d+)\.mlp\.gate_proj\.weight$",
+        "layers.${1}.ffn.gate.base.weight",
+    ),
+    (
+        r"^model\.layers\.(\d+)\.mlp\.up_proj\.weight$",
+        "layers.${1}.ffn.up.base.weight",
+    ),
+    (
+        r"^model\.layers\.(\d+)\.mlp\.down_proj\.weight$",
+        "layers.${1}.ffn.down.base.weight",
+    ),
+];
+
+pub fn load_hf_weights<B: Backend>(model: &mut Llama<B>, dir: impl AsRef<Path>) -> Result<()> {
+    let path = dir.as_ref().join("model.safetensors");
+    if !path.exists() {
+        bail!("{} does not exist", path.display());
+    }
+
+    let mut store = SafetensorsStore::from_file(&path)
+        // BF16 -> F32, then transpose the Linear weights:
+        // HF stores [out, in], burn expects [in, out]
+        .with_from_adapter(Bf16ToF32.chain(PyTorchToBurnAdapter))
+        // the LoRA tensors are not in the file; they keep their init
+        .allow_partial(true);
+
+    for (from, to) in REMAP {
+        store = store.with_key_remapping(*from, *to);
+    }
+
+    let result = model.load_from(&mut store)?;
+
+    println!("weights applied: {}", result.applied.len());
+    if !result.missing.is_empty() {
+        println!("missing        : {}", result.missing.len());
+        for (k, _) in result.missing.iter().take(5) {
+            println!("  - {k}");
+        }
+    }
+    if !result.unused.is_empty() {
+        println!("unused         : {}", result.unused.len());
+        for k in result.unused.iter().take(5) {
+            println!("  - {k}");
+        }
+    }
+    Ok(())
+}
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);
+        }
+    }
+}
diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..8d0bf4b
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,217 @@
+mod adapter;
+mod config;
+mod data;
+mod loader;
+mod lora;
+mod model;
+mod train;
+
+use anyhow::Result;
+use burn::backend::ndarray::NdArrayDevice;
+use burn::backend::{Autodiff, NdArray};
+use burn::module::AutodiffModule;
+use burn::prelude::*;
+use burn::tensor::TensorData;
+use clap::{Parser, Subcommand};
+use std::path::PathBuf;
+use tokenizers::Tokenizer;
+
+use config::{HfConfig, LoraConfig};
+use model::Llama;
+
+const MODEL_DIR: &str = "models/smollm2-135m-instruct";
+const DATA: &str = "data/dolly-15k/databricks-dolly-15k.jsonl";
+
+type Back = Autodiff<NdArray<f32>>;
+type Inner = NdArray<f32>;
+
+#[derive(Parser)]
+#[command(about = "LoRA fine-tuning for Llama models, written in Rust with burn")]
+struct Cli {
+    /// Hugging Face model directory (config.json, tokenizer.json, model.safetensors).
+    #[arg(long, global = true, default_value = MODEL_DIR)]
+    model: PathBuf,
+    #[command(subcommand)]
+    cmd: Cmd,
+}
+
+#[derive(Subcommand)]
+enum Cmd {
+    /// Load the weights and generate text, to check the model is wired correctly.
+    Check {
+        #[arg(default_value = "The capital of France is")]
+        prompt: String,
+        #[arg(long, default_value_t = 12)]
+        tokens: usize,
+        /// Directory of a LoRA adapter to apply on top of the base model.
+        #[arg(long)]
+        adapter: Option<PathBuf>,
+    },
+    /// Train a LoRA adapter on dolly-15k.
+    Train {
+        /// dolly-15k JSONL file.
+        #[arg(long, default_value = DATA)]
+        data: PathBuf,
+        #[arg(long, default_value_t = 20)]
+        steps: usize,
+        #[arg(long, default_value_t = 1)]
+        batch: usize,
+        #[arg(long, default_value_t = 1e-4)]
+        lr: f64,
+        #[arg(long, default_value_t = 256)]
+        max_len: usize,
+        #[arg(long, default_value_t = 256)]
+        examples: usize,
+        /// Where to save the trained adapter.
+        #[arg(long, default_value = "adapters/dolly")]
+        out: PathBuf,
+    },
+}
+
+fn generate(
+    model: &Llama<Inner>,
+    tok: &Tokenizer,
+    cfg: &HfConfig,
+    prompt: &str,
+    n: usize,
+    device: &NdArrayDevice,
+) -> Result<()> {
+    let mut ids: Vec<i64> = tok
+        .encode(prompt, false)
+        .map_err(|e| anyhow::anyhow!("{e}"))?
+        .get_ids()
+        .iter()
+        .map(|&i| i as i64)
+        .collect();
+
+    print!("{prompt}");
+    use std::io::Write;
+    for _ in 0..n {
+        let len = ids.len();
+        let input =
+            Tensor::<Inner, 2, Int>::from_data(TensorData::new(ids.clone(), [1, len]), device);
+        let next = model
+            .forward(input)
+            .slice([0..1, len - 1..len, 0..cfg.vocab_size])
+            .reshape([cfg.vocab_size])
+            .argmax(0)
+            .into_scalar() as i64;
+        // SmolLM2: 0 = <|endoftext|>, 2 = <|im_end|>
+        if next == 2 || next == 0 {
+            break;
+        }
+        ids.push(next);
+        print!(
+            "{}",
+            tok.decode(&[next as u32], false)
+                .map_err(|e| anyhow::anyhow!("{e}"))?
+        );
+        std::io::stdout().flush().ok();
+    }
+    println!("\n");
+    Ok(())
+}
+
+fn main() -> Result<()> {
+    let cli = Cli::parse();
+    let device = Default::default();
+
+    let model_dir = cli.model.as_path();
+    let cfg = HfConfig::load(model_dir)?;
+    let tok = Tokenizer::from_file(model_dir.join("tokenizer.json"))
+        .map_err(|e| anyhow::anyhow!("{e}"))?;
+
+    println!(
+        "model : {} layers, hidden {}, {} heads / {} kv, head_dim {}, vocab {}",
+        cfg.num_hidden_layers,
+        cfg.hidden_size,
+        cfg.num_attention_heads,
+        cfg.num_key_value_heads,
+        cfg.head_dim(),
+        cfg.vocab_size
+    );
+    let describe = |l: &LoraConfig| {
+        println!(
+            "lora  : rank {}, alpha {}, scaling {:.2}, targets {:?}",
+            l.rank,
+            l.alpha,
+            l.scaling(),
+            l.targets
+        );
+    };
+
+    match cli.cmd {
+        Cmd::Check {
+            prompt,
+            tokens,
+            adapter: adapter_dir,
+        } => {
+            // The adapter config wins: it sets the rank and targets,
+            // and therefore the shape of the tensors to create.
+            let lora = match &adapter_dir {
+                Some(dir) => adapter::load_config(dir)?,
+                None => LoraConfig::default(),
+            };
+            describe(&lora);
+
+            let mut model = Llama::<Inner>::new(&cfg, &lora, cfg.max_seq(), &device);
+            loader::load_hf_weights(&mut model, model_dir)?;
+
+            if let Some(dir) = &adapter_dir {
+                let n = adapter::load(&mut model, dir)?;
+                println!(
+                    "adapter        : {n} tensors from {}  |B| = {:.6}",
+                    dir.display(),
+                    model.lora_norm()
+                );
+            }
+            println!();
+            generate(&model, &tok, &cfg, &prompt, tokens, &device)?;
+        }
+        Cmd::Train {
+            data,
+            steps,
+            batch,
+            lr,
+            max_len,
+            examples,
+            out,
+        } => {
+            let lora = LoraConfig::default();
+            describe(&lora);
+            let mut model = Llama::<Back>::new(&cfg, &lora, cfg.max_seq(), &device);
+            loader::load_hf_weights(&mut model, model_dir)?;
+            // freeze AFTER loading: the base weights stop requiring gradients
+            let model = model.freeze_base();
+
+            let data = data::load_dolly(&data, &tok, max_len, examples)?;
+            println!("examples loaded: {}\n", data.len());
+
+            let cfg_t = train::TrainCfg {
+                steps,
+                batch_size: batch,
+                lr,
+            };
+            let model = train::train(model, &data, cfg_t, &device)?;
+
+            adapter::save(&model, &lora, &out)?;
+            println!(
+                "\nadapter saved to {}/  (|B| = {:.6})",
+                out.display(),
+                model.lora_norm()
+            );
+
+            println!("\n--- sample after training ---");
+            let infer = model.valid();
+            generate(
+                &infer,
+                &tok,
+                &cfg,
+                "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n",
+                40,
+                &device,
+            )?;
+        }
+    }
+    Ok(())
+}
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
+    }
+}
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`