MiniCPM5-1B-rust (v0.2.0)

MiniCPM5-1B 模型的 Rust 推理实现,基于 Burn 框架。

特性

  • 🔥 Wgpu GPU 加速:支持 CUDA / Metal / Vulkan / DX12 等多种 GPU 后端
  • 🚀 算子融合:启用 fusion + autotune 优化,自动选择最优 kernel
  • 📝 流式输出:实时生成文本,支持逐字输出
  • 🧠 思考模式:支持开启模型的推理思考过程
  • 📦 纯 Rust:无 C/C++ 依赖,跨平台编译简单

项目结构

MiniCPM5-1B-rust/
├── crates/
│   ├── minicpm-core/        # 核心模型实现(Transformer、Attention、FFN
│   ├── minicpm-inference/   # 推理引擎(生成、采样、tokenizer
│   └── minicpm-convert/     # 模型转换库
└── examples/
    ├── convert/             # 模型转换示例(safetensors → Burn 格式)
    └── wgpu-backend/        # Wgpu GPU 推理示例

快速开始

1. 准备模型

将 MiniCPM5-1B safetensors 模型放到 MiniCPM5-1B/ 目录:

MiniCPM5-1B/
├── model-00000-of-00001.safetensors
├── config.json
└── tokenizer.json

2. 转换模型

cargo run --release -p convert

转换完成后会生成 model/ 目录(全精度 f32,Burn mpk 格式)。

3. 运行推理

# 普通模式
cargo run --release -p wgpu-backend

# 开启思考模式
cargo run --release -p wgpu-backend -- --think

使用说明

wgpu-backend 命令行参数

参数 说明
--think 开启思考模式,模型会在回答前进行推理思考

GenerationConfig

GenerationConfig {
    max_new_tokens: Some(1200), // 最大生成 token 数
    temperature: 0.7,            // 温度,越高越随机
    top_p: 0.95,                 // nucleus sampling 参数
}

作为库使用

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);
}

发布构建

# 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

S
Description
No description provided
Readme 2.1 GiB
Languages
Rust 83.6%
Jinja 16.4%