refactor(项目结构) 重构为 workspace 多成员项目
- 将项目拆分为三个 crate:minicpm-core(核心模型)、minicpm-convert(转换功能)、minicpm-inference(推理功能) - 添加两个示例:minimal-inference(最小推理)和 convert(模型转换) - 转换后自动拷贝 config.json 和 tokenizer.json 到 model 目录 - 更新 README 说明 workspace 结构和使用方式
This commit is contained in:
@@ -1,107 +1,186 @@
|
||||
# MiniCPM5-1B Burn
|
||||
# MiniCPM5-1B-rust
|
||||
|
||||
使用 [Burn](https://github.com/tracel-ai/burn) 框架从零实现的 MiniCPM5-1B 模型。
|
||||
基于 [Burn](https://burn.dev/) 深度学习框架实现的 MiniCPM5-1B 大语言模型推理库,纯 Rust 编写,支持 GPU 加速。
|
||||
|
||||
## 功能
|
||||
## 模型架构参数
|
||||
|
||||
- **模型转换**:将 safetensors 格式转换为 Burn 原生 bincode 格式
|
||||
- **模型推理**:使用 WGPU 后端运行 MiniCPM5-1B 进行文本生成,支持无上限输出
|
||||
- **独立输出**:转换后的模型文件放在 `model/` 目录,方便移植到其他项目
|
||||
| 参数 | 值 |
|
||||
|------|-----|
|
||||
| 层数 | 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 结构
|
||||
|
||||
- 24 层 Transformer
|
||||
- GQA:16 个查询头,2 个 KV 头
|
||||
- RoPE 旋转位置编码(rope_theta = 5000000)
|
||||
- RMSNorm 归一化
|
||||
- SiLU 激活函数(FFN: gate + up + down)
|
||||
- KV Cache 优化
|
||||
项目采用 Cargo workspace 多成员结构,包含三个 crate:
|
||||
|
||||
## 使用
|
||||
### minicpm-core
|
||||
|
||||
### 准备模型
|
||||
核心模型定义,包含:
|
||||
|
||||
将 MiniCPM5-1B 模型放到 `MiniCPM5-1B/` 目录:
|
||||
```
|
||||
MiniCPM5-1B/
|
||||
├── config.json
|
||||
├── tokenizer.json
|
||||
└── model-00000-of-00001.safetensors
|
||||
```
|
||||
- `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. 转换模型
|
||||
|
||||
```bash
|
||||
# 默认输出到 model/model.bin
|
||||
cargo run --release -- convert
|
||||
首先将 HuggingFace 格式的 MiniCPM5-1B 模型转换为 Burn MPK 格式:
|
||||
|
||||
# 指定输出路径
|
||||
cargo run --release -- convert --output my_model/model.bin
|
||||
```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(())
|
||||
}
|
||||
```
|
||||
|
||||
### 2. 运行推理
|
||||
转换完成后,`MiniCPM5-1B-burn/` 目录下会生成 `model.mp` 文件。
|
||||
|
||||
```bash
|
||||
# 无上限生成(直到遇到 eos token)
|
||||
cargo run --release -- run --text "你好"
|
||||
### 2. 推理使用
|
||||
|
||||
# 限制最大生成 token 数
|
||||
cargo run --release -- run --text "你好" --max-tokens 100
|
||||
使用 `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(())
|
||||
}
|
||||
```
|
||||
|
||||
## 移植到其他项目
|
||||
|
||||
转换后的文件:
|
||||
```
|
||||
model/
|
||||
└── model.bin # Burn bincode 格式模型
|
||||
### 依赖配置
|
||||
|
||||
在你的 `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 burn::record::{BinFileRecorder, FullPrecisionSettings};
|
||||
use minicpm_burn_lib::{LlamaConfig, LlamaForCausalLM};
|
||||
use minicpm_inference::{MiniCPM, GenerationConfig};
|
||||
use burn::backend::Wgpu;
|
||||
|
||||
let config = LlamaConfig::from_json("config.json")?;
|
||||
let model = LlamaForCausalLM::<Backend>::new(config, &device);
|
||||
let model = model.load_file("model/model.bin", &BinFileRecorder::<FullPrecisionSettings>::new())?;
|
||||
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/` 目录需要包含以下文件:
|
||||
|
||||
```
|
||||
src/
|
||||
├── lib.rs # 库入口
|
||||
├── main.rs # CLI 工具
|
||||
├── config.rs # 模型配置
|
||||
├── exporter.rs # 模型转换
|
||||
├── model/ # 模型架构
|
||||
│ ├── mod.rs
|
||||
│ ├── attention.rs # GQA 注意力 + KV Cache
|
||||
│ ├── decoder.rs # 解码器层
|
||||
│ ├── ffn.rs # 前馈网络
|
||||
│ ├── model.rs # LlamaForCausalLM
|
||||
│ ├── norm.rs # RMSNorm
|
||||
│ └── rope.rs # RoPE 位置编码
|
||||
└── inference/ # 推理工具
|
||||
├── mod.rs
|
||||
├── generation.rs # 文本生成(greedy + KV Cache)
|
||||
├── loader.rs # safetensors 加载器
|
||||
└── tokenizer.rs # 分词器封装
|
||||
MiniCPM5-1B/
|
||||
├── config.json # 模型配置文件
|
||||
├── model.safetensors # safetensors 格式权重
|
||||
└── tokenizer.json # tokenizer 配置
|
||||
```
|
||||
|
||||
## 依赖
|
||||
转换后 Burn 模型目录结构:
|
||||
|
||||
- `burn` 0.21 - 深度学习框架(WGPU 后端)
|
||||
- `burn-store` - safetensors 支持
|
||||
- `tokenizers` - HuggingFace 分词器
|
||||
- `clap` - CLI 参数解析
|
||||
```
|
||||
MiniCPM5-1B-burn/
|
||||
└── model.mp # Burn MPK 格式权重
|
||||
```
|
||||
|
||||
## 性能
|
||||
## 性能说明
|
||||
|
||||
- WGPU Vulkan 后端:约 10 tokens/s
|
||||
- 模型大小:约 4.0 GB(F32)
|
||||
| Backend | 设备 | 速度 | 显存占用 |
|
||||
|---------|------|------|----------|
|
||||
| WGPU (Vulkan) | RTX 4060 | ~10 tokens/s | ~4 GB (F32) |
|
||||
|
||||
> 以上数据仅供参考,实际性能因硬件配置、生成长度等因素而异。
|
||||
|
||||
## 许可证
|
||||
|
||||
MIT
|
||||
MIT
|
||||
|
||||
Reference in New Issue
Block a user