千元显卡玩转百亿大模型， 清华推出工具包BMInf让模型推理轻而易举

大家好，我是对白。

今天给大家推荐一下我校计算机系NLP实验室和智源团队联合发布的一款低资源大模型推理工具包BMInf，在最低配置为NVIDIA GTX 1060 6G的千元级显卡上便可以进行百亿模型的高效推理。

大家快来一起试用一下吧~

最近在工业界与学术界，最热门的方向莫过于预训练语言模型。而具有百亿乃至千亿参数的大规模预训练语言模型，更是业界与学术界发力的热点。

但现在大模型的应用却有着较高的门槛，排队申请或需要付费的API、较长的模型响应速度、推理所需要的较为昂贵的算力资源……种种因素都影响着大模型的快速应用与落地。对于普通的研究者与开发者来说，大模型可以说是看得见，却很难摸得着。

近日，清华大学自然语言处理实验室团队与北京智源研究院联合发布了一款低资源大模型推理工具包BMInf，在最低配置为NVIDIA GTX 1060 6G的千元级显卡上便可以进行百亿模型的高效推理。BMInf全名为Big Model Inference，它具有如下特点：

1、硬件友好。BMInf最低支持在NVIDIA GTX 1060单卡运行百亿大模型，使用更好的GPU会有更好的运行性能。在显存支持进行大模型推理的情况下（如V100或A100显卡），BMInf的实现较现有PyTorch版本仍有较大性能提升。

2、开源共享。模型参数开源共享，用户在本地即可部署运行，无需访问或申请API。

3、能力全面。支持生成模型CPM1 [1]、通用模型CPM2 [2]、对话模型EVA [3]，模型能力覆盖文本补全、文本生成与对话。

4、模型升级。基于持续学习推出百亿模型新升级CPM2.1，文本生成能力大幅提高

5、应用便捷。基于工具包可以快速开发大模型相关下游应用。

项目地址：https://github.com/OpenBMB/BMInf

对于大模型使用感兴趣的同学，快来试用体验吧！

使用方式

工具包的使用安装较为便捷，官方README里提供了pip、源代码与Docker三种安装方式。

图1 工具包三种安装方式

安装完成后，便可以通过运行examples里面的脚本来体验大模型的快速推理了。运行过程中工具包首先会自动下载模型，之后便能够为事先预设的句子生成预测的结果。

图2 工具包运行演示

（样例可以说是紧贴时事了，不过模型为特定日门票的定价倒是便宜了一些）

除此之外，工具包还提供了基于三种模型的网页版展示Demo，用户可以方便地通过本地运行服务器来可视化地体验大模型的推理结果。目前Demo提供了三项功能，分别是文本填空、智能对话与故事生成，通过网站用户可以非常快捷地体验大模型的趣味应用。

图3 文本填空功能随机样例

图4 以“哈利波特第一次去霍格沃茨魔法学校上学，地址在北京通州”开头生成故事

背后技术

项目的官方文档中以CPM2-11B（悟道·文源）模型为例介绍了工具包背后所用到的技术。

具体来说，CPM2-11B是一个巨大的预训练语言模型，它包含了110亿参数，总共的模型文件大小达到了22GB。对于一块普通的机械硬盘来说，光是读取这样大小的模型就已经会花费数分钟的时间，更不用说将这样庞大的数据放进内存和GPU中进行后续的运算。

在工具包的设计过程中，开发者预设了GTX 1060 6GB显存和16GB内存这样的最低配置要求。开发者使用了一系列的技术让一个大规模的预训练语言模型可以运行在一台普通的电脑上，并进行相关的推理运算。

模型压缩

如果不做任何特殊处理，运行一个22GB的模型需要一块显存大小至少为22GB的GPU。满足这样条件的GPU通常是很昂贵的（例如 V100 32GB, A100, RTX 3090，市场价均超过2万元），为了能让模型在更小显存的GPU上运行，开发者在保留模型原有结构的基础上，将模型中占比最大的线性层参数（占比99%）从16比特浮点数转换为了int8格式。为了让压缩后的模型更贴近于原来的效果，开发者在将参数转换后进行了几千次迭代的微调让模型适应新的参数精度，微调后的模型基本上已经达到了和原始模型相近的能力。在具体的PPL指标中，压缩后的模型相比于压缩前只相差了5~10左右。

显存调度

在使用了模型压缩技术后，原本大小22GB的模型被压缩到了11GB，对于NVIDIA旗舰级别GPU来说（如GTX 1080Ti, RTX 2080Ti），11GB显存已经是可以摸到的门槛了，但是考虑到在推理过程中还需要使用一些额外的空间来存储中间变量，这样的显存容量依然不够。另外，能够拥有这样旗舰级别显卡的个人用户仍然较少，像GTX 1060 6G这样甜点级的GPU有着更加广泛的市场占有率。

要让一个11GB大小的模型运行在显存容量只有6GB的GPU上，开发者使用了显存和内存的优化与调度技术。在运行过程中将需要用于进行运算的模型参数提前准备好并放置在显存中，而对于暂时不需要用到的部分，则放置在CPU内存中。

图5 模型中不同层参数的调度示意

在具体的实现过程中，开发者使用了两个独立的内存池，交替地用于计算和加载参数，并通过将计算的时间和参数加载的时间重叠来提升整体的推理效率。通过调整每次加载和运算的层数，可以将大模型适配到不同显存容量的GPU。实验表明，将更多的模型层驻留在GPU中通常可以带来更好的效率。在整个模型推理的过程中，瓶颈往往在于内存和显存的拷贝，这也是工具包之后会优化的方向。

性能测试

开发者在不同NVIDIA显卡上进行了百亿模型解码的速度测试。即使在1060显卡上，模型仍然能够达到1.6个字符每秒的解码速度，达到了类似于打字机一样的运行效率。拥有较强性能显卡的用户仍然会获得更快的解码速度。在支持直接运行大模型推理的显卡上（如单卡V100或A100），BMInf的解码速度仍明显优于现有的基于PyTorch的实现（注：该PyTorch实现仍使用公开版本模型参数）。

图6 工具包在不同显卡上的解码速度测试

参考文献

1. 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. 2020.

2. 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. 2021.

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. 2021.

你好，我是对白，硕士毕业于清华，现大厂算法工程师，拿过八家大厂的SSP级以上offer。

本科时独立创业五年，成立两家公司，并拿过总计三百多万元融资(已到账)，项目入选南京321高层次创业人才引进计划。创业做过无人机、机器人和互联网教育，保研后选择退出。

我每周至少更新一篇原创，分享自己的算法技术、创业心得和人生感悟。