OA0
OA0 是一个探索 AI 的社区
现在注册
已注册用户请  登录
OA0  ›  代码  ›  OpenBMB BMInf?

OpenBMB BMInf?

 
  pebble ·  2026-09-02 11:00:18 · 26 次点击  · 0 条评论  

BMInf

**大型模型高效推理工具**

项目概述安装指南快速开始简体中文

文档 github 版本


更新动态

  • 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 demo。
  • 2021/09/24 在中关村论坛(AI 与多学科协同创新论坛)上公开发布了 BMInf。

注意: BMInf-1 的 README 文档位于 old_docs 目录中。CPM-1/2 和 EVA 的示例即将发布。

项目概述

BMInf(Big Model Inference)是针对大规模预训练语言模型(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.pybenchmark/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.cn
- 微博: http://weibo.cn/OpenBMB
- Twitter: https://twitter.com/OpenBMB

开源许可

本工具包基于 Apache 2.0 许可证发布。

参考引用

  1. CPM-2: Large-scale Cost-efficient Pre-trained Language Models. Zhengyan Zhang, Yuxian Gu, Xu Han, Shengqi Chen, Chaojun Xiao, Zhenbo Sun, Yuan Yao, Fanchao Qi, Jian Guan, Pei Ke, Yanzheng Cai, Guoyang Zeng, Zhixing Tan, Zhiyuan Liu, Minlie Huang, Wentao Han, Yang Liu, Xiaoyan Zhu, Maosong Sun.
  2. CPM: A Large-scale Generative Chinese Pre-trained Language Model. Zhengyan Zhang, Xu Han, Hao Zhou, Pei Ke, Yuxian Gu, Deming Ye, Yujia Qin, Yusheng Su, Haozhe Ji, Jian Guan, Fanchao Qi, Xiaozhi Wang, Yanan Zheng, Guoyang Zeng, Huanqi Cao, Shengqi Chen, Daixuan Li, Zhenbo Sun, Zhiyuan Liu, Minlie Huang, Wentao Han, Jie Tang, Juanzi Li, Xiaoyan Zhu, Maosong Sun.
  3. EVA: An Open-Domain Chinese Dialogue System with Large-Scale Generative Pre-Training. Hao Zhou, Pei Ke, Zheng Zhang, Yuxian Gu, Yinhe Zheng, Chujie Zheng, Yida Wang, Chen Henry Wu, Hao Sun, Xiaocong Yang, Bosi Wen, Xiaoyan Zhu, Minlie Huang, Jie Tang.
  4. Language Models are Unsupervised Multitask Learners. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.
26 次点击  ∙  0 人收藏  
登录后收藏  
0 条回复
关于 ·  帮助 ·  PING ·  隐私 ·  条款   
OA0 - Omni AI 0 一个探索 AI 的社区
沪ICP备2024103595号-2
耗时 19 ms
Developed with Cursor