PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一
前言 PyTorch团队让大模型推理速度加快了10倍。且只用了不到1000行的纯原生PyTorch代码!
本文转载自量子位
仅用于学术分享,若侵权请联系删除
欢迎关注公众号CV技术指南,专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。
【CV技术指南】CV全栈指导班、基础入门班、论文指导班 全面上线!!
项目名为GPT-fast,加速效果观感是这样婶儿的:
![](https://pic4.zhimg.com/80/v2-f2fbdbed74f9a87a9d855e3995215c1b_1440w.webp)
通畅,属实通畅!
重点是,团队直接放出了代码以及详细“教程”。还是简笔画版的那种,特别好理解。
开发团队成员@Horace He表示:
我们不把它看作是库或者框架,更希望大家能把它当成个例子,根据自己的需求“复制粘贴”。
![](https://pic4.zhimg.com/80/v2-12318130637f168fbd76e4bc9ca1c8f3_1440w.webp)
网友直接炸开锅,英伟达AI科学家Jim Fan评价道:
这是自Andrej Karpathy发布的minGPT以来最棒的教程式repo之一!
开源世界需要更多minGPT、GPT-Fast这样的项目!
![](https://pic2.zhimg.com/80/v2-29610c4add03e66babc114159b8181a1_1440w.webp)
那么GPT-fast究竟是如何给大模型提速的?
开盒大模型“加速包”
总的来说,用到这几种方法:
- Torch.compile:一个专门为PyTorch模型设计的编译器,可以提升模型运行效率。
- GPU量化:通过减少计算的精度来加速模型的运算速度。
- 推测性解码:使用一个较小的模型来预测较大模型的输出,以此加快大语言模型的运算。
- 张量并行性:通过在多个硬件设备上分布模型的运算来加速处理速度。
下面我们来一一展开。
开发团队一开始使用简单的PyTorch来实现,但效果不佳(25.5 tok/s):
![](https://pic1.zhimg.com/80/v2-c2c10fd277817782b1ea9583d631e920_1440w.webp)
他们查看跟踪后发现,一个原因是推理性能由于CPU过多占用而受限。
![](https://pic1.zhimg.com/80/v2-3df80e92475fd20a39b1828e168119ac_1440w.webp)
那么如何解决呢?
可以想象这样一个场景,GPU是一个庞大的工厂(拥有大量可用的算力),而CPU则是一个小推车,来回为工厂“供货”。
在很多情况下,CPU无法足够快地“喂”GPU。
因此,开发团队建议给GPU更多的工作量,或者说一次性给它更大“块”的任务来处理。
![](https://pic3.zhimg.com/80/v2-407cc8b8f9c9632073cb0260065a0102_1440w.webp)
在推理过程中要做到这一点,可以引入torch.compile。
torch.compile能够捕获模型中更大的区域,并将其编译成单一的编译区域。特别是当以“reduce-overhead”模式运行时,它非常有效地减少了CPU的开销。
效果立竿见影,性能直接提升了4倍,从25 tok/s提高到107 tok/s:
![](https://pic1.zhimg.com/80/v2-001c1787213191d03f2f1a65818cf478_1440w.webp)
接下来,开发团队想进一步提升速度,但遇到了内存带宽瓶颈。
开发团队计算了模型的带宽利用率,结果已经达到了72%:
![](https://pic4.zhimg.com/80/v2-aad4776692a0e3798d3dc8af91a29a53_1440w.webp)
也就是说进一步提高速度的空间可能有限。
重新审视上面的方程式,团队发现虽然实际上不能改变模型参数量,也不能改变GPU的内存带宽(至少在不花更多钱的情况下),但可以改变存储每个参数所用的字节数。
![](https://pic4.zhimg.com/80/v2-1a660aedb47f9ac0e986d97288061be3_1440w.webp)
这意味着,虽然无法改变模型的大小或者升级硬件来提高性能,但可以通过减少存储模型参数所需的数据量来提高效率。
通常可以通过量化技术来实现,即减少表示每个参数所需的位数。
![](https://pic1.zhimg.com/80/v2-b7da3b912ebb7562f2f518f80e316634_1440w.webp)
由此,开发团队引入了下一个技术——int8量化。
采用int8权重量化减少了内存负载,进一步提升了性能(157.4 tok/s):
![](https://pic4.zhimg.com/80/v2-a31be1e9dba56df46a6ead2e7f9bd403_1440w.webp)
使用量化后还有一个问题:要生成100个token,必须加载(或调用)模型权重100次。频繁加载模型权重也会导致效率低下。
![](https://pic2.zhimg.com/80/v2-d01d5078c01a55b8613a9d6c022bf161_1440w.webp)
乍一看,好像没有什么解决的法子,因为在自回归生成模式中存在着严格的序列依赖关系。
但开发团队指出,通过利用推测性解码可以打破这种严格的序列依赖关系。
![](https://pic2.zhimg.com/80/v2-e5eabcaf820e1c87acfb94f8e74eb659_1440w.webp)
再来打个比方,想象有一个资深工程师Verity,他在技术决策上总是正确,但编写代码的速度相对较慢。
同时,还有一个初级工程师Drake,和Verity相反,不擅长技术决策,但编写代码的速度更快、成本也更低。
![](https://pic1.zhimg.com/80/v2-47e07d2c4b9cc63856c27e914dd15084_1440w.webp)
那么如何利用不同人的优势来提高整体效率?
方法很简单,先让Drake编写代码,并在此过程中做出技术决策。接下来,将代码交给Verity进行审查,不对的地方就让Drake重做。
![](https://pic1.zhimg.com/80/v2-1c7dabb7f67285b5efaa182332e3e7c8_1440w.webp)
在Transformer模型推理中,大型的验证模型即为Verity角色,Drake则是一个更小的、能更快生成文本的草稿模型。
开发团队使用草稿模型生成8个token,然后使用验证模型并行处理,丢弃不匹配的部分。
由此一来,打破了串行依赖,再次提高速度。
值得一提的是,推测性解码不会改变输出的质量。只要使用草稿模型生成token+验证这些token所需的时间少于单独生成这些token所需的时间,这种方法就是有效的。
而且使用原生PyTorch实现这种技术实际上非常简单,整个实现过程只需要大约50行原生PyTorch代码。
![](https://pic1.zhimg.com/80/v2-f09c0f3b801509bdbce85400674bb718_1440w.webp)
由于AMD也支持Triton和torch.compile后端,因此之前在Nvidia GPU上应用的所有优化也可以在AMD GPU上重新应用。
开发团队观察到int8量化的加速从22 tok/s达到102 tok/s:
![](https://pic1.zhimg.com/80/v2-35dc177024929335fdb444bd4f485620_1440w.webp)
之后开发团队又用了int4量化,进一步提升速度,但模型准确性有所下降。
因此使用了分组量化和GPTQ降低权重大小。
![](https://pic4.zhimg.com/80/v2-bf91ec4d9e28f89c113c0642dd7b80bb_1440w.webp)
最后在保证准确性的前提下,速度提升至202.1 tok/s:
![](https://pic2.zhimg.com/80/v2-52908f31e8979d1b5d6483a46c097bb9_1440w.webp)
将以上技术结合使用,达到更高速度244.7 tok/s:
![](https://pic4.zhimg.com/80/v2-d559453f9b84fbded4f188b50daa4a97_1440w.webp)
到目前为止,研发团队一直都是在单个GPU上提速。但其实很多情况下是可以使用多个GPU的。
而使用多个GPU可以增加内存带宽,从而提高模型的整体性能。
![](https://pic1.zhimg.com/80/v2-c4b297a0531c56a9d049a1395c99b408_1440w.webp)
在选择并行处理策略时,需要在多个设备上分割一个token的处理过程,所以需要使用张量并行性。
而PyTorch也提供了用于张量并行性的底层工具,可以与torch.compile结合使用。
开发团队还透露也正在开发用于表达张量并行性的更高级别的API。
然而,即使没有更高级别的API,添加张量并行性也很容易,150行代码即可实现,且不需要对模型进行任何改变。
![](https://pic1.zhimg.com/80/v2-cdcfc92e6b51fdea8081afe37c66aa64_1440w.webp)
之前提到的所有优化都可以与张量并行性相结合。将这些优化结合起来,能够以55 tokens/s的速度为Llama-70B提供int8量化。
![](https://pic1.zhimg.com/80/v2-c01aa32778aeff2f0b5ec1d6aa25f8e0_1440w.webp)
最后总结成果,忽略量化,仅用766行代码(model.py 244行代码,generate.py 371行代码,tp.py 151行代码),就实现了快速推理、推测性解码和张量并行性。
对于Llama-7B,使用compile+int4量化+推测性解码速度达到241 tok/s。对于Llama-70B,通过加入张量并行性,达到80 tok/s。
这些性能都接近或超越了当前SOTA。
参考链接:
[1]https://pytorch.org/blog/accelerating-generative-ai-2/?utm_content=273712248&utm_medium=social&utm_source=twitter&hss_channel=tw-776585502606721024
[2]https://twitter.com/DrJimFan/status/1730298947376443698
[3]https://twitter.com/cHHillee/status/1730293330213531844
欢迎关注公众号CV技术指南,专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。
【技术文档】《从零搭建pytorch模型教程》122页PDF下载
QQ交流群:470899183。群内有大佬负责解答大家的日常学习、科研、代码问题。
其它文章
LSKA注意力 | 重新思考和设计大卷积核注意力,性能优于ConvNeXt、SWin、RepLKNet以及VAN
CVPR 2023 | TinyMIM:微软亚洲研究院用知识蒸馏改进小型ViT
ICCV2023|涨点神器!目标检测蒸馏学习新方法,浙大、海康威视等提出
ICCV 2023 Oral | 突破性图像融合与分割研究:全时多模态基准与多交互特征学习
HDRUNet | 深圳先进院董超团队提出带降噪与反量化功能的单帧HDR重建算法
南科大提出ORCTrack | 解决DeepSORT等跟踪方法的遮挡问题,即插即用真的很香
1800亿参数,世界顶级开源大模型Falcon官宣!碾压LLaMA 2,性能直逼GPT-4
SAM-Med2D:打破自然图像与医学图像的领域鸿沟,医疗版 SAM 开源了!
GhostSR|针对图像超分的特征冗余,华为诺亚&北大联合提出GhostSR
Meta推出像素级动作追踪模型,简易版在线可玩 | GitHub 1.4K星
CSUNet | 完美缝合Transformer和CNN,性能达到UNet家族的巅峰!