1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
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)
}
}
|
