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

1 files changed, 67 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())
+}