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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
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(())
}
|
