aboutsummaryrefslogtreecommitdiffstats
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/adapter.rs67
-rw-r--r--src/config.rs73
-rw-r--r--src/data.rs104
-rw-r--r--src/loader.rs116
-rw-r--r--src/lora.rs75
-rw-r--r--src/main.rs217
-rw-r--r--src/model.rs333
-rw-r--r--src/train.rs111
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` 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,