BMInf：千元显卡玩转百亿大模型

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

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

由OpenBMB团队与北京智源研究院联合发布的一款低资源大模型推理工具包BMInf，一面世就获得了广大发烧友的喜爱，该项目已在GitHub获得237个Star，并且论文已被ACL 2022收录。下面，我将带领大家重温BMInf。

工具特点

在最低配置为NVIDIA GTX 1060 6G的千元级显卡上便可以进行百亿模型的高效推理。BMInf全名为Big Model Inference，它具有如下特点：

工具特点

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

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

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

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

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

背后技术

项目的官方文档中以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内存中。

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

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

性能测试

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

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

参考文献

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.