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

1 files changed, 116 insertions, 0 deletions

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