MiniCPM5-1B-rust

基于 Burn 深度学习框架实现的 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 格式：

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 进行文本生成：

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 中添加：

[dependencies]
minicpm-inference = { path = "../MiniCPM5-1B-rust/crates/minicpm-inference" }
burn = { version = "0.21", features = ["std", "wgpu"] }
anyhow = "1.0"

也可以根据需要选择其他 backend（如 tch-gpu、cuda 等）。

最小示例

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