概述 • 安装 • 快速入门 • 简体中文

新增内容

• 2022/07/31 (BMInf 2.0.0) BMInf 现在可以应用于任何基于 Transformer 的模型。
• 2021/12/21 (BMInf 1.0.0) 现在该包不再依赖 cupy，并支持 PyTorch 反向传播。
• 2021/10/18 我们更新了 generate 接口，并新增了 CPM 2.1 示例。
• 2021/09/24 我们在 2021 年中关村论坛（AI 与多学科协同创新论坛）上公开发布了 BMInf。

注意: BMInf-1 的 README 文档可在 old_docs 目录中找到。CPM-1/2 和 EVA 的示例将很快发布。

概述

BMInf（大模型推理）是一个用于大规模预训练语言模型（PLMs）的低资源推理工具包。

BMInf 支持在最低配置的单块 NVIDIA GTX 1060 GPU 上运行参数量超过 100 亿的模型。使用更强大的 GPU 能够获得更好的性能。即使在 GPU 显存足以支持大型模型推理的情况下（如 V100 或 A100），BMInf 仍然比现有的 PyTorch 实现有显著的性能提升。

如果您使用此代码，请引用以下论文：

@inproceedings{han2022bminf,
	title={BMInf: An Efficient Toolkit for Big Model Inference and Tuning},
	author={Han, Xu and Zeng, Guoyang and Zhao, Weilin and Liu, Zhiyuan and Zhang, Zhengyan and Zhou, Jie and Zhang, Jun and Chao, Jia and Sun, Maosong},
	booktitle={Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics: System Demonstrations},
	pages={224--230},
	year={2022}
}

安装

• 通过 pip 安装：pip install bminf

• 从源代码安装：下载软件包并运行 python setup.py install

硬件要求

以下是运行 BMInf 的最低和推荐配置。

	最低配置	推荐配置
内存	16GB	24GB
GPU	NVIDIA GeForce GTX 1060 6GB	NVIDIA Tesla V100 16GB
PCI-E	PCI-E 3.0 x16	PCI-E 3.0 x16

BMInf 支持计算能力为 6.1 或更高的 GPU。请参考 表格 查看您的 GPU 是否受支持。

软件要求

BMInf 需要 CUDA 版本 >= 10.1，所有依赖项均可在安装过程中自动安装。

• python >= 3.6
• torch >= 1.7.1
• cpm_kernels >= 1.0.9

快速入门

使用 bminf.wrapper 自动转换您的模型。

import bminf

# 在 CPU 上初始化您的模型
model = MyModel()

# 使用 wrapper 前加载 state_dict
model.load_state_dict(model_checkpoint)

# 应用 wrapper
with torch.cuda.device(CUDA_DEVICE_INDEX):
    model = bminf.wrapper(model)

如果 bminf.wrapper 不适合您的模型，您可以采用以下方法手动替换。

• 将 torch.nn.ModuleList 替换为 bminf.TransformerBlockList。

module_list = bminf.TransformerBlockList([
	# ...
], [CUDA_DEVICE_INDEX])

• 将 torch.nn.Linear 替换为 bminf.QuantizedLinear。

linear = bminf.QuantizedLinear(torch.nn.Linear(...))

性能

以下是我们测试的不同平台上 CPM2 编码器和解码器的速度。您也可以运行 benchmark/cpm2/encoder.py 和 benchmark/cpm2/decoder.py 来测试您机器上的速度！

实现	GPU	编码器速度 (tokens/s)	解码器速度 (tokens/s)
BMInf	NVIDIA GeForce GTX 1060	718	4.4
BMInf	NVIDIA GeForce GTX 1080Ti	1200	12
BMInf	NVIDIA GeForce GTX 2080Ti	2275	19
BMInf	NVIDIA Tesla V100	2966	20
BMInf	NVIDIA Tesla A100	4365	26
PyTorch	NVIDIA Tesla V100	-	3
PyTorch	NVIDIA Tesla A100	-	7

社区

我们欢迎所有人根据我们的 贡献指南 贡献代码。

您也可以在其他平台找到我们：

• QQ 群：735930538
• 微信公众号：OpenBMB
• 官网：https://www.openbmb.org
• 微博：http://weibo.cn/OpenBMB
• Twitter：https://twitter.com/OpenBMB

许可证

本软件包采用 Apache 2.0 许可证发布。

参考文献

1. CPM-2：大规模低成本预训练语言模型。 张正言、顾宇贤、韩旭、陈圣奇、肖超军、孙振波、姚元、齐凡超、关健、柯培、蔡彦政、曾国阳、谭志兴、刘志远、黄敏列、韩文涛、刘洋、朱晓燕、孙茂松。
2. CPM：大规模生成式中文预训练语言模型。 张正言、韩旭、周浩、柯培、顾宇贤、叶德明、秦雨佳、苏雨生、季浩哲、关健、齐凡超、王小智、郑雅楠、曾国阳、曹焕琪、陈圣奇、李代轩、孙振波、刘志远、黄敏列、韩文涛、唐杰、李娟子、朱晓燕、孙茂松。
3. EVA：具有大规模生成式预训练的开放域中文对话系统。 周浩、柯培、张正言、顾宇贤、郑银鹤、郑楚洁、王一达、吴亨利、孙浩、杨小聪、温博思、朱晓燕、黄敏列、唐杰。
4. 语言模型是无监督的多任务学习者。 亚历克·拉德福德、杰弗里·吴、瑞文·柴尔德、大卫·卢安、达里奥·阿莫迪以及伊利亚·萨茨克韦尔。

BMInf 快速上手指南

BMInf (Big Model Inference) 是一个专为大规模预训练语言模型设计的低资源推理工具包。它支持在单张消费级显卡（如 GTX 1060）上运行参数量超过 100 亿的模型，并在高端显卡上提供比原生 PyTorch 更优的推理性能。

环境准备

在开始之前，请确保您的硬件和软件环境满足以下要求：

硬件要求

• GPU: 计算能力 (Compute Capability) 6.1 或更高（支持 NVIDIA GeForce GTX 1060 及以上型号）。
• 显存: 最低 6GB (推荐 16GB 或以上)。
• 内存: 最低 16GB (推荐 24GB)。
• 接口: PCI-E 3.0 x16。

软件依赖

• 操作系统: Linux (推荐)
• CUDA: 版本 >= 10.1
• Python: 版本 >= 3.6
• PyTorch: 版本 >= 1.7.1
• cpm_kernels: 版本 >= 1.0.9

提示：国内用户安装 PyTorch 时，建议使用清华或中科大镜像源加速下载。

安装步骤

您可以通过 pip 直接安装，这是最便捷的方式。

方式一：使用 pip 安装（推荐）

pip install bminf

国内加速安装命令：

pip install bminf -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装

如果您需要最新开发版或自定义编译：

git clone https://github.com/OpenBMB/BMInf.git
cd BMInf
python setup.py install

基本使用

BMInf 的核心功能是通过 bminf.wrapper 自动包装您的 PyTorch 模型，从而实现低显存推理。

最简单的使用示例

假设您已经定义了一个继承自 torch.nn.Module 的模型类 MyModel 并拥有对应的权重文件。

import torch
import bminf

# 1. 在 CPU 上初始化您的模型
model = MyModel()

# 2. 加载预训练权重 (state_dict)
# 请务必在使用 wrapper 之前完成权重加载
model.load_state_dict(model_checkpoint)

# 3. 应用 BMInf 包装器
# 将模型移至指定 GPU 并进行量化/优化处理
CUDA_DEVICE_INDEX = 0  # 替换为您实际的 GPU 索引
with torch.cuda.device(CUDA_DEVICE_INDEX):
    model = bminf.wrapper(model)

# 4. 现在可以像往常一样使用 model 进行推理
# outputs = model(inputs)

手动替换模式（可选）

如果自动包装器 bminf.wrapper 无法完美适配您的模型结构，您可以手动替换关键组件：

• 替换模块列表：将 torch.nn.ModuleList 替换为 bminf.TransformerBlockList。

  module_list = bminf.TransformerBlockList([
      # ... 您的 Transformer 块
  ], [CUDA_DEVICE_INDEX])
  

• 替换线性层：将 torch.nn.Linear 替换为 bminf.QuantizedLinear。

  linear = bminf.QuantizedLinear(torch.nn.Linear(...))