macrocc
69b37ced07
- 将项目拆分为三个 crate：minicpm-core（核心模型）、minicpm-convert（转换功能）、minicpm-inference（推理功能） - 添加两个示例：minimal-inference（最小推理）和 convert（模型转换） - 转换后自动拷贝 config.json 和 tokenizer.json 到 model 目录 - 更新 README 说明 workspace 结构和使用方式
187 lines
4.4 KiB
Markdown
187 lines
4.4 KiB
Markdown
# MiniCPM5-1B-rust
|
||
|
||
基于 [Burn](https://burn.dev/) 深度学习框架实现的 MiniCPM5-1B 大语言模型推理库，纯 Rust 编写，支持 GPU 加速。
|
||
|
||
## 模型架构参数
|
||
|
||
| 参数 | 值 |
|
||
|------|-----|
|
||
| 层数 | 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 | 支持 |
|
||
|
||
## 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 解码
|
||
|
||
## 快速开始
|
||
|
||
### 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/
|
||
├── config.json # 模型配置文件
|
||
├── model.safetensors # safetensors 格式权重
|
||
└── tokenizer.json # tokenizer 配置
|
||
```
|
||
|
||
转换后 Burn 模型目录结构：
|
||
|
||
```
|
||
MiniCPM5-1B-burn/
|
||
└── model.mp # Burn MPK 格式权重
|
||
```
|
||
|
||
## 性能说明
|
||
|
||
| Backend | 设备 | 速度 | 显存占用 |
|
||
|---------|------|------|----------|
|
||
| WGPU (Vulkan) | RTX 4060 | ~10 tokens/s | ~4 GB (F32) |
|
||
|
||
> 以上数据仅供参考，实际性能因硬件配置、生成长度等因素而异。
|
||
|
||
## 许可证
|
||
|
||
MIT
|