Files
MiniCPM5-1B-rust/readme.md
T
macrocc 7e0585421f feat(inference): 添加模型加载和分词器功能
feat(model): 实现 Llama 模型及其解码器层
feat(norm): 添加 RMSNorm 归一化层
feat(rope): 实现 RoPE 位置编码
2026-07-01 10:56:43 +08:00

2.7 KiB
Raw Blame History

MiniCPM5-1B Burn

使用 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. 转换模型

# 默认输出到 model/model.bin
cargo run --release -- convert

# 指定输出路径
cargo run --release -- convert --output my_model/model.bin

2. 运行推理

# 无上限生成(直到遇到 eos token)
cargo run --release -- run --text "你好"

# 限制最大生成 token 数
cargo run --release -- run --text "你好" --max-tokens 100

移植到其他项目

转换后的文件:

model/
└── model.bin    # Burn bincode 格式模型

在其他项目中加载:

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 GBF32

许可证

MIT