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