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

1 files changed, 73 insertions, 0 deletions

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