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-gpucuda 等)。

最小示例

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

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