feat(inference): 添加模型加载和分词器功能
feat(model): 实现 Llama 模型及其解码器层 feat(norm): 添加 RMSNorm 归一化层 feat(rope): 实现 RoPE 位置编码
This commit is contained in:
@@ -0,0 +1,226 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::model::LlamaForCausalLM;
|
||||
use burn::module::{Module, Param};
|
||||
use burn::nn::{EmbeddingRecord, LinearRecord};
|
||||
use burn::tensor::backend::Backend;
|
||||
use burn::tensor::{Float, Shape, Tensor, TensorData};
|
||||
use memmap2::Mmap;
|
||||
use std::path::Path;
|
||||
|
||||
pub fn load_safetensors<B: Backend>(
|
||||
model: LlamaForCausalLM<B>,
|
||||
path: &Path,
|
||||
device: &B::Device,
|
||||
) -> anyhow::Result<LlamaForCausalLM<B>> {
|
||||
let file = std::fs::File::open(path)?;
|
||||
let mmap = unsafe { Mmap::map(&file)? };
|
||||
|
||||
let header_len = u64::from_le_bytes(mmap[..8].try_into().unwrap()) as usize;
|
||||
let header_bytes = &mmap[8..8 + header_len];
|
||||
let header: serde_json::Value = serde_json::from_slice(header_bytes)?;
|
||||
|
||||
let mut tensors = HashMap::new();
|
||||
let mut shapes = HashMap::new();
|
||||
let mut dtypes = HashMap::new();
|
||||
|
||||
let data_offset = 8 + header_len;
|
||||
|
||||
if let Some(obj) = header.as_object() {
|
||||
for (name, info) in obj {
|
||||
if name == "__metadata__" {
|
||||
continue;
|
||||
}
|
||||
if let Some(info_obj) = info.as_object() {
|
||||
if let Some(offsets) = info_obj.get("data_offsets") {
|
||||
if let Some(arr) = offsets.as_array() {
|
||||
let start = arr[0].as_u64().unwrap() as usize;
|
||||
let end = arr[1].as_u64().unwrap() as usize;
|
||||
let data = &mmap[data_offset + start..data_offset + end];
|
||||
tensors.insert(name.clone(), data.to_vec());
|
||||
}
|
||||
}
|
||||
if let Some(shape) = info_obj.get("shape") {
|
||||
if let Some(arr) = shape.as_array() {
|
||||
let shape_vec: Vec<usize> = arr
|
||||
.iter()
|
||||
.map(|v| v.as_u64().unwrap() as usize)
|
||||
.collect();
|
||||
shapes.insert(name.clone(), shape_vec);
|
||||
}
|
||||
}
|
||||
if let Some(dtype) = info_obj.get("dtype") {
|
||||
if let Some(s) = dtype.as_str() {
|
||||
dtypes.insert(name.clone(), s.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut model = model;
|
||||
|
||||
if let Some(weight) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, "model.embed_tokens.weight", device) {
|
||||
let record = EmbeddingRecord {
|
||||
weight: Param::from_tensor(weight),
|
||||
};
|
||||
model.model.embed_tokens = model.model.embed_tokens.clone().load_record(record);
|
||||
}
|
||||
|
||||
for i in 0..model.config.num_hidden_layers {
|
||||
let prefix = format!("model.layers.{i}");
|
||||
|
||||
if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.q_proj.weight"), device) {
|
||||
let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
|
||||
model.model.layers[i].self_attn.q_proj = model.model.layers[i].self_attn.q_proj.clone().load_record(record);
|
||||
}
|
||||
if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.k_proj.weight"), device) {
|
||||
let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
|
||||
model.model.layers[i].self_attn.k_proj = model.model.layers[i].self_attn.k_proj.clone().load_record(record);
|
||||
}
|
||||
if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.v_proj.weight"), device) {
|
||||
let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
|
||||
model.model.layers[i].self_attn.v_proj = model.model.layers[i].self_attn.v_proj.clone().load_record(record);
|
||||
}
|
||||
if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.o_proj.weight"), device) {
|
||||
let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
|
||||
model.model.layers[i].self_attn.o_proj = model.model.layers[i].self_attn.o_proj.clone().load_record(record);
|
||||
}
|
||||
|
||||
if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.gate_proj.weight"), device) {
|
||||
let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
|
||||
model.model.layers[i].mlp.gate_proj = model.model.layers[i].mlp.gate_proj.clone().load_record(record);
|
||||
}
|
||||
if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.up_proj.weight"), device) {
|
||||
let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
|
||||
model.model.layers[i].mlp.up_proj = model.model.layers[i].mlp.up_proj.clone().load_record(record);
|
||||
}
|
||||
if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.down_proj.weight"), device) {
|
||||
let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
|
||||
model.model.layers[i].mlp.down_proj = model.model.layers[i].mlp.down_proj.clone().load_record(record);
|
||||
}
|
||||
|
||||
if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.input_layernorm.weight"), device) {
|
||||
model.model.layers[i].input_layernorm.weight = Param::from_tensor(w);
|
||||
}
|
||||
if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.post_attention_layernorm.weight"), device) {
|
||||
model.model.layers[i].post_attention_layernorm.weight = Param::from_tensor(w);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, "model.norm.weight", device) {
|
||||
model.model.norm.weight = Param::from_tensor(w);
|
||||
}
|
||||
|
||||
if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, "lm_head.weight", device) {
|
||||
let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
|
||||
model.lm_head = model.lm_head.clone().load_record(record);
|
||||
} else if model.config.tie_word_embeddings {
|
||||
let embed_weight = model.model.embed_tokens.weight.val();
|
||||
let record = LinearRecord { weight: Param::from_tensor(embed_weight.transpose()), bias: None };
|
||||
model.lm_head = model.lm_head.clone().load_record(record);
|
||||
}
|
||||
|
||||
Ok(model)
|
||||
}
|
||||
|
||||
fn load_tensor_2d<B: Backend>(
|
||||
tensors: &HashMap<String, Vec<u8>>,
|
||||
shapes: &HashMap<String, Vec<usize>>,
|
||||
dtypes: &HashMap<String, String>,
|
||||
name: &str,
|
||||
device: &B::Device,
|
||||
) -> Option<Tensor<B, 2, Float>> {
|
||||
let data = tensors.get(name)?;
|
||||
let shape = shapes.get(name)?;
|
||||
let dtype = dtypes.get(name)?;
|
||||
|
||||
assert_eq!(shape.len(), 2, "Expected 2D tensor for {name}");
|
||||
|
||||
let f32_data = convert_to_f32(data, dtype);
|
||||
let shape_arr: [usize; 2] = [shape[0], shape[1]];
|
||||
let tensor_data = TensorData::new(f32_data, Shape::new(shape_arr));
|
||||
Some(Tensor::from_data(tensor_data, device))
|
||||
}
|
||||
|
||||
fn load_tensor_1d<B: Backend>(
|
||||
tensors: &HashMap<String, Vec<u8>>,
|
||||
shapes: &HashMap<String, Vec<usize>>,
|
||||
dtypes: &HashMap<String, String>,
|
||||
name: &str,
|
||||
device: &B::Device,
|
||||
) -> Option<Tensor<B, 1, Float>> {
|
||||
let data = tensors.get(name)?;
|
||||
let shape = shapes.get(name)?;
|
||||
let dtype = dtypes.get(name)?;
|
||||
|
||||
assert_eq!(shape.len(), 1, "Expected 1D tensor for {name}");
|
||||
|
||||
let f32_data = convert_to_f32(data, dtype);
|
||||
let shape_arr: [usize; 1] = [shape[0]];
|
||||
let tensor_data = TensorData::new(f32_data, Shape::new(shape_arr));
|
||||
Some(Tensor::from_data(tensor_data, device))
|
||||
}
|
||||
|
||||
fn convert_to_f32(data: &[u8], dtype: &str) -> Vec<f32> {
|
||||
match dtype {
|
||||
"BF16" => data
|
||||
.chunks_exact(2)
|
||||
.map(|chunk| {
|
||||
let bytes: [u8; 2] = chunk.try_into().unwrap();
|
||||
bf16_to_f32(u16::from_le_bytes(bytes))
|
||||
})
|
||||
.collect(),
|
||||
"F16" => data
|
||||
.chunks_exact(2)
|
||||
.map(|chunk| {
|
||||
let bytes: [u8; 2] = chunk.try_into().unwrap();
|
||||
f16_to_f32(u16::from_le_bytes(bytes))
|
||||
})
|
||||
.collect(),
|
||||
"F32" => data
|
||||
.chunks_exact(4)
|
||||
.map(|chunk| {
|
||||
let bytes: [u8; 4] = chunk.try_into().unwrap();
|
||||
f32::from_le_bytes(bytes)
|
||||
})
|
||||
.collect(),
|
||||
_ => panic!("Unsupported dtype: {dtype}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn bf16_to_f32(bf16: u16) -> f32 {
|
||||
let f32_bits = (bf16 as u32) << 16;
|
||||
f32::from_bits(f32_bits)
|
||||
}
|
||||
|
||||
fn f16_to_f32(f16: u16) -> f32 {
|
||||
let sign = (f16 >> 15) & 1;
|
||||
let exp = (f16 >> 10) & 0x1F;
|
||||
let mant = f16 & 0x3FF;
|
||||
|
||||
if exp == 0 {
|
||||
if mant == 0 {
|
||||
return if sign == 0 { 0.0 } else { -0.0 };
|
||||
}
|
||||
let mut m = mant as u32;
|
||||
let mut e = 0u32;
|
||||
while (m & 0x400) == 0 {
|
||||
m <<= 1;
|
||||
e += 1;
|
||||
}
|
||||
m &= 0x3FF;
|
||||
let f32_exp = 127 - 14 - e;
|
||||
let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((m as u32) << 13);
|
||||
return f32::from_bits(f32_bits);
|
||||
}
|
||||
|
||||
if exp == 0x1F {
|
||||
let f32_bits = ((sign as u32) << 31) | (0xFFu32 << 23) | ((mant as u32) << 13);
|
||||
return f32::from_bits(f32_bits);
|
||||
}
|
||||
|
||||
let f32_exp = (exp as u32) + 127 - 15;
|
||||
let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((mant as u32) << 13);
|
||||
f32::from_bits(f32_bits)
|
||||
}
|
||||
Reference in New Issue
Block a user