BMInf
cupy,并支持 PyTorch 反向传播。generate 接口,并新增了 CPM 2.1 demo。注意: 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,所有依赖项均可通过安装过程自动完成安装。
使用 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.cn
- 微博: http://weibo.cn/OpenBMB
- Twitter: https://twitter.com/OpenBMB
本工具包基于 Apache 2.0 许可证发布。