diff options
Diffstat (limited to 'src/main.rs')
| -rw-r--r-- | src/main.rs | 217 |
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(()) +} |