077大型语言模型推理优化
大型语言模型推理优化¶
大型语言模型(LLMs)通过生成具有高度理解和流畅性的文本,将聊天和代码补全等文本生成应用提升到了新的水平。然而,使 LLMs 如此强大的一个因素——它们的体积——也为推理带来了挑战。
基本的推理速度较慢,因为每次生成下一个标记时都需要反复调用 LLM。随着生成的进行,输入序列会逐渐变长,这使得 LLM 的处理时间越来越长。此外,LLMs 拥有数十亿个参数,这些参数在内存中存储和处理非常困难。
本指南将展示如何使用 Transformers 中的优化技术来加速 LLM 推理。
Hugging Face 还提供了 文本生成推理 (TGI),这是一个专门用于部署和提供高度优化的 LLM 推理的库。它包括一些 Transformers 不包含的面向部署的优化功能,例如连续批处理以提高吞吐量和多 GPU 推理的张量并行性。
静态键值缓存和 torch.compile¶
在解码过程中,LLM 会计算每个输入标记的键值对(kv 值)。由于它是自回归的,每次生成的输出都会成为新的输入,因此会重复计算相同的 kv 值。这效率不高,因为它每次都在重新计算相同的 kv 值。
为了优化这一点,可以使用 kv 缓存来存储过去的键和值,而不是每次都重新计算。然而,由于 kv 缓存在每次生成步骤中会增长,并且是动态的,这阻碍了你利用 torch.compile,这是一个强大的优化工具,可以将 PyTorch 代码融合成快速且优化的内核。我们有一整篇关于 kv 缓存的指南 在这里。
静态 kv 缓存解决了这个问题,通过预先分配最大值来预分配 kv 缓存的大小,允许你将其与 torch.compile 结合使用,从而实现最高 4 倍的速度提升。具体速度提升取决于模型大小(较大的模型速度提升较小)和硬件。
目前,只有 Llama 和少数其他模型支持静态 kv 缓存和 torch.compile。请查看 这个 issues 获取实时的模型兼容列表。
根据任务的复杂性,静态 kv 缓存有三种使用方式:
- 基本使用:只需在
generation_config中设置一个标志(推荐); - 高级使用:处理一个多轮生成或自定义生成循环的缓存对象;
- 高级使用:如果单图对你有意义,可以将整个
generate函数编译成单个图。
选择下面的正确标签以获取每种方式的详细说明。
无论使用哪种策略与 torch.compile 结合,如果你将 LLM 输入左填充到有限的几个值,可以避免与形状相关的重新编译。pad_to_multiple_of 分词器标志 是你的朋友!
基本使用:generation_config¶
高级使用:控制静态缓存¶
高级使用:端到端生成编译¶
以下示例使用 Gemma 模型。所有需要做的就是:
- 访问模型的
generation_config属性并将cache_implementation设置为“static”; - 调用
torch.compile对模型进行编译,以便在静态 kv 缓存下运行前向传播。
就是这样!