107 lines
2.7 KiB
Markdown
107 lines
2.7 KiB
Markdown
|
# MiniCPM5-1B Burn
|
|||
|
|
|
|||
|
|
使用 [Burn](https://github.com/tracel-ai/burn) 框架从零实现的 MiniCPM5-1B 模型。
|
|||
|
|
|
|||
|
|
## 功能
|
|||
|
|
|
|||
|
|
- **模型转换**：将 safetensors 格式转换为 Burn 原生 bincode 格式
|
|||
|
|
- **模型推理**：使用 WGPU 后端运行 MiniCPM5-1B 进行文本生成，支持无上限输出
|
|||
|
|
- **独立输出**：转换后的模型文件放在 `model/` 目录，方便移植到其他项目
|
|||
|
|
|
|||
|
|
## 架构
|
|||
|
|
|
|||
|
|
- 24 层 Transformer
|
|||
|
|
- GQA：16 个查询头，2 个 KV 头
|
|||
|
|
- RoPE 旋转位置编码（rope_theta = 5000000）
|
|||
|
|
- RMSNorm 归一化
|
|||
|
|
- SiLU 激活函数（FFN: gate + up + down）
|
|||
|
|
- KV Cache 优化
|
|||
|
|
|
|||
|
|
## 使用
|
|||
|
|
|
|||
|
|
### 准备模型
|
|||
|
|
|
|||
|
|
将 MiniCPM5-1B 模型放到 `MiniCPM5-1B/` 目录：
|
|||
|
|
```
|
|||
|
|
MiniCPM5-1B/
|
|||
|
|
├── config.json
|
|||
|
|
├── tokenizer.json
|
|||
|
|
└── model-00000-of-00001.safetensors
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
### 1. 转换模型
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
# 默认输出到 model/model.bin
|
|||
|
|
cargo run --release -- convert
|
|||
|
|
|
|||
|
|
# 指定输出路径
|
|||
|
|
cargo run --release -- convert --output my_model/model.bin
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
### 2. 运行推理
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
# 无上限生成（直到遇到 eos token）
|
|||
|
|
cargo run --release -- run --text "你好"
|
|||
|
|
|
|||
|
|
# 限制最大生成 token 数
|
|||
|
|
cargo run --release -- run --text "你好" --max-tokens 100
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
## 移植到其他项目
|
|||
|
|
|
|||
|
|
转换后的文件：
|
|||
|
|
```
|
|||
|
|
model/
|
|||
|
|
└── model.bin # Burn bincode 格式模型
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
在其他项目中加载：
|
|||
|
|
```rust
|
|||
|
|
use burn::record::{BinFileRecorder, FullPrecisionSettings};
|
|||
|
|
use minicpm_burn_lib::{LlamaConfig, LlamaForCausalLM};
|
|||
|
|
|
|||
|
|
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())?;
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
## 项目结构
|
|||
|
|
|
|||
|
|
```
|
|||
|
|
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 # 分词器封装
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
## 依赖
|
|||
|
|
|
|||
|
|
- `burn` 0.21 - 深度学习框架（WGPU 后端）
|
|||
|
|
- `burn-store` - safetensors 支持
|
|||
|
|
- `tokenizers` - HuggingFace 分词器
|
|||
|
|
- `clap` - CLI 参数解析
|
|||
|
|
|
|||
|
|
## 性能
|
|||
|
|
|
|||
|
|
- WGPU Vulkan 后端：约 10 tokens/s
|
|||
|
|
- 模型大小：约 4.0 GB（F32）
|
|||
|
|
|
|||
|
|
## 许可证
|
|||
|
|
|
|||
|
|
MIT
|