PyTorch搭建大型模型“加速包”

日期:2023-12-04 19:09:12 / 人气:129

PyTorch搭建大型模型“加速包”,比1000行代码快10倍!“西风自凹寺来。
量子比特|微信官方账号QbitAI
PyTorch团队将大型模型的推理速度提高了10倍。
而且只需要不到1000行纯原生PyTorch代码!
该项目名为GPT-fast,对加速效果的感知如下:
平稳,非常平稳!
重点是团队直接发布了代码和详细的“教程”。还是简笔画版,特别好理解。
开发小组成员@贺瑞斯说:
我们不把它当成一个库,一个框架,而是希望大家以它为例,根据自己的需求“复制粘贴”。
网友们直接炸开了锅,英伟达AI科学家Jim Fan评论道:
这是自Andrej Karpathy发布minGPT以来最好的教程报告之一!
开源世界需要更多像minGPT和GPT这样的项目——快!
那么,GPT-fast如何加快大型模型的速度呢?
开箱大型号“加速包”
一般来说,使用这些方法:
Torch.compile:专门为PyTorch模型设计的编译器,可以提高模型的运行效率。
GPU量化:通过降低计算的精度来加快模型的运算速度。
预测解码:使用较小的模型预测较大模型的输出,从而加快大型语言模型的运行速度。
张量并行:通过将模型的运算分布在多个硬件设备上,加快处理速度。
让我们逐一展开。
一开始开发团队用的是简单的PyTorch,效果并不好(25.5 tok/s):
他们查看了痕迹后发现,一个原因是由于CPU占用过多导致推理性能受限。
那么如何解决呢?
你可以想象一个场景,GPU是一个巨大的工厂(拥有大量可用的计算能力),CPU是一辆大车,来回“补给”工厂。
很多情况下,CPU无法足够快地“喂”GPU。
所以开发团队建议给GPU更多的工作量,或者一次性给它更大的“块”任务去处理。
在推理过程中要做到这一点,可以引入torch.compile。
Torch.compile可以在模型中捕获一个更大的区域,编译成一个单独的编译区域。特别是在“减少开销”模式下运行时,有效减少了CPU的开销。
效果立竿见影,性能直接提升4倍,从25 tok/s提升到107 Tok/s;
接下来开发团队想进一步提升速度,却遇到了内存带宽的瓶颈。
开发团队计算了模型的带宽利用率,结果已经达到72%:
也就是说,进一步提速的空间可能是有限的。
重新检查上面的等式,团队发现虽然我们实际上不能改变模型参数或GPU的内存带宽(至少不花更多的钱),但我们可以改变用于存储每个参数的字节数。
这意味着,虽然不可能改变模型的大小或升级硬件来提高性能,但可以通过减少存储模型参数所需的数据量来提高效率。
通常可以通过量化技术来实现,即减少表示每个参数所需的比特数。
于是,开发团队引入了下一项技术——int 8量化。
用int8进行权重量化,降低了内存负荷,进一步提升了性能(157.4 Tok/s);
使用量化后,仍然存在一个问题:要生成100个token,模型权重必须被加载(或调用)100次。频繁加载模型权重也会导致效率低下。
乍一看似乎无解,因为自回归生成模型中存在严格的序列依赖性。
但开发团队指出,这种严格的序列依赖可以通过使用推测解码来打破。
再举个例子,想象一下一个高级工程师Verity,他在技术决策上总是正确的,但是写代码的速度相对较慢。
同时还有一个初级工程师德雷克,他和Verity相反,不擅长技术决策,但是写代码更快,成本更低。
那么如何利用不同人的优势来提高整体效率呢?
方法简单。让Drake先写代码,在这个过程中做技术决策。接下来把代码交给Verity审核,让Drake重做错误。
在Transformer模型的推理中,大规模验证模型是Verity角色,而Drake是更小的草稿模型,可以更快地生成文本。
开发团队使用草案模型生成8个token,然后使用验证模型并行处理它们,丢弃不匹配的部分。
结果,串行依赖被打破,速度再次提高。
值得一提的是,推测解码不会改变输出的质量。只要使用草案模型生成令牌+验证这些令牌所需的时间少于单独生成这些令牌所需的时间,这种方法就是有效的。
而且使用原生PyTorch实现这项技术其实非常简单,整个实现过程只需要50行左右的原生PyTorch代码。
由于AMD还支持Triton和torch.compile后端,所以之前在Nvidia GPU上应用的所有优化也可以在AMD GPU上重新应用。
开发团队观察到int8的量子化从22 tok/s加速到102 tok/s:
后来开发团队用int4量化进一步提高了速度,但是模型精度下降了。
因此,使用块量化和GPTQ来降低权重。
最后在保证精度的前提下,速度提升到202.1 tok/s:
结合以上技术,实现244.7 tok/s的更高速度:
到目前为止,R&D团队一直在单个GPU上加速。但其实很多情况下,可以使用多个GPU。
使用多个GPU可以增加内存带宽,从而提高模型的整体性能。
在选择并行处理策略时,需要将一个令牌的处理过程划分在多个设备上,因此需要张量并行。
PyTorch还提供了张量并行的底层工具,可以与torch.compile结合使用。
开发团队还透露,用于表达张量并行性的更高级API也正在开发中。
然而,即使没有更高级别的API,也很容易添加张量并行,它可以在150行代码中实现,而无需对模型进行任何更改。
前面提到的所有优化都可以和张量并行结合。结合这些优化,我们可以以55令牌/秒的速度为Llama-70B提供int8量化。
最后对结果进行总结,忽略量化。仅用766行代码(model.py 244行,generate.py 371行,tp.py 151行)实现了快速推理、推测解码和张量并行。
对于Llama-7B,用compile+int4进行量化和推测解码的速度达到241 Tok/s,对于Llama-70B,通过加入张量并行,可以达到80 tok/s。

作者:安信娱乐平台官网




现在致电 8888910 OR 查看更多联系方式 →

COPYRIGHT 安信注册登录娱乐中国站 版权所有