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

1 files changed, 217 insertions, 0 deletions

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(())
+}