chore(代码变更): 1. 清理无用代码 2. 更新依赖库 3. 优化项目结构 4. 修复小型bug 5. 增加注释以提高可读性
This commit is contained in:
@@ -0,0 +1,31 @@
|
||||
use burn::module::Module;
|
||||
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
|
||||
use burn::tensor::backend::Backend;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::loader::{load_into_model, TensorStore};
|
||||
use minicpm_core::config::LlamaConfig;
|
||||
use minicpm_core::model::LlamaForCausalLM;
|
||||
|
||||
/// 导出全精度模型（f32）
|
||||
pub fn export<B: Backend>(
|
||||
store: &TensorStore,
|
||||
config: &LlamaConfig,
|
||||
output_dir: &Path,
|
||||
device: &B::Device,
|
||||
) -> anyhow::Result<()> {
|
||||
println!("创建模型结构...");
|
||||
let model = LlamaForCausalLM::<B>::new(config.clone(), device);
|
||||
|
||||
println!("加载 safetensors 权重...");
|
||||
let model = load_into_model(model, store, device);
|
||||
|
||||
println!("保存为 MPK 格式（全精度）...");
|
||||
std::fs::create_dir_all(output_dir)?;
|
||||
let output_path = output_dir.join("model");
|
||||
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
|
||||
model.save_file(&output_path, &recorder)?;
|
||||
|
||||
println!("全精度模型已导出到: {:?}", output_dir);
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
pub mod full;
|
||||
|
||||
/// 导出格式
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum Format {
|
||||
/// 全精度 f32
|
||||
Full,
|
||||
}
|
||||
|
||||
impl Format {
|
||||
pub fn name(&self) -> &'static str {
|
||||
match self {
|
||||
Format::Full => "全精度",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn dir_name(&self) -> &'static str {
|
||||
match self {
|
||||
Format::Full => "model",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 导出任务配置
|
||||
pub struct ExportTask {
|
||||
pub format: Format,
|
||||
pub output_dir: String,
|
||||
}
|
||||
|
||||
impl ExportTask {
|
||||
pub fn new(format: Format) -> Self {
|
||||
Self {
|
||||
output_dir: format.dir_name().to_string(),
|
||||
format,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_dir(format: Format, output_dir: impl Into<String>) -> Self {
|
||||
Self {
|
||||
output_dir: output_dir.into(),
|
||||
format,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,511 +1,46 @@
|
||||
use minicpm_core::config::LlamaConfig;
|
||||
use minicpm_core::model::LlamaForCausalLM;
|
||||
use burn::module::{Module, Param};
|
||||
use burn::nn::{EmbeddingRecord, LinearRecord};
|
||||
use burn::record::{FullPrecisionSettings, HalfPrecisionSettings, NamedMpkFileRecorder};
|
||||
pub mod format;
|
||||
pub mod loader;
|
||||
pub mod utils;
|
||||
|
||||
pub use format::{ExportTask, Format};
|
||||
pub use loader::TensorStore;
|
||||
|
||||
use burn::tensor::backend::Backend;
|
||||
use burn::tensor::{Float, Shape, Tensor, TensorData};
|
||||
use memmap2::Mmap;
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
pub fn export_model<B: Backend>(
|
||||
use minicpm_core::config::LlamaConfig;
|
||||
|
||||
/// 执行导出任务
|
||||
pub fn run_export<B: Backend>(
|
||||
safetensors_path: &Path,
|
||||
config_path: &Path,
|
||||
tokenizer_path: &Path,
|
||||
output_dir: &Path,
|
||||
tasks: &[ExportTask],
|
||||
device: &B::Device,
|
||||
) -> anyhow::Result<()> {
|
||||
println!("开始转换 MiniCPM 模型为 Burn 格式...");
|
||||
println!("开始转换 MiniCPM 模型...");
|
||||
println!("源文件: {:?}", safetensors_path);
|
||||
|
||||
println!("加载配置文件: {:?}", config_path);
|
||||
println!("加载配置文件...");
|
||||
let config = LlamaConfig::from_json(config_path.to_str().unwrap())?;
|
||||
|
||||
println!("创建模型结构...");
|
||||
let model = LlamaForCausalLM::<B>::new(config, device);
|
||||
println!("解析 safetensors 文件...");
|
||||
let store = TensorStore::from_file(safetensors_path)?;
|
||||
|
||||
println!("加载 safetensors 权重...");
|
||||
let model = load_safetensors(model, safetensors_path, device)?;
|
||||
for task in tasks {
|
||||
println!("\n--- 导出 {} ---", task.format.name());
|
||||
let output_dir = Path::new(&task.output_dir);
|
||||
|
||||
println!("保存为 MPK 格式...");
|
||||
std::fs::create_dir_all(output_dir)?;
|
||||
let output_path = output_dir.join("model");
|
||||
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
|
||||
model.save_file(&output_path, &recorder)?;
|
||||
match task.format {
|
||||
Format::Full => {
|
||||
format::full::export::<B>(&store, &config, output_dir, device)?;
|
||||
}
|
||||
}
|
||||
|
||||
println!("拷贝配置文件和 tokenizer...");
|
||||
std::fs::copy(config_path, output_dir.join("config.json"))?;
|
||||
std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
|
||||
std::fs::copy(config_path, output_dir.join("config.json"))?;
|
||||
std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
|
||||
}
|
||||
|
||||
println!("模型已成功导出到: {:?}", output_dir);
|
||||
println!("\n所有导出任务完成！");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 导出模型为半精度（f16），输出文件约为全精度的一半大小。
|
||||
pub fn export_model_half<B: Backend>(
|
||||
safetensors_path: &Path,
|
||||
config_path: &Path,
|
||||
tokenizer_path: &Path,
|
||||
output_dir: &Path,
|
||||
device: &B::Device,
|
||||
) -> anyhow::Result<()> {
|
||||
println!("开始转换 MiniCPM 模型为 Burn 格式（半精度 f16）...");
|
||||
|
||||
println!("加载配置文件: {:?}", config_path);
|
||||
let config = LlamaConfig::from_json(config_path.to_str().unwrap())?;
|
||||
|
||||
println!("创建模型结构...");
|
||||
let model = LlamaForCausalLM::<B>::new(config, device);
|
||||
|
||||
println!("加载 safetensors 权重...");
|
||||
let model = load_safetensors(model, safetensors_path, device)?;
|
||||
|
||||
println!("保存为半精度 MPK 格式...");
|
||||
std::fs::create_dir_all(output_dir)?;
|
||||
let output_path = output_dir.join("model");
|
||||
let recorder = NamedMpkFileRecorder::<HalfPrecisionSettings>::new();
|
||||
model.save_file(&output_path, &recorder)?;
|
||||
|
||||
println!("拷贝配置文件和 tokenizer...");
|
||||
std::fs::copy(config_path, output_dir.join("config.json"))?;
|
||||
std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
|
||||
|
||||
println!("半精度模型已成功导出到: {:?}", output_dir);
|
||||
Ok(())
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
// ==================== Q8 量化导出 ====================
|
||||
|
||||
/// Q8 量化单个 f32 切片，返回 (scale, i8_data)
|
||||
fn quantize_q8(f32_data: &[f32]) -> (f32, Vec<i8>) {
|
||||
let max_abs = f32_data
|
||||
.iter()
|
||||
.map(|&x| x.abs())
|
||||
.fold(0.0f32, |a, b| a.max(b));
|
||||
let scale = if max_abs > 0.0 { max_abs / 127.0 } else { 1.0 };
|
||||
let q8: Vec<i8> = f32_data
|
||||
.iter()
|
||||
.map(|&x| (x / scale).round().clamp(-127.0, 127.0) as i8)
|
||||
.collect();
|
||||
(scale, q8)
|
||||
}
|
||||
|
||||
/// Q8 导出元数据中的 tensor 条目
|
||||
#[derive(serde::Serialize, serde::Deserialize)]
|
||||
struct Q8TensorMeta {
|
||||
shape: Vec<usize>,
|
||||
scale: f32,
|
||||
offset: u64,
|
||||
numel: usize,
|
||||
}
|
||||
|
||||
/// Q8 导出元数据
|
||||
#[derive(serde::Serialize, serde::Deserialize)]
|
||||
struct Q8Metadata {
|
||||
tensors: std::collections::HashMap<String, Q8TensorMeta>,
|
||||
}
|
||||
|
||||
/// 导出模型为 Q8 量化格式（INT8），输出 model.q8.json + model.q8.bin。
|
||||
/// 2D 权重（Linear/Embedding）量化为 INT8；1D norm 权重存为 F16。
|
||||
pub fn export_model_q8(
|
||||
safetensors_path: &Path,
|
||||
config_path: &Path,
|
||||
tokenizer_path: &Path,
|
||||
output_dir: &Path,
|
||||
) -> anyhow::Result<()> {
|
||||
println!("开始转换 MiniCPM 模型为 Q8 量化格式...");
|
||||
|
||||
println!("解析 safetensors 文件: {:?}", safetensors_path);
|
||||
let file = std::fs::File::open(safetensors_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 data_offset = 8 + header_len;
|
||||
|
||||
// 收集 tensor 信息
|
||||
let mut tensor_infos: Vec<(String, Vec<usize>, String, usize, usize)> = Vec::new();
|
||||
if let Some(obj) = header.as_object() {
|
||||
for (name, info) in obj {
|
||||
if name == "__metadata__" {
|
||||
continue;
|
||||
}
|
||||
if let Some(info_obj) = info.as_object() {
|
||||
let mut shape = Vec::new();
|
||||
if let Some(s) = info_obj.get("shape") {
|
||||
if let Some(arr) = s.as_array() {
|
||||
shape = arr.iter().map(|v| v.as_u64().unwrap() as usize).collect();
|
||||
}
|
||||
}
|
||||
let mut dtype = String::new();
|
||||
if let Some(d) = info_obj.get("dtype") {
|
||||
dtype = d.as_str().unwrap_or("").to_string();
|
||||
}
|
||||
let mut start = 0;
|
||||
let mut end = 0;
|
||||
if let Some(offsets) = info_obj.get("data_offsets") {
|
||||
if let Some(arr) = offsets.as_array() {
|
||||
start = arr[0].as_u64().unwrap() as usize;
|
||||
end = arr[1].as_u64().unwrap() as usize;
|
||||
}
|
||||
}
|
||||
tensor_infos.push((name.clone(), shape, dtype, start, end));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 定义需要量化的 2D 权重名称（与 load_safetensors 中的列表对应）
|
||||
let mut q8_tensors: Vec<String> = Vec::new();
|
||||
for i in 0..32 {
|
||||
let p = format!("model.layers.{i}");
|
||||
q8_tensors.push(format!("{p}.self_attn.q_proj.weight"));
|
||||
q8_tensors.push(format!("{p}.self_attn.k_proj.weight"));
|
||||
q8_tensors.push(format!("{p}.self_attn.v_proj.weight"));
|
||||
q8_tensors.push(format!("{p}.self_attn.o_proj.weight"));
|
||||
q8_tensors.push(format!("{p}.mlp.gate_proj.weight"));
|
||||
q8_tensors.push(format!("{p}.mlp.up_proj.weight"));
|
||||
q8_tensors.push(format!("{p}.mlp.down_proj.weight"));
|
||||
}
|
||||
q8_tensors.push("model.embed_tokens.weight".to_string());
|
||||
q8_tensors.push("lm_head.weight".to_string());
|
||||
|
||||
let q8_set: std::collections::HashSet<String> = q8_tensors.into_iter().collect();
|
||||
|
||||
std::fs::create_dir_all(output_dir)?;
|
||||
let json_path = output_dir.join("model.q8.json");
|
||||
let bin_path = output_dir.join("model.q8.bin");
|
||||
|
||||
let mut metadata: std::collections::HashMap<String, Q8TensorMeta> = std::collections::HashMap::new();
|
||||
let mut bin_file = std::fs::File::create(&bin_path)?;
|
||||
let mut current_offset: u64 = 0;
|
||||
|
||||
// 先处理 2D 权重：量化为 INT8
|
||||
for (name, shape, dtype, start, end) in &tensor_infos {
|
||||
if shape.len() != 2 {
|
||||
continue;
|
||||
}
|
||||
if !q8_set.contains(name.as_str()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
println!("量化 2D 权重: {name} shape={shape:?} dtype={dtype}");
|
||||
let raw = &mmap[data_offset + *start..data_offset + *end];
|
||||
let f32_data = convert_to_f32(raw, dtype);
|
||||
let numel = f32_data.len();
|
||||
let (scale, q8_data) = quantize_q8(&f32_data);
|
||||
|
||||
// 写入 INT8 数据
|
||||
let bytes: &[u8] = unsafe {
|
||||
std::slice::from_raw_parts(q8_data.as_ptr() as *const u8, q8_data.len())
|
||||
};
|
||||
use std::io::Write;
|
||||
bin_file.write_all(bytes)?;
|
||||
|
||||
metadata.insert(
|
||||
name.clone(),
|
||||
Q8TensorMeta {
|
||||
shape: shape.clone(),
|
||||
scale,
|
||||
offset: current_offset,
|
||||
numel,
|
||||
},
|
||||
);
|
||||
current_offset += q8_data.len() as u64;
|
||||
}
|
||||
|
||||
// 再处理 1D norm 权重：存为 F16
|
||||
for (name, shape, dtype, start, end) in &tensor_infos {
|
||||
if shape.len() != 1 {
|
||||
continue;
|
||||
}
|
||||
let raw = &mmap[data_offset + *start..data_offset + *end];
|
||||
let f32_data = convert_to_f32(raw, dtype);
|
||||
let numel = f32_data.len();
|
||||
|
||||
// 转 F16 存储
|
||||
let mut f16_bytes = Vec::with_capacity(numel * 2);
|
||||
for &v in &f32_data {
|
||||
let f16_bits = f32_to_f16_bits(v);
|
||||
f16_bytes.extend_from_slice(&f16_bits.to_le_bytes());
|
||||
}
|
||||
|
||||
use std::io::Write;
|
||||
bin_file.write_all(&f16_bytes)?;
|
||||
|
||||
// 用负数 scale 表示 F16 类型（加载时 scale < 0 表示 F16）
|
||||
metadata.insert(
|
||||
name.clone(),
|
||||
Q8TensorMeta {
|
||||
shape: shape.clone(),
|
||||
scale: -1.0, // 标记为 F16
|
||||
offset: current_offset,
|
||||
numel,
|
||||
},
|
||||
);
|
||||
current_offset += f16_bytes.len() as u64;
|
||||
}
|
||||
|
||||
// 写 JSON 元数据
|
||||
let meta = Q8Metadata { tensors: metadata };
|
||||
let json_str = serde_json::to_string_pretty(&meta)?;
|
||||
std::fs::write(&json_path, json_str)?;
|
||||
|
||||
println!("拷贝配置文件和 tokenizer...");
|
||||
std::fs::copy(config_path, output_dir.join("config.json"))?;
|
||||
std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
|
||||
|
||||
println!("Q8 量化模型已导出到: {:?}", output_dir);
|
||||
println!(" - model.q8.json (元数据)",);
|
||||
println!(" - model.q8.bin (INT8 + F16 权重)",);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// f32 → f16 位表示（返回 u16）
|
||||
fn f32_to_f16_bits(v: f32) -> u16 {
|
||||
let bits = v.to_bits();
|
||||
let sign = bits >> 31;
|
||||
let exp = (bits >> 23) & 0xFF;
|
||||
let mant = bits & 0x7FFFFF;
|
||||
|
||||
let res: u32 = if exp == 0xFF {
|
||||
// NaN / Inf
|
||||
let f16_exp = 0x1F;
|
||||
let f16_mant = if mant == 0 { 0 } else { 0x3FF };
|
||||
(sign << 15) | (f16_exp << 10) | f16_mant
|
||||
} else if exp < 103 {
|
||||
// 下溢，返回 0
|
||||
0
|
||||
} else {
|
||||
let new_exp = (exp as i32) - 127 + 15;
|
||||
if new_exp >= 31 {
|
||||
// 上溢
|
||||
(sign << 15) | (0x1F << 10)
|
||||
} else if new_exp <= 0 {
|
||||
// 次正规数
|
||||
let mant_new = (mant | 0x800000) >> (113 - new_exp);
|
||||
(sign << 15) | (mant_new & 0x3FF)
|
||||
} else {
|
||||
let f16_mant = mant >> 13;
|
||||
(sign << 15) | ((new_exp as u32) << 10) | f16_mant
|
||||
}
|
||||
};
|
||||
res as u16
|
||||
}
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
pub mod safetensors;
|
||||
|
||||
pub use safetensors::{load_into_model, TensorStore};
|
||||
@@ -0,0 +1,185 @@
|
||||
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::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::utils::convert_to_f32;
|
||||
use minicpm_core::model::LlamaForCausalLM;
|
||||
|
||||
/// 从 safetensors 文件解析出的张量数据
|
||||
pub struct TensorStore {
|
||||
tensors: HashMap<String, Vec<u8>>,
|
||||
shapes: HashMap<String, Vec<usize>>,
|
||||
dtypes: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl TensorStore {
|
||||
/// 从 safetensors 文件加载所有张量
|
||||
pub fn from_file(path: &Path) -> anyhow::Result<Self> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
tensors,
|
||||
shapes,
|
||||
dtypes,
|
||||
})
|
||||
}
|
||||
|
||||
/// 加载 2D 张量
|
||||
pub fn load_2d<B: Backend>(&self, name: &str, device: &B::Device) -> Option<Tensor<B, 2, Float>> {
|
||||
let data = self.tensors.get(name)?;
|
||||
let shape = self.shapes.get(name)?;
|
||||
let dtype = self.dtypes.get(name)?;
|
||||
|
||||
assert_eq!(shape.len(), 2, "期望 2D 张量: {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))
|
||||
}
|
||||
|
||||
/// 加载 1D 张量
|
||||
pub fn load_1d<B: Backend>(&self, name: &str, device: &B::Device) -> Option<Tensor<B, 1, Float>> {
|
||||
let data = self.tensors.get(name)?;
|
||||
let shape = self.shapes.get(name)?;
|
||||
let dtype = self.dtypes.get(name)?;
|
||||
|
||||
assert_eq!(shape.len(), 1, "期望 1D 张量: {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))
|
||||
}
|
||||
|
||||
/// 获取原始字节数据
|
||||
pub fn raw(&self, name: &str) -> Option<&[u8]> {
|
||||
self.tensors.get(name).map(|v| v.as_slice())
|
||||
}
|
||||
|
||||
pub fn shape(&self, name: &str) -> Option<&Vec<usize>> {
|
||||
self.shapes.get(name)
|
||||
}
|
||||
|
||||
pub fn dtype(&self, name: &str) -> Option<&String> {
|
||||
self.dtypes.get(name)
|
||||
}
|
||||
|
||||
pub fn names(&self) -> impl Iterator<Item = &String> {
|
||||
self.tensors.keys()
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 safetensors 权重加载到模型中
|
||||
pub fn load_into_model<B: Backend>(
|
||||
model: LlamaForCausalLM<B>,
|
||||
store: &TensorStore,
|
||||
device: &B::Device,
|
||||
) -> LlamaForCausalLM<B> {
|
||||
let mut model = model;
|
||||
|
||||
if let Some(w) = store.load_2d("model.embed_tokens.weight", device) {
|
||||
let record = EmbeddingRecord {
|
||||
weight: Param::from_tensor(w),
|
||||
};
|
||||
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}");
|
||||
|
||||
load_linear(&mut model, store, device, i, &format!("{prefix}.self_attn.q_proj.weight"), |m| &mut m.model.layers[i].self_attn.q_proj);
|
||||
load_linear(&mut model, store, device, i, &format!("{prefix}.self_attn.k_proj.weight"), |m| &mut m.model.layers[i].self_attn.k_proj);
|
||||
load_linear(&mut model, store, device, i, &format!("{prefix}.self_attn.v_proj.weight"), |m| &mut m.model.layers[i].self_attn.v_proj);
|
||||
load_linear(&mut model, store, device, i, &format!("{prefix}.self_attn.o_proj.weight"), |m| &mut m.model.layers[i].self_attn.o_proj);
|
||||
load_linear(&mut model, store, device, i, &format!("{prefix}.mlp.gate_proj.weight"), |m| &mut m.model.layers[i].mlp.gate_proj);
|
||||
load_linear(&mut model, store, device, i, &format!("{prefix}.mlp.up_proj.weight"), |m| &mut m.model.layers[i].mlp.up_proj);
|
||||
load_linear(&mut model, store, device, i, &format!("{prefix}.mlp.down_proj.weight"), |m| &mut m.model.layers[i].mlp.down_proj);
|
||||
|
||||
if let Some(w) = store.load_1d(&format!("{prefix}.input_layernorm.weight"), device) {
|
||||
model.model.layers[i].input_layernorm.weight = Param::from_tensor(w);
|
||||
}
|
||||
if let Some(w) = store.load_1d(&format!("{prefix}.post_attention_layernorm.weight"), device) {
|
||||
model.model.layers[i].post_attention_layernorm.weight = Param::from_tensor(w);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(w) = store.load_1d("model.norm.weight", device) {
|
||||
model.model.norm.weight = Param::from_tensor(w);
|
||||
}
|
||||
|
||||
if let Some(w) = store.load_2d("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);
|
||||
}
|
||||
|
||||
model
|
||||
}
|
||||
|
||||
fn load_linear<B: Backend, F>(
|
||||
model: &mut LlamaForCausalLM<B>,
|
||||
store: &TensorStore,
|
||||
device: &B::Device,
|
||||
_layer_idx: usize,
|
||||
name: &str,
|
||||
get_field: F,
|
||||
) where
|
||||
F: FnOnce(&mut LlamaForCausalLM<B>) -> &mut burn::nn::Linear<B>,
|
||||
{
|
||||
if let Some(w) = store.load_2d(name, device) {
|
||||
let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
|
||||
let field = get_field(model);
|
||||
*field = field.clone().load_record(record);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
/// BF16 → F32
|
||||
pub fn bf16_to_f32(bf16: u16) -> f32 {
|
||||
let f32_bits = (bf16 as u32) << 16;
|
||||
f32::from_bits(f32_bits)
|
||||
}
|
||||
|
||||
/// F16 → F32
|
||||
pub 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)
|
||||
}
|
||||
|
||||
/// F32 → F16 位表示
|
||||
pub fn f32_to_f16_bits(v: f32) -> u16 {
|
||||
let bits = v.to_bits();
|
||||
let sign = bits >> 31;
|
||||
let exp = (bits >> 23) & 0xFF;
|
||||
let mant = bits & 0x7FFFFF;
|
||||
|
||||
let res: u32 = if exp == 0xFF {
|
||||
let f16_exp = 0x1F;
|
||||
let f16_mant = if mant == 0 { 0 } else { 0x3FF };
|
||||
(sign << 15) | (f16_exp << 10) | f16_mant
|
||||
} else if exp < 103 {
|
||||
0
|
||||
} else {
|
||||
let new_exp = (exp as i32) - 127 + 15;
|
||||
if new_exp >= 31 {
|
||||
(sign << 15) | (0x1F << 10)
|
||||
} else if new_exp <= 0 {
|
||||
let mant_new = (mant | 0x800000) >> (113 - new_exp);
|
||||
(sign << 15) | (mant_new & 0x3FF)
|
||||
} else {
|
||||
let f16_mant = mant >> 13;
|
||||
(sign << 15) | ((new_exp as u32) << 10) | f16_mant
|
||||
}
|
||||
};
|
||||
res as u16
|
||||
}
|
||||
|
||||
/// 将原始字节按 dtype 转换为 f32 向量
|
||||
pub fn convert_to_f32(data: &[u8], dtype: &str) -> Vec<f32> {
|
||||
match dtype {
|
||||
"BF16" => data
|
||||
.chunks_exact(2)
|
||||
.map(|chunk| bf16_to_f32(u16::from_le_bytes(chunk.try_into().unwrap())))
|
||||
.collect(),
|
||||
"F16" => data
|
||||
.chunks_exact(2)
|
||||
.map(|chunk| f16_to_f32(u16::from_le_bytes(chunk.try_into().unwrap())))
|
||||
.collect(),
|
||||
"F32" => data
|
||||
.chunks_exact(4)
|
||||
.map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
|
||||
.collect(),
|
||||
_ => panic!("不支持的 dtype: {dtype}"),
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,10 @@ version = "0.1.2"
|
||||
edition = "2021"
|
||||
publish = ["gitea"]
|
||||
|
||||
[features]
|
||||
default = ["avx2"]
|
||||
avx2 = []
|
||||
|
||||
[dependencies]
|
||||
burn = { version = "0.21", default-features = false, features = ["std"] }
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
|
||||
@@ -129,7 +129,6 @@ impl<B: Backend> Attention<B> {
|
||||
let q = rope.apply(q, offset);
|
||||
let k = rope.apply(k, offset);
|
||||
|
||||
// 合并 KV Cache
|
||||
let (k_full, v_full) = match cache {
|
||||
Some(cache) => {
|
||||
let k_full = Tensor::cat(vec![cache.k.clone(), k], 2);
|
||||
|
||||
@@ -4,15 +4,15 @@ use burn::tensor::{Float, Tensor};
|
||||
|
||||
use super::attention::{Attention, KVCache};
|
||||
use super::ffn::FeedForward;
|
||||
use super::norm::RmsNorm;
|
||||
use super::norm::RMSNorm;
|
||||
use super::rope::RoPE;
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct DecoderLayer<B: Backend> {
|
||||
pub self_attn: Attention<B>,
|
||||
pub mlp: FeedForward<B>,
|
||||
pub input_layernorm: RmsNorm<B>,
|
||||
pub post_attention_layernorm: RmsNorm<B>,
|
||||
pub input_layernorm: RMSNorm<B>,
|
||||
pub post_attention_layernorm: RMSNorm<B>,
|
||||
}
|
||||
|
||||
impl<B: Backend> DecoderLayer<B> {
|
||||
@@ -27,8 +27,8 @@ impl<B: Backend> DecoderLayer<B> {
|
||||
) -> Self {
|
||||
let self_attn = Attention::new(hidden_size, n_heads, n_kv_heads, head_dim, device);
|
||||
let mlp = FeedForward::new(hidden_size, intermediate_size, device);
|
||||
let input_layernorm = RmsNorm::new(hidden_size, rms_norm_eps, device);
|
||||
let post_attention_layernorm = RmsNorm::new(hidden_size, rms_norm_eps, device);
|
||||
let input_layernorm = RMSNorm::new(hidden_size, rms_norm_eps, device);
|
||||
let post_attention_layernorm = RMSNorm::new(hidden_size, rms_norm_eps, device);
|
||||
|
||||
Self {
|
||||
self_attn,
|
||||
|
||||
@@ -6,7 +6,7 @@ use burn::tensor::{Float, Int, Tensor};
|
||||
use super::attention::KVCache;
|
||||
use crate::config::LlamaConfig;
|
||||
use super::decoder::DecoderLayer;
|
||||
use super::norm::RmsNorm;
|
||||
use super::norm::RMSNorm;
|
||||
use super::rope::RoPE;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -26,7 +26,7 @@ impl<B: Backend> LlamaKVCache<B> {
|
||||
pub struct LlamaModel<B: Backend> {
|
||||
pub embed_tokens: Embedding<B>,
|
||||
pub layers: Vec<DecoderLayer<B>>,
|
||||
pub norm: RmsNorm<B>,
|
||||
pub norm: RMSNorm<B>,
|
||||
pub rope: RoPE,
|
||||
}
|
||||
|
||||
@@ -58,7 +58,7 @@ impl<B: Backend> LlamaForCausalLM<B> {
|
||||
));
|
||||
}
|
||||
|
||||
let norm = RmsNorm::new(config.hidden_size, config.rms_norm_eps, device);
|
||||
let norm = RMSNorm::new(config.hidden_size, config.rms_norm_eps, device);
|
||||
|
||||
let lm_head = LinearConfig::new(config.hidden_size, config.vocab_size).init(device);
|
||||
|
||||
|
||||
@@ -4,12 +4,12 @@ use burn::tensor::backend::Backend;
|
||||
use burn::tensor::{Float, Tensor};
|
||||
|
||||
#[derive(Module, Debug)]
|
||||
pub struct RmsNorm<B: Backend> {
|
||||
pub struct RMSNorm<B: Backend> {
|
||||
pub weight: Param<Tensor<B, 1, Float>>,
|
||||
pub eps: f64,
|
||||
}
|
||||
|
||||
impl<B: Backend> RmsNorm<B> {
|
||||
impl<B: Backend> RMSNorm<B> {
|
||||
pub fn new(dim: usize, eps: f64, device: &B::Device) -> Self {
|
||||
let weight = Initializer::Ones.init([dim], device);
|
||||
Self {
|
||||
|
||||
@@ -6,10 +6,9 @@ publish = ["gitea"]
|
||||
|
||||
[dependencies]
|
||||
minicpm-core = { path = "../minicpm-core", version = "0.1.2", registry = "gitea" }
|
||||
burn = { version = "0.21", default-features = false, features = ["std"] }
|
||||
burn = { version = "0.21", default-features = false, features = ["std", "autotune"] }
|
||||
tokenizers = "0.20"
|
||||
rand = "0.8"
|
||||
anyhow = "1.0"
|
||||
memmap2 = "0.9"
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
@@ -5,18 +5,8 @@ use minicpm_core::config::EosTokenId;
|
||||
use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
|
||||
|
||||
use crate::config::GenerationConfig;
|
||||
use crate::sampling::{sample, sample_last};
|
||||
use crate::sampling::{CachedSampler, sample_last};
|
||||
|
||||
/// 流式生成迭代器 —— 用户可以自行遍历处理每个 token
|
||||
///
|
||||
/// 用法示例：
|
||||
/// ```ignore
|
||||
/// let stream = TokenStream::new(model, input_ids, config, eos_token_id, device);
|
||||
/// for token in stream {
|
||||
/// // 自行处理每个 token
|
||||
/// println!("{}", token);
|
||||
/// }
|
||||
/// ```
|
||||
pub struct TokenStream<'a, B: Backend> {
|
||||
model: &'a LlamaForCausalLM<B>,
|
||||
config: &'a GenerationConfig,
|
||||
@@ -28,6 +18,7 @@ pub struct TokenStream<'a, B: Backend> {
|
||||
finished: bool,
|
||||
first_step: bool,
|
||||
input_ids: Vec<u32>,
|
||||
sampler: CachedSampler,
|
||||
}
|
||||
|
||||
impl<'a, B: Backend> TokenStream<'a, B> {
|
||||
@@ -49,6 +40,7 @@ impl<'a, B: Backend> TokenStream<'a, B> {
|
||||
finished: false,
|
||||
first_step: true,
|
||||
input_ids: input_ids.to_vec(),
|
||||
sampler: CachedSampler::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -122,11 +114,12 @@ impl<'a, B: Backend> TokenStream<'a, B> {
|
||||
.slice([0..1, 0..1, 0..vocab_size])
|
||||
.reshape([vocab_size]);
|
||||
|
||||
Some(sample(&next_token_logits, self.config.temperature, self.config.top_p))
|
||||
let data = next_token_logits.to_data();
|
||||
let values: Vec<f32> = data.as_slice().unwrap().to_vec();
|
||||
Some(self.sampler.sample(&values, self.config.temperature, self.config.top_p))
|
||||
}
|
||||
}
|
||||
|
||||
/// 非流式生成 —— 返回所有生成的 token id
|
||||
pub fn generate_tokens<B: Backend>(
|
||||
model: &LlamaForCausalLM<B>,
|
||||
input_ids: &[u32],
|
||||
@@ -136,4 +129,4 @@ pub fn generate_tokens<B: Backend>(
|
||||
) -> Vec<u32> {
|
||||
let stream = TokenStream::new(model, input_ids, config, eos_token_id, device);
|
||||
stream.collect()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,14 +5,14 @@ pub mod sampling;
|
||||
pub mod tokenizer;
|
||||
|
||||
pub use config::GenerationConfig;
|
||||
pub use generator::TokenStream;
|
||||
pub use generator::{TokenStream};
|
||||
pub use minicpm_core::config::{EosTokenId, LlamaConfig};
|
||||
pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
|
||||
pub use tokenizer::TokenizerWrapper;
|
||||
|
||||
use burn::tensor::backend::Backend;
|
||||
|
||||
use loader::{load_model, load_model_q8_from_dir};
|
||||
use loader::load_model;
|
||||
|
||||
pub struct MiniCPM<B: Backend> {
|
||||
model: LlamaForCausalLM<B>,
|
||||
@@ -40,24 +40,6 @@ impl<B: Backend> MiniCPM<B> {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn load_q8(
|
||||
model_dir: &str,
|
||||
config_path: &str,
|
||||
tokenizer_path: &str,
|
||||
device: &B::Device,
|
||||
) -> anyhow::Result<Self> {
|
||||
let config = LlamaConfig::from_json(config_path)?;
|
||||
let tokenizer = TokenizerWrapper::from_file(tokenizer_path)?;
|
||||
let model = load_model_q8_from_dir::<B>(model_dir, &config, device)?;
|
||||
|
||||
Ok(Self {
|
||||
model,
|
||||
config,
|
||||
tokenizer,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn generate(
|
||||
&self,
|
||||
prompt: &str,
|
||||
@@ -111,15 +93,6 @@ impl<B: Backend> MiniCPM<B> {
|
||||
}
|
||||
}
|
||||
|
||||
/// 文本流式迭代器 —— 解码 token 并返回文本片段
|
||||
///
|
||||
/// 用法示例：
|
||||
/// ```ignore
|
||||
/// let mut stream = model.generate_stream("你好", false, &config)?;
|
||||
/// for text in stream {
|
||||
/// print!("{}", text);
|
||||
/// }
|
||||
/// ```
|
||||
pub struct TextStream<'a, B: Backend> {
|
||||
token_stream: TokenStream<'a, B>,
|
||||
tokenizer: &'a TokenizerWrapper,
|
||||
@@ -147,4 +120,4 @@ impl<'a, B: Backend> TextStream<'a, B> {
|
||||
pub fn all_tokens(&self) -> &[u32] {
|
||||
&self.buffer
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
pub mod q8;
|
||||
|
||||
use burn::module::Module;
|
||||
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
|
||||
use burn::tensor::backend::Backend;
|
||||
@@ -7,8 +5,6 @@ use burn::tensor::backend::Backend;
|
||||
use minicpm_core::config::LlamaConfig;
|
||||
use minicpm_core::model::LlamaForCausalLM;
|
||||
|
||||
use self::q8::load_model_q8;
|
||||
|
||||
pub fn load_model<B: Backend>(
|
||||
model_path: &str,
|
||||
config: &LlamaConfig,
|
||||
@@ -19,11 +15,3 @@ pub fn load_model<B: Backend>(
|
||||
let model = model.load_file(model_path, &recorder, device)?;
|
||||
Ok(model)
|
||||
}
|
||||
|
||||
pub fn load_model_q8_from_dir<B: Backend>(
|
||||
model_dir: &str,
|
||||
config: &LlamaConfig,
|
||||
device: &B::Device,
|
||||
) -> anyhow::Result<LlamaForCausalLM<B>> {
|
||||
load_model_q8(model_dir, config, device)
|
||||
}
|
||||
@@ -1,201 +0,0 @@
|
||||
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::collections::HashMap;
|
||||
|
||||
use minicpm_core::config::LlamaConfig;
|
||||
use minicpm_core::model::LlamaForCausalLM;
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct Q8TensorMeta {
|
||||
shape: Vec<usize>,
|
||||
scale: f32,
|
||||
offset: u64,
|
||||
numel: usize,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct Q8Metadata {
|
||||
tensors: HashMap<String, Q8TensorMeta>,
|
||||
}
|
||||
|
||||
pub fn load_model_q8<B: Backend>(
|
||||
model_dir: &str,
|
||||
config: &LlamaConfig,
|
||||
device: &B::Device,
|
||||
) -> anyhow::Result<LlamaForCausalLM<B>> {
|
||||
let json_path = std::path::Path::new(model_dir).join("model.q8.json");
|
||||
let bin_path = std::path::Path::new(model_dir).join("model.q8.bin");
|
||||
|
||||
let json_str = std::fs::read_to_string(&json_path)?;
|
||||
let metadata: Q8Metadata = serde_json::from_str(&json_str)?;
|
||||
|
||||
let file = std::fs::File::open(&bin_path)?;
|
||||
let mmap = unsafe { Mmap::map(&file)? };
|
||||
|
||||
let mut model = LlamaForCausalLM::<B>::new(config.clone(), device);
|
||||
|
||||
if let Some(meta) = metadata.tensors.get("model.embed_tokens.weight") {
|
||||
let w = if meta.scale >= 0.0 {
|
||||
load_q8_tensor_2d::<B>(&mmap, meta, device)
|
||||
} else {
|
||||
load_f16_tensor_2d::<B>(&mmap, meta, device)
|
||||
};
|
||||
let record = EmbeddingRecord { weight: Param::from_tensor(w) };
|
||||
model.model.embed_tokens = model.model.embed_tokens.clone().load_record(record);
|
||||
}
|
||||
|
||||
for i in 0..config.num_hidden_layers {
|
||||
let p = format!("model.layers.{i}");
|
||||
|
||||
if let Some(meta) = metadata.tensors.get(&format!("{p}.self_attn.q_proj.weight")) {
|
||||
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||
let record = LinearRecord { weight: Param::from_tensor(w), 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(meta) = metadata.tensors.get(&format!("{p}.self_attn.k_proj.weight")) {
|
||||
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||
let record = LinearRecord { weight: Param::from_tensor(w), 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(meta) = metadata.tensors.get(&format!("{p}.self_attn.v_proj.weight")) {
|
||||
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||
let record = LinearRecord { weight: Param::from_tensor(w), 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(meta) = metadata.tensors.get(&format!("{p}.self_attn.o_proj.weight")) {
|
||||
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||
let record = LinearRecord { weight: Param::from_tensor(w), 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(meta) = metadata.tensors.get(&format!("{p}.mlp.gate_proj.weight")) {
|
||||
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||
let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
|
||||
model.model.layers[i].mlp.gate_proj = model.model.layers[i].mlp.gate_proj.clone().load_record(record);
|
||||
}
|
||||
if let Some(meta) = metadata.tensors.get(&format!("{p}.mlp.up_proj.weight")) {
|
||||
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||
let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
|
||||
model.model.layers[i].mlp.up_proj = model.model.layers[i].mlp.up_proj.clone().load_record(record);
|
||||
}
|
||||
if let Some(meta) = metadata.tensors.get(&format!("{p}.mlp.down_proj.weight")) {
|
||||
let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
|
||||
let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
|
||||
model.model.layers[i].mlp.down_proj = model.model.layers[i].mlp.down_proj.clone().load_record(record);
|
||||
}
|
||||
|
||||
if let Some(meta) = metadata.tensors.get(&format!("{p}.input_layernorm.weight")) {
|
||||
let w = load_f16_tensor_1d::<B>(&mmap, meta, device);
|
||||
model.model.layers[i].input_layernorm.weight = Param::from_tensor(w);
|
||||
}
|
||||
if let Some(meta) = metadata.tensors.get(&format!("{p}.post_attention_layernorm.weight")) {
|
||||
let w = load_f16_tensor_1d::<B>(&mmap, meta, device);
|
||||
model.model.layers[i].post_attention_layernorm.weight = Param::from_tensor(w);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(meta) = metadata.tensors.get("model.norm.weight") {
|
||||
let w = load_f16_tensor_1d::<B>(&mmap, meta, device);
|
||||
model.model.norm.weight = Param::from_tensor(w);
|
||||
}
|
||||
|
||||
if let Some(meta) = metadata.tensors.get("lm_head.weight") {
|
||||
let w = if meta.scale >= 0.0 {
|
||||
load_q8_tensor_2d::<B>(&mmap, meta, device)
|
||||
} else {
|
||||
load_f16_tensor_2d::<B>(&mmap, meta, device)
|
||||
};
|
||||
let w_t = w.transpose();
|
||||
let record = LinearRecord { weight: Param::from_tensor(w_t), bias: None };
|
||||
model.lm_head = model.lm_head.clone().load_record(record);
|
||||
} else if config.tie_word_embeddings {
|
||||
let embed_weight = model.model.embed_tokens.weight.val();
|
||||
let w_t = embed_weight.transpose();
|
||||
let record = LinearRecord { weight: Param::from_tensor(w_t), bias: None };
|
||||
model.lm_head = model.lm_head.clone().load_record(record);
|
||||
}
|
||||
|
||||
Ok(model)
|
||||
}
|
||||
|
||||
fn load_q8_tensor_2d<B: Backend>(
|
||||
mmap: &Mmap,
|
||||
meta: &Q8TensorMeta,
|
||||
device: &B::Device,
|
||||
) -> Tensor<B, 2, Float> {
|
||||
let bytes = &mmap[meta.offset as usize..meta.offset as usize + meta.numel];
|
||||
let q8: &[i8] = unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const i8, meta.numel) };
|
||||
let f32_data: Vec<f32> = q8.iter().map(|&q| (q as f32) * meta.scale).collect();
|
||||
let shape: [usize; 2] = [meta.shape[0], meta.shape[1]];
|
||||
Tensor::from_data(TensorData::new(f32_data, Shape::new(shape)), device)
|
||||
}
|
||||
|
||||
fn load_f16_tensor_1d<B: Backend>(
|
||||
mmap: &Mmap,
|
||||
meta: &Q8TensorMeta,
|
||||
device: &B::Device,
|
||||
) -> Tensor<B, 1, Float> {
|
||||
let start = meta.offset as usize;
|
||||
let bytes = &mmap[start..start + meta.numel * 2];
|
||||
let f32_data: Vec<f32> = bytes
|
||||
.chunks_exact(2)
|
||||
.map(|c| {
|
||||
let bits = u16::from_le_bytes([c[0], c[1]]);
|
||||
f16_to_f32(bits)
|
||||
})
|
||||
.collect();
|
||||
let shape: [usize; 1] = [meta.shape[0]];
|
||||
Tensor::from_data(TensorData::new(f32_data, Shape::new(shape)), device)
|
||||
}
|
||||
|
||||
fn load_f16_tensor_2d<B: Backend>(
|
||||
mmap: &Mmap,
|
||||
meta: &Q8TensorMeta,
|
||||
device: &B::Device,
|
||||
) -> Tensor<B, 2, Float> {
|
||||
let start = meta.offset as usize;
|
||||
let bytes = &mmap[start..start + meta.numel * 2];
|
||||
let f32_data: Vec<f32> = bytes
|
||||
.chunks_exact(2)
|
||||
.map(|c| {
|
||||
let bits = u16::from_le_bytes([c[0], c[1]]);
|
||||
f16_to_f32(bits)
|
||||
})
|
||||
.collect();
|
||||
let shape: [usize; 2] = [meta.shape[0], meta.shape[1]];
|
||||
Tensor::from_data(TensorData::new(f32_data, Shape::new(shape)), device)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
@@ -1,6 +1,91 @@
|
||||
use burn::tensor::backend::Backend;
|
||||
use burn::tensor::{Float, Tensor};
|
||||
use rand::Rng;
|
||||
use rand::{rngs::ThreadRng, Rng};
|
||||
|
||||
struct Sampler {
|
||||
rng: ThreadRng,
|
||||
}
|
||||
|
||||
impl Sampler {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
rng: rand::thread_rng(),
|
||||
}
|
||||
}
|
||||
|
||||
fn sample(&mut self, logits: &[f32], temperature: f32, top_p: f32) -> u32 {
|
||||
if temperature <= 0.001 {
|
||||
let mut max_idx = 0;
|
||||
let mut max_val = f32::NEG_INFINITY;
|
||||
for (i, &val) in logits.iter().enumerate() {
|
||||
if val > max_val {
|
||||
max_val = val;
|
||||
max_idx = i;
|
||||
}
|
||||
}
|
||||
return max_idx as u32;
|
||||
}
|
||||
|
||||
let max_val = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||
let exp_sum: f32 = logits.iter().map(|&v| (v - max_val).exp()).sum();
|
||||
let mut probs: Vec<f32> = logits.iter().map(|&v| (v - max_val).exp() / exp_sum).collect();
|
||||
|
||||
if top_p < 1.0 {
|
||||
let mut indexed: Vec<(usize, f32)> = probs.iter().enumerate().map(|(i, &p)| (i, p)).collect();
|
||||
indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
|
||||
let mut cumsum = 0.0;
|
||||
let mut cutoff = indexed.len();
|
||||
for (i, &(_, p)) in indexed.iter().enumerate() {
|
||||
cumsum += p;
|
||||
if cumsum > top_p {
|
||||
cutoff = i + 1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let keep: std::collections::HashSet<usize> = indexed[..cutoff].iter().map(|(i, _)| *i).collect();
|
||||
for (i, p) in probs.iter_mut().enumerate() {
|
||||
if !keep.contains(&i) {
|
||||
*p = 0.0;
|
||||
}
|
||||
}
|
||||
|
||||
let sum: f32 = probs.iter().sum();
|
||||
if sum > 0.0 {
|
||||
for p in probs.iter_mut() {
|
||||
*p /= sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let r: f32 = self.rng.gen();
|
||||
let mut cumsum = 0.0;
|
||||
for (i, &p) in probs.iter().enumerate() {
|
||||
cumsum += p;
|
||||
if r < cumsum {
|
||||
return i as u32;
|
||||
}
|
||||
}
|
||||
probs.len() as u32 - 1
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CachedSampler {
|
||||
sampler: std::cell::RefCell<Sampler>,
|
||||
}
|
||||
|
||||
impl CachedSampler {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
sampler: std::cell::RefCell::new(Sampler::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sample(&self, logits: &[f32], temperature: f32, top_p: f32) -> u32 {
|
||||
self.sampler.borrow_mut().sample(logits, temperature, top_p)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sample_last<B: Backend>(logits: &Tensor<B, 3, Float>, temperature: f32, top_p: f32) -> u32 {
|
||||
let shape = logits.shape();
|
||||
@@ -32,11 +117,9 @@ pub fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p:
|
||||
return max_idx as u32;
|
||||
}
|
||||
|
||||
let scaled: Vec<f32> = values.iter().map(|&v| v / temperature).collect();
|
||||
|
||||
let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||
let exp_sum: f32 = scaled.iter().map(|&v| (v - max_val).exp()).sum();
|
||||
let mut probs: Vec<f32> = scaled.iter().map(|&v| (v - max_val).exp() / exp_sum).collect();
|
||||
let max_val = values.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||
let exp_sum: f32 = values.iter().map(|&v| (v - max_val).exp()).sum();
|
||||
let mut probs: Vec<f32> = values.iter().map(|&v| (v - max_val).exp() / exp_sum).collect();
|
||||
|
||||
if top_p < 1.0 {
|
||||
let mut indexed: Vec<(usize, f32)> = probs.iter().enumerate().map(|(i, &p)| (i, p)).collect();
|
||||
@@ -77,4 +160,4 @@ pub fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p:
|
||||
}
|
||||
}
|
||||
probs.len() as u32 - 1
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user