.hd-box .hd-fr

PyTorch 造大模型“加速包”,不到 1000 行代码提速 10 倍!英伟达科学家:minGPT 以来最好的教程式 repo 之一

2023-12-03 20:26量子位(西风)9评

PyTorch 团队让大模型推理速度加快了10 倍。且只用了不到1000 行的纯原生 PyTorch 代码

项目名为GPT-fast,加速效果观感是这样婶儿的:

通畅,属实通畅!

重点是,团队直接放出了代码以及详细“教程”。还是简笔画版的那种,特别好理解。

开发团队成员 @Horace He 表示:

我们不把它看作是库或者框架,更希望大家能把它当成个例子,根据自己的需求“复制粘贴”。

网友直接炸开锅,英伟达 AI 科学家 Jim Fan 评价道:

这是自 Andrej Karpathy 发布的 minGPT 以来最棒的教程式 repo 之一!

开源世界需要更多 minGPT、GPT-Fast 这样的项目!

那么 GPT-fast 究竟是如何给大模型提速的?

开盒大模型“加速包”

总的来说,用到这几种方法:

下面我们来一一展开。

开发团队一开始使用简单的 PyTorch 来实现,但效果不佳(25.5 tok / s):

他们查看跟踪后发现,一个原因是推理性能由于 CPU 过多占用而受限。

那么如何解决呢?

可以想象这样一个场景,GPU 是一个庞大的工厂(拥有大量可用的算力),而 CPU 则是一个小推车,来回为工厂“供货”。

在很多情况下,CPU 无法足够快地“喂”GPU。

因此,开发团队建议给 GPU 更多的工作量,或者说一次性给它更大“块”的任务来处理。

在推理过程中要做到这一点,可以引入torch.compile

torch.compile能够捕获模型中更大的区域,并将其编译成单一的编译区域。特别是当以“reduce-overhead”模式运行时,它非常有效地减少了 CPU 的开销。

效果立竿见影,性能直接提升了 4 倍,从 25 tok / s 提高到 107 tok / s:

接下来,开发团队想进一步提升速度,但遇到了内存带宽瓶颈。

开发团队计算了模型的带宽利用率,结果已经达到了 72%:

也就是说进一步提高速度的空间可能有限。

重新审视上面的方程式,团队发现虽然实际上不能改变模型参数量,也不能改变 GPU 的内存带宽(至少在不花更多钱的情况下),但可以改变存储每个参数所用的字节数

这意味着,虽然无法改变模型的大小或者升级硬件来提高性能,但可以通过减少存储模型参数所需的数据量来提高效率。

通常可以通过量化技术来实现,即减少表示每个参数所需的位数。

由此,开发团队引入了下一个技术 ——int8 量化

采用 int8 权重量化减少了内存负载,进一步提升了性能(157.4 tok / s):

使用量化后还有一个问题:要生成 100 个 token,必须加载(或调用)模型权重 100 次。频繁加载模型权重也会导致效率低下。

乍一看,好像没有什么解决的法子,因为在自回归生成模式中存在着严格的序列依赖关系。

但开发团队指出,通过利用推测性解码可以打破这种严格的序列依赖关系。

再来打个比方,想象有一个资深工程师 Verity,他在技术决策上总是正确,但编写代码的速度相对较慢。

同时,还有一个初级工程师 Drake,和 Verity 相反,不擅长技术决策,但编写代码的速度更快、成本也更低。

那么如何利用不同人的优势来提高整体效率?

方法很简单,先让 Drake 编写代码,并在此过程中做出技术决策。接下来,将代码交给 Verity 进行审查,不对的地方就让 Drake 重做。

在 Transformer 模型推理中,大型的验证模型即为 Verity 角色,Drake 则是一个更小的、能更快生成文本的草稿模型。

开发团队使用草稿模型生成 8 个 token,然后使用验证模型并行处理,丢弃不匹配的部分。

由此一来,打破了串行依赖,再次提高速度。

值得一提的是,推测性解码不会改变输出的质量。只要使用草稿模型生成 token + 验证这些 token 所需的时间少于单独生成这些 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:

到目前为止,研发团队一直都是在单个 GPU 上提速。但其实很多情况下是可以使用多个 GPU 的。

而使用多个 GPU 可以增加内存带宽,从而提高模型的整体性能。

在选择并行处理策略时,需要在多个设备上分割一个 token 的处理过程,所以需要使用张量并行性。

而 PyTorch 也提供了用于张量并行性的底层工具,可以与torch.compile结合使用。

开发团队还透露也正在开发用于表达张量并行性的更高级别的 API。

然而,即使没有更高级别的 API,添加张量并行性也很容易,150 行代码即可实现,且不需要对模型进行任何改变。

之前提到的所有优化都可以与张量并行性相结合。将这些优化结合起来,能够以 55 tokens / s 的速度为 Llama-70B 提供 int8 量化。

最后总结成果,忽略量化,仅用 766 行代码(model.py244 行代码,generate.py371 行代码,tp.py151 行代码),就实现了快速推理、推测性解码和张量并行性。

对于 Llama-7B,使用 compile+int4 量化 + 推测性解码速度达到 241 tok / s。对于 Llama-70B,通过加入张量并行性,达到 80 tok / s。

这些性能都接近或超越了当前 SOTA。

参考链接:

广告声明:文内含有的对外跳转链接(包括不限于超链接、二维码、口令等形式),用于传递更多信息,节省甄选时间,结果仅供参考,IT之家所有文章均包含本声明。

下载IT之家APP,分享赚金币换豪礼
相关文章
大家都在买广告
热门评论
查看更多评论