Compare commits

...

6 Commits

29 changed files with 1108 additions and 2453 deletions
+14
View File
@@ -19,3 +19,17 @@
- `minicpm-inference/src/lib.rs`: 新增 `generate_stream` 函数和 `MiniCPM::generate_stream` 方法,每 token 通过回调 `impl FnMut(&str)` 输出解码文本
- `examples/minimal-inference/src/main.rs`: 改为流式调用,`print!` + `flush` 实时输出
## 升级 0.1.2 重新发布(含流式功能)
- minicpm-core, minicpm-convert, minicpm-inference 全部升级到 0.1.2
- 发布全部成功
## Q8 量化导出和加载
- `minicpm-convert`: 新增 `export_model_q8`2D 权重 per-tensor INT8 量化(scale = max(|f|)/127),1D norm 权重存 F16;输出 `model.q8.json` + `model.q8.bin`
- `minicpm-inference`: 新增 `MiniCPM::load_q8`,加载 Q8 模型并在内存中反量化为 f32 推理;新增依赖 `memmap2``serde``serde_json`
- `examples/convert`: 支持 Q8 导出
- `examples/minimal-inference`: 支持 `--full|--half|--q8` 三种模式
- 量化方式:`q = clamp(round(f / scale), -127, 127)``f' = q * scale`
Generated
+88 -1643
View File
File diff suppressed because it is too large Load Diff
+6 -1
View File
@@ -3,7 +3,12 @@ members = [
"crates/minicpm-core",
"crates/minicpm-convert",
"crates/minicpm-inference",
"examples/minimal-inference",
"examples/convert",
"examples/wgpu-backend",
]
resolver = "2"
[profile.release]
opt-level = 3
lto = true
codegen-units = 1
+4 -3
View File
@@ -1,12 +1,13 @@
[package]
name = "minicpm-convert"
version = "0.1.1"
version = "0.2.0"
edition = "2021"
publish = ["gitea"]
[dependencies]
minicpm-core = { path = "../minicpm-core", version = "0.1.1", registry = "gitea" }
minicpm-core = { path = "../minicpm-core", version = "0.2.0", registry = "gitea" }
burn = { version = "0.21", default-features = false, features = ["std"] }
memmap2 = "0.9"
anyhow = "1.0"
serde_json = "1.0"
serde_json = "1.0"
serde = { version = "1.0", features = ["derive"] }
+31
View File
@@ -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(())
}
+44
View File
@@ -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,
}
}
}
+29 -241
View File
@@ -1,258 +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, 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(())
}
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)
}
+3
View File
@@ -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);
}
}
+84
View File
@@ -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}"),
}
}
+5 -1
View File
@@ -1,9 +1,13 @@
[package]
name = "minicpm-core"
version = "0.1.1"
version = "0.2.0"
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);
+5 -5
View File
@@ -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,
+3 -3
View File
@@ -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);
+2 -2
View File
@@ -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 -4
View File
@@ -1,12 +1,14 @@
[package]
name = "minicpm-inference"
version = "0.1.1"
version = "0.2.0"
edition = "2021"
publish = ["gitea"]
[dependencies]
minicpm-core = { path = "../minicpm-core", version = "0.1.1", registry = "gitea" }
burn = { version = "0.21", default-features = false, features = ["std"] }
minicpm-core = { path = "../minicpm-core", version = "0.2.0", registry = "gitea" }
burn = { version = "0.21", default-features = false, features = ["std", "autotune"] }
tokenizers = "0.20"
rand = "0.8"
anyhow = "1.0"
anyhow = "1.0"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
+16
View File
@@ -0,0 +1,16 @@
#[derive(Debug, Clone)]
pub struct GenerationConfig {
pub max_new_tokens: Option<usize>,
pub temperature: f32,
pub top_p: f32,
}
impl Default for GenerationConfig {
fn default() -> Self {
Self {
max_new_tokens: None,
temperature: 1.0,
top_p: 1.0,
}
}
}
+132
View File
@@ -0,0 +1,132 @@
use burn::tensor::backend::Backend;
use burn::tensor::{Int, Tensor};
use minicpm_core::config::EosTokenId;
use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
use crate::config::GenerationConfig;
use crate::sampling::{CachedSampler, sample_last};
pub struct TokenStream<'a, B: Backend> {
model: &'a LlamaForCausalLM<B>,
config: &'a GenerationConfig,
eos_token_id: &'a EosTokenId,
device: &'a B::Device,
cache: Option<LlamaKVCache<B>>,
last_token: Option<u32>,
count: usize,
finished: bool,
first_step: bool,
input_ids: Vec<u32>,
sampler: CachedSampler,
}
impl<'a, B: Backend> TokenStream<'a, B> {
pub fn new(
model: &'a LlamaForCausalLM<B>,
input_ids: &[u32],
config: &'a GenerationConfig,
eos_token_id: &'a EosTokenId,
device: &'a B::Device,
) -> Self {
Self {
model,
config,
eos_token_id,
device,
cache: None,
last_token: None,
count: 0,
finished: false,
first_step: true,
input_ids: input_ids.to_vec(),
sampler: CachedSampler::new(),
}
}
}
impl<'a, B: Backend> Iterator for TokenStream<'a, B> {
type Item = u32;
fn next(&mut self) -> Option<Self::Item> {
if self.finished {
return None;
}
if let Some(max) = self.config.max_new_tokens {
if self.count >= max {
self.finished = true;
return None;
}
}
let next_token = if self.first_step {
self.first_step = false;
self.step_first()
} else {
self.step_next()
};
match next_token {
Some(token) => {
self.count += 1;
self.last_token = Some(token);
if self.eos_token_id.contains(token) {
self.finished = true;
}
Some(token)
}
None => {
self.finished = true;
None
}
}
}
}
impl<'a, B: Backend> TokenStream<'a, B> {
fn step_first(&mut self) -> Option<u32> {
let input_ints: Vec<i64> = self.input_ids.iter().map(|&x| x as i64).collect();
let input_tensor =
Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), self.device)
.unsqueeze::<2>();
let (logits, new_cache) = self.model.forward_with_cache(input_tensor, self.cache.as_ref());
self.cache = Some(new_cache);
Some(sample_last(&logits, self.config.temperature, self.config.top_p))
}
fn step_next(&mut self) -> Option<u32> {
let last_token = self.last_token?;
let input_ints: Vec<i64> = vec![last_token as i64];
let input_tensor =
Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), self.device)
.unsqueeze::<2>();
let (logits, new_cache) = self.model.forward_with_cache(input_tensor, self.cache.as_ref());
self.cache = Some(new_cache);
let vocab_size = self.model.config.vocab_size;
let next_token_logits = logits
.slice([0..1, 0..1, 0..vocab_size])
.reshape([vocab_size]);
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))
}
}
pub fn generate_tokens<B: Backend>(
model: &LlamaForCausalLM<B>,
input_ids: &[u32],
config: &GenerationConfig,
eos_token_id: &EosTokenId,
device: &B::Device,
) -> Vec<u32> {
let stream = TokenStream::new(model, input_ids, config, eos_token_id, device);
stream.collect()
}
+52 -314
View File
@@ -1,305 +1,18 @@
pub mod config;
pub mod generator;
pub mod loader;
pub mod sampling;
pub mod tokenizer;
pub use config::GenerationConfig;
pub use generator::{TokenStream};
pub use minicpm_core::config::{EosTokenId, LlamaConfig};
pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
pub use tokenizer::TokenizerWrapper;
use burn::module::Module;
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
use burn::tensor::backend::Backend;
use burn::tensor::{Float, Int, Tensor};
use rand::Rng;
use tokenizers::Tokenizer;
pub struct TokenizerWrapper {
tokenizer: Tokenizer,
}
impl TokenizerWrapper {
pub fn from_file(path: &str) -> anyhow::Result<Self> {
let tokenizer = Tokenizer::from_file(path)
.map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e))?;
Ok(Self { tokenizer })
}
pub fn encode(&self, text: &str, add_special_tokens: bool) -> anyhow::Result<Vec<u32>> {
let encoding = self
.tokenizer
.encode(text, add_special_tokens)
.map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?;
Ok(encoding.get_ids().to_vec())
}
pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> anyhow::Result<String> {
let text = self
.tokenizer
.decode(ids, skip_special_tokens)
.map_err(|e| anyhow::anyhow!("Failed to decode: {}", e))?;
Ok(text)
}
pub fn vocab_size(&self) -> usize {
self.tokenizer.get_vocab_size(true)
}
/// 应用 MiniCPM5 chat template
pub fn apply_chat_template(&self, user_msg: &str, enable_thinking: bool) -> anyhow::Result<Vec<u32>> {
let prompt = if enable_thinking {
format!(
"<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n",
user_msg
)
} else {
format!(
"<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
user_msg
)
};
self.encode(&prompt, false)
}
}
pub struct GenerationConfig {
pub max_new_tokens: Option<usize>,
pub temperature: f32,
pub top_p: f32,
}
impl Default for GenerationConfig {
fn default() -> Self {
Self {
max_new_tokens: None,
temperature: 1.0,
top_p: 1.0,
}
}
}
pub fn generate_with_cache<B: Backend>(
model: &LlamaForCausalLM<B>,
input_ids: &[u32],
config: &GenerationConfig,
eos_token_id: &EosTokenId,
device: &B::Device,
) -> Vec<u32> {
let mut output_ids = input_ids.to_vec();
let mut cache: Option<LlamaKVCache<B>> = None;
// 第一步:输入完整 prompt,建立初始 cache
{
let input_ints: Vec<i64> = input_ids.iter().map(|&x| x as i64).collect();
let input_tensor =
Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), device)
.unsqueeze::<2>();
let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref());
cache = Some(new_cache);
let next_token = sample_last(&logits, config.temperature, config.top_p);
output_ids.push(next_token);
if eos_token_id.contains(next_token) {
return output_ids;
}
}
// 后续步骤:每次只输入 1 个 token,使用 cache
let mut count = 1;
loop {
if let Some(max) = config.max_new_tokens {
if count >= max {
break;
}
}
let last_token = output_ids[output_ids.len() - 1];
let input_ints: Vec<i64> = vec![last_token as i64];
let input_tensor =
Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), device)
.unsqueeze::<2>();
let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref());
cache = Some(new_cache);
let vocab_size = model.config.vocab_size;
let next_token_logits = logits
.clone()
.slice([0..1, 0..1, 0..vocab_size])
.reshape([vocab_size]);
let next_token = sample(&next_token_logits, config.temperature, config.top_p);
output_ids.push(next_token);
if eos_token_id.contains(next_token) {
break;
}
count += 1;
}
output_ids
}
/// 流式生成:每生成一个新 token,立即调用 `on_token` 回调输出该 token 的解码文本。
/// 返回完整的 output token IDs。
pub fn generate_stream<B: Backend>(
model: &LlamaForCausalLM<B>,
tokenizer: &TokenizerWrapper,
input_ids: &[u32],
config: &GenerationConfig,
eos_token_id: &EosTokenId,
device: &B::Device,
mut on_token: impl FnMut(&str),
) -> Vec<u32> {
let mut output_ids = input_ids.to_vec();
let mut cache: Option<LlamaKVCache<B>> = None;
// 第一步:完整 prompt 输入
{
let input_ints: Vec<i64> = input_ids.iter().map(|&x| x as i64).collect();
let input_tensor =
Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), device)
.unsqueeze::<2>();
let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref());
cache = Some(new_cache);
let next_token = sample_last(&logits, config.temperature, config.top_p);
output_ids.push(next_token);
// 流式输出第一个 token
if let Ok(text) = tokenizer.decode(&[next_token], true) {
on_token(&text);
}
if eos_token_id.contains(next_token) {
return output_ids;
}
}
// 后续步骤:每步生成一个 token
let mut count = 1;
loop {
if let Some(max) = config.max_new_tokens {
if count >= max {
break;
}
}
let last_token = output_ids[output_ids.len() - 1];
let input_ints: Vec<i64> = vec![last_token as i64];
let input_tensor =
Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), device)
.unsqueeze::<2>();
let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref());
cache = Some(new_cache);
let vocab_size = model.config.vocab_size;
let next_token_logits = logits
.clone()
.slice([0..1, 0..1, 0..vocab_size])
.reshape([vocab_size]);
let next_token = sample(&next_token_logits, config.temperature, config.top_p);
output_ids.push(next_token);
// 流式输出新 token
if let Ok(text) = tokenizer.decode(&[next_token], true) {
on_token(&text);
}
if eos_token_id.contains(next_token) {
break;
}
count += 1;
}
output_ids
}
fn sample_last<B: Backend>(logits: &Tensor<B, 3, Float>, temperature: f32, top_p: f32) -> u32 {
let shape = logits.shape();
let dims: [usize; 3] = shape.dims();
let seq_len = dims[1];
let vocab_size = dims[2];
let last_logits = logits
.clone()
.slice([0..1, seq_len - 1..seq_len, 0..vocab_size])
.reshape([vocab_size]);
sample(&last_logits, temperature, top_p)
}
fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p: f32) -> u32 {
let data = logits.clone().to_data();
let values: Vec<f32> = data.as_slice().unwrap().to_vec();
// temperature = 0 时退化为 greedy
if temperature <= 0.001 {
let mut max_idx = 0;
let mut max_val = f32::NEG_INFINITY;
for (i, &val) in values.iter().enumerate() {
if val > max_val {
max_val = val;
max_idx = i;
}
}
return max_idx as u32;
}
// 应用 temperature
let scaled: Vec<f32> = values.iter().map(|&v| v / temperature).collect();
// softmax
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();
// top_p 过滤
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 mut rng = rand::thread_rng();
let r: f32 = 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
}
use loader::load_model;
pub struct MiniCPM<B: Backend> {
model: LlamaForCausalLM<B>,
@@ -317,10 +30,7 @@ impl<B: Backend> MiniCPM<B> {
) -> anyhow::Result<Self> {
let config = LlamaConfig::from_json(config_path)?;
let tokenizer = TokenizerWrapper::from_file(tokenizer_path)?;
let model = LlamaForCausalLM::<B>::new(config.clone(), device);
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
let model = model.load_file(model_path, &recorder, device)?;
let model = load_model::<B>(model_path, &config, device)?;
Ok(Self {
model,
@@ -337,7 +47,7 @@ impl<B: Backend> MiniCPM<B> {
config: &GenerationConfig,
) -> anyhow::Result<String> {
let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
let output_ids = generate_with_cache(
let output_ids = generator::generate_tokens(
&self.model,
&input_ids,
config,
@@ -348,27 +58,26 @@ impl<B: Backend> MiniCPM<B> {
self.tokenizer.decode(new_ids, true)
}
/// 流式生成:每生成一个 token 立即调用 `on_token` 回调,
/// 参数为该 token 解码后的文本。
pub fn generate_stream(
&self,
pub fn generate_stream<'a>(
&'a self,
prompt: &str,
think: bool,
config: &GenerationConfig,
on_token: impl FnMut(&str),
) -> anyhow::Result<String> {
config: &'a GenerationConfig,
) -> anyhow::Result<TextStream<'a, B>> {
let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
let output_ids = generate_stream(
let token_stream = TokenStream::new(
&self.model,
&self.tokenizer,
&input_ids,
config,
&self.config.eos_token_id,
&self.device,
on_token,
);
let new_ids = &output_ids[input_ids.len()..];
self.tokenizer.decode(new_ids, true)
Ok(TextStream {
token_stream,
tokenizer: &self.tokenizer,
buffer: Vec::new(),
})
}
pub fn config(&self) -> &LlamaConfig {
@@ -383,3 +92,32 @@ impl<B: Backend> MiniCPM<B> {
&self.model
}
}
pub struct TextStream<'a, B: Backend> {
token_stream: TokenStream<'a, B>,
tokenizer: &'a TokenizerWrapper,
buffer: Vec<u32>,
}
impl<'a, B: Backend> Iterator for TextStream<'a, B> {
type Item = String;
fn next(&mut self) -> Option<Self::Item> {
match self.token_stream.next() {
Some(token) => {
self.buffer.push(token);
match self.tokenizer.decode(&[token], true) {
Ok(text) => Some(text),
Err(_) => Some(String::new()),
}
}
None => None,
}
}
}
impl<'a, B: Backend> TextStream<'a, B> {
pub fn all_tokens(&self) -> &[u32] {
&self.buffer
}
}
@@ -0,0 +1,17 @@
use burn::module::Module;
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
use burn::tensor::backend::Backend;
use minicpm_core::config::LlamaConfig;
use minicpm_core::model::LlamaForCausalLM;
pub fn load_model<B: Backend>(
model_path: &str,
config: &LlamaConfig,
device: &B::Device,
) -> anyhow::Result<LlamaForCausalLM<B>> {
let model = LlamaForCausalLM::<B>::new(config.clone(), device);
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
let model = model.load_file(model_path, &recorder, device)?;
Ok(model)
}
+163
View File
@@ -0,0 +1,163 @@
use burn::tensor::backend::Backend;
use burn::tensor::{Float, Tensor};
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();
let dims: [usize; 3] = shape.dims();
let seq_len = dims[1];
let vocab_size = dims[2];
let last_logits = logits
.clone()
.slice([0..1, seq_len - 1..seq_len, 0..vocab_size])
.reshape([vocab_size]);
sample(&last_logits, temperature, top_p)
}
pub fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p: f32) -> u32 {
let data = logits.clone().to_data();
let values: Vec<f32> = data.as_slice().unwrap().to_vec();
if temperature <= 0.001 {
let mut max_idx = 0;
let mut max_val = f32::NEG_INFINITY;
for (i, &val) in values.iter().enumerate() {
if val > max_val {
max_val = val;
max_idx = i;
}
}
return max_idx as u32;
}
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();
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 mut rng = rand::thread_rng();
let r: f32 = 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
}
+48
View File
@@ -0,0 +1,48 @@
use tokenizers::Tokenizer;
pub struct TokenizerWrapper {
tokenizer: Tokenizer,
}
impl TokenizerWrapper {
pub fn from_file(path: &str) -> anyhow::Result<Self> {
let tokenizer = Tokenizer::from_file(path)
.map_err(|e| anyhow::anyhow!("加载 tokenizer 失败: {}", e))?;
Ok(Self { tokenizer })
}
pub fn encode(&self, text: &str, add_special_tokens: bool) -> anyhow::Result<Vec<u32>> {
let encoding = self
.tokenizer
.encode(text, add_special_tokens)
.map_err(|e| anyhow::anyhow!("编码失败: {}", e))?;
Ok(encoding.get_ids().to_vec())
}
pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> anyhow::Result<String> {
let text = self
.tokenizer
.decode(ids, skip_special_tokens)
.map_err(|e| anyhow::anyhow!("解码失败: {}", e))?;
Ok(text)
}
pub fn vocab_size(&self) -> usize {
self.tokenizer.get_vocab_size(true)
}
pub fn apply_chat_template(&self, user_msg: &str, enable_thinking: bool) -> anyhow::Result<Vec<u32>> {
let prompt = if enable_thinking {
format!(
"<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n",
user_msg
)
} else {
format!(
"<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
user_msg
)
};
self.encode(&prompt, false)
}
}
+5 -1
View File
@@ -3,7 +3,11 @@ name = "convert"
version = "0.1.0"
edition = "2021"
[[bin]]
name = "convert"
path = "src/main.rs"
[dependencies]
minicpm-convert = { path = "../../crates/minicpm-convert" }
burn = { version = "0.21", features = ["std", "wgpu"] }
burn = { version = "0.21", default-features = false, features = ["std", "wgpu", "fusion", "autotune"] }
anyhow = "1.0"
+10 -14
View File
@@ -1,24 +1,20 @@
use burn::backend::Wgpu;
use minicpm_convert::export_model;
use burn::backend::{wgpu::WgpuDevice, Wgpu};
use minicpm_convert::{run_export, ExportTask, Format};
use std::path::Path;
fn main() -> anyhow::Result<()> {
let device = Default::default();
let safetensors_path = Path::new("MiniCPM5-1B/model-00000-of-00001.safetensors");
let config_path = Path::new("MiniCPM5-1B/config.json");
let tokenizer_path = Path::new("MiniCPM5-1B/tokenizer.json");
let output_dir = Path::new("model");
export_model::<Wgpu>(
safetensors_path,
config_path,
tokenizer_path,
output_dir,
&device,
)?;
println!("模型转换工具 (Wgpu GPU 加速)");
println!("源文件: {:?}", safetensors_path);
println!();
println!("模型转换完成!");
let device = WgpuDevice::default();
let tasks = vec![ExportTask::new(Format::Full)];
run_export::<Wgpu>(safetensors_path, config_path, tokenizer_path, &tasks, &device)?;
Ok(())
}
}
-54
View File
@@ -1,54 +0,0 @@
// 运行方式:
// 1. 确保已在项目根目录
// 2. 运行:cargo run --release -p minimal-inference
// 3. 确保模型文件存在:
// - model/model.mpk (模型权重)
// - MiniCPM5-1B/config.json (模型配置)
// - MiniCPM5-1B/tokenizer.json (分词器)
use burn::backend::Wgpu;
use minicpm_inference::{GenerationConfig, MiniCPM};
use std::time::Instant;
fn main() -> anyhow::Result<()> {
let device = Default::default();
println!("正在加载模型...");
let start = Instant::now();
let model = MiniCPM::<Wgpu>::load(
"model/model",
"MiniCPM5-1B/config.json",
"MiniCPM5-1B/tokenizer.json",
&device,
)?;
println!("模型加载完成,耗时: {:.2?}", start.elapsed());
let prompt = "我怕人家一说要改前端文件,明天咔嚓甩给我了";
println!("\n用户: {}", prompt);
println!("Assistant: ");
let config = GenerationConfig {
max_new_tokens: Some(200),
temperature: 0.7,
top_p: 0.95,
};
let start = Instant::now();
let _response = model.generate_stream(
prompt,
false,
&config,
|token| {
print!("{}", token);
std::io::Write::flush(&mut std::io::stdout()).ok();
},
)?;
let elapsed = start.elapsed();
println!();
println!("\n生成耗时: {:.2?}", elapsed);
Ok(())
}
@@ -1,9 +1,9 @@
[package]
name = "minimal-inference"
name = "wgpu-backend"
version = "0.1.0"
edition = "2021"
[dependencies]
minicpm-inference = { path = "../../crates/minicpm-inference" }
burn = { version = "0.21", features = ["std", "wgpu"] }
burn = { version = "0.21", default-features = false, features = ["std", "wgpu", "fusion", "autotune"] }
anyhow = "1.0"
+50
View File
@@ -0,0 +1,50 @@
use burn::backend::{wgpu::WgpuDevice, Wgpu};
use minicpm_inference::{GenerationConfig, MiniCPM};
use std::io::Write;
use std::time::Instant;
fn main() -> anyhow::Result<()> {
let args: Vec<String> = std::env::args().collect();
let think = args.iter().any(|s| s == "--think");
println!("========================================");
println!("后端: Wgpu (GPU 加速)");
println!("模型目录: model");
if think {
println!("思考模式: 开启");
} else {
println!("思考模式: 关闭");
}
println!("========================================");
let start = Instant::now();
let device = WgpuDevice::default();
let model = MiniCPM::<Wgpu>::load("model/model", "MiniCPM5-1B/config.json", "MiniCPM5-1B/tokenizer.json", &device)?;
println!("模型加载完成,耗时: {:.2?}", start.elapsed());
let prompt = "需要处理生成一个河南的旅游计划,要求如下:\n1. 旅游计划为三天两夜的行程安排。\n2. 每天的行程安排包括景点、餐饮和住宿建议。\n3. 景点安排要考虑交通便利性和时间合理性。\n请生成详细的旅游计划,包含每天的行程安排、景点介绍、餐饮推荐和住宿建议。\n";
println!("\n用户: {prompt}");
print!("Assistant: ");
std::io::stdout().flush()?;
let config = GenerationConfig {
max_new_tokens: Some(1200),
temperature: 0.7,
top_p: 0.95,
};
let start = Instant::now();
let stream = model.generate_stream(prompt, think, &config)?;
for text in stream {
print!("{text}");
std::io::stdout().flush().ok();
}
let elapsed = start.elapsed();
println!();
println!("\n生成耗时: {:.2?}", elapsed);
Ok(())
}
BIN
View File
Binary file not shown.
+104 -164
View File
@@ -1,186 +1,126 @@
# MiniCPM5-1B-rust
# MiniCPM5-1B-rust (v0.2.0)
基于 [Burn](https://burn.dev/) 深度学习框架实现的 MiniCPM5-1B 大语言模型推理库,纯 Rust 编写,支持 GPU 加速
MiniCPM5-1B 模型的 Rust 推理实现,基于 [Burn](https://github.com/tracel-ai/burn) 框架
## 模型架构参数
## 特性
| 参数 | 值 |
|------|-----|
| 层数 | 24 |
| 注意力机制 | GQA (16 Query / 2 KV) |
| head_dim | 128 |
| hidden_size | 1536 |
| intermediate_size | 8960 |
| vocab_size | 151936 |
| RoPE theta | 5000000 |
| 归一化 | RMSNorm |
| 激活函数 | SiLU (SwiGLU) |
| KV Cache | 支持 |
- 🔥 **Wgpu GPU 加速**:支持 CUDA / Metal / Vulkan / DX12 等多种 GPU 后端
- 🚀 **算子融合**:启用 fusion + autotune 优化,自动选择最优 kernel
- 📝 **流式输出**:实时生成文本,支持逐字输出
- 🧠 **思考模式**:支持开启模型的推理思考过程
- 📦 **纯 Rust**:无 C/C++ 依赖,跨平台编译简单
## Workspace 结构
## 项目结构
项目采用 Cargo workspace 多成员结构,包含三个 crate:
### minicpm-core
核心模型定义,包含:
- `LlamaConfig` — 模型配置(从 `config.json` 加载)
- `LlamaForCausalLM` — 因果语言模型主体
- `LlamaKVCache` — KV 缓存结构
- `EosTokenId` — EOS token 标识(支持单个或多个)
- 各模块实现:Attention、Decoder、FFN、RMSNorm、RoPE
### minicpm-convert
模型格式转换工具,负责将 HuggingFace safetensors 格式的权重转换为 Burn MPK 格式:
- `export_model()` — 完整转换流程(加载 safetensors → 构建模型 → 导出 MPK)
- 支持 BF16 / F16 / F32 精度输入
- 支持 `tie_word_embeddings` 权重共享
### minicpm-inference
推理功能封装,提供高层 API
- `MiniCPM` — 高层封装,整合模型 + tokenizer + 推理逻辑
- `TokenizerWrapper` — tokenizer 封装,支持 MiniCPM5 chat template
- `GenerationConfig` — 生成配置(max_new_tokens、temperature、top_p
- `generate_with_cache()` — 带 KV Cache 的自回归生成
- 支持 temperature 采样、top-p 采样、greedy 解码
```
MiniCPM5-1B-rust/
├── crates/
│ ├── minicpm-core/ # 核心模型实现(Transformer、Attention、FFN
│ ├── minicpm-inference/ # 推理引擎(生成、采样、tokenizer)
│ └── minicpm-convert/ # 模型转换库
└── examples/
├── convert/ # 模型转换示例(safetensors → Burn 格式)
└── wgpu-backend/ # Wgpu GPU 推理示例
```
## 快速开始
### 1. 转换模型
### 1. 准备模型
首先将 HuggingFace 格式的 MiniCPM5-1B 模型转换为 Burn MPK 格式
```rust
use minicpm_convert::export_model;
use burn::backend::Wgpu;
use std::path::Path;
fn main() -> anyhow::Result<()> {
let device = Default::default();
export_model::<Wgpu>(
Path::new("MiniCPM5-1B/model.safetensors"),
Path::new("MiniCPM5-1B/config.json"),
Path::new("MiniCPM5-1B-burn"),
&device,
)?;
Ok(())
}
```
转换完成后,`MiniCPM5-1B-burn/` 目录下会生成 `model.mp` 文件。
### 2. 推理使用
使用 `MiniCPM` 高层 API 进行文本生成:
```rust
use minicpm_inference::{MiniCPM, GenerationConfig};
use burn::backend::Wgpu;
fn main() -> anyhow::Result<()> {
let device = Default::default();
// 加载模型
let model = MiniCPM::<Wgpu>::load(
"MiniCPM5-1B-burn/model",
"MiniCPM5-1B/config.json",
"MiniCPM5-1B/tokenizer.json",
&device,
)?;
// 生成配置
let gen_config = GenerationConfig {
max_new_tokens: Some(512),
temperature: 0.7,
top_p: 0.8,
};
// 生成回答(think = true 启用思考模式)
let response = model.generate("用 Rust 写一个 Hello World", true, &gen_config)?;
println!("{}", response);
Ok(())
}
```
## 移植到其他项目
### 依赖配置
在你的 `Cargo.toml` 中添加:
```toml
[dependencies]
minicpm-inference = { path = "../MiniCPM5-1B-rust/crates/minicpm-inference" }
burn = { version = "0.21", features = ["std", "wgpu"] }
anyhow = "1.0"
```
> 也可以根据需要选择其他 backend(如 `tch-gpu`、`cuda` 等)。
### 最小示例
```rust
use minicpm_inference::{MiniCPM, GenerationConfig};
use burn::backend::Wgpu;
fn main() -> anyhow::Result<()> {
let device = Default::default();
let model = MiniCPM::<Wgpu>::load(
"model",
"config.json",
"tokenizer.json",
&device,
)?;
let config = GenerationConfig {
max_new_tokens: Some(256),
temperature: 1.0,
top_p: 1.0,
};
let output = model.generate("你好", false, &config)?;
println!("{}", output);
Ok(())
}
```
## 模型文件准备
`MiniCPM5-1B/` 目录需要包含以下文件:
将 MiniCPM5-1B safetensors 模型放到 `MiniCPM5-1B/` 目录
```
MiniCPM5-1B/
├── config.json # 模型配置文件
├── model.safetensors # safetensors 格式权重
└── tokenizer.json # tokenizer 配置
├── model-00000-of-00001.safetensors
├── config.json
└── tokenizer.json
```
转换后 Burn 模型目录结构:
### 2. 转换模型
```
MiniCPM5-1B-burn/
└── model.mp # Burn MPK 格式权重
```bash
cargo run --release -p convert
```
## 性能说明
转换完成后会生成 `model/` 目录(全精度 f32,Burn mpk 格式)。
| Backend | 设备 | 速度 | 显存占用 |
|---------|------|------|----------|
| WGPU (Vulkan) | RTX 4060 | ~10 tokens/s | ~4 GB (F32) |
### 3. 运行推理
> 以上数据仅供参考,实际性能因硬件配置、生成长度等因素而异。
```bash
# 普通模式
cargo run --release -p wgpu-backend
## 许可证
# 开启思考模式
cargo run --release -p wgpu-backend -- --think
```
## 使用说明
### wgpu-backend 命令行参数
| 参数 | 说明 |
|------|------|
| `--think` | 开启思考模式,模型会在回答前进行推理思考 |
### GenerationConfig
```rust
GenerationConfig {
max_new_tokens: Some(1200), // 最大生成 token 数
temperature: 0.7, // 温度,越高越随机
top_p: 0.95, // nucleus sampling 参数
}
```
### 作为库使用
```rust
use burn::backend::{wgpu::WgpuDevice, Wgpu};
use minicpm_inference::{GenerationConfig, MiniCPM};
let device = WgpuDevice::default();
let model = MiniCPM::<Wgpu>::load(
"model/model",
"config.json",
"tokenizer.json",
&device,
)?;
let config = GenerationConfig {
max_new_tokens: Some(500),
temperature: 0.7,
top_p: 0.95,
};
// 流式生成
let stream = model.generate_stream("你好", false, &config)?;
for text in stream {
print!("{}", text);
}
```
## 发布构建
```bash
# Release 构建(启用 LTO + 最高优化级别)
cargo build --release
# 发布到 gitea registry
cargo publish -p minicpm-core --registry gitea
cargo publish -p minicpm-inference --registry gitea
cargo publish -p minicpm-convert --registry gitea
```
Release profile 配置:
- `opt-level = 3`:最高优化级别
- `lto = true`:链接时优化
- `codegen-units = 1`:单代码生成单元,最大化优化
## 依赖
- Rust 1.75+
- 支持 Wgpu 的 GPUCUDA / Metal / Vulkan / DX12
## License
MIT