对 Intel Gaudi 更快的辅助生成支持

发布于 2024 年 6 月 4 日
在 GitHub 上更新

随着模型规模的增长,生成式人工智能的实现需要大量的推理资源。这不仅增加了每次生成的成本,还增加了处理此类请求的功耗。

文本生成推理优化对于降低延迟、基础设施成本和功耗至关重要。这可以改善用户体验并提高文本生成任务的效率。

辅助解码是一种加速文本生成的流行方法。我们针对 Intel Gaudi 对其进行了调整和优化,如上一篇文章所示,它提供了与 Nvidia H100 GPU 相似的性能,而其价格与 Nvidia A100 80GB GPU 相似。这项工作现已成为 Optimum Habana 的一部分,Optimum Habana 扩展了各种 Hugging Face 库(如 Transformers 和 Diffusers),以便您的 AI 工作流针对 Intel Gaudi 处理器进行全面优化。

推测采样 - 辅助解码

推测采样是一种用于加速文本生成的技术。它的工作原理是生成一个草稿模型,该模型生成 K 个 token,然后这些 token 在目标模型中进行评估。如果草稿模型被拒绝,则使用目标模型生成下一个 token。这个过程重复进行。通过使用推测采样,我们可以提高文本生成的速度,并实现与自回归采样相似的采样质量。该技术允许我们在生成文本时指定一个草稿模型。这种方法已被证明可以为大型基于 Transformer 的模型提供大约 2 倍的加速。总的来说,这些技术可以加速文本生成并提高 Intel Gaudi 处理器上的性能。

然而,草稿模型和目标模型具有不同大小的 KV 缓存,因此挑战在于同时利用独立的优化策略。本文假设一个量化模型,并结合推测采样利用 KV 缓存。请注意,每个模型都有自己的 KV 缓存,草稿模型用于生成 K 个 token,然后这些 token 在目标模型中进行评估。当草稿模型被拒绝时,目标模型用于生成下一个 token。草稿模型用于生成接下来的 K 个 token,然后重复这个过程。

请注意,作者 [2] 证明在执行推测采样时目标分布是恢复的——这保证了与目标本身上的自回归采样相同的采样质量。因此,不利用推测采样不值得的情况与草稿模型的相对大小没有足够的节省或草稿模型的接受率不足以从草稿模型较小的尺寸中受益有关。

有一种与推测采样类似的技术,称为辅助生成。这是在大致相同的时间 [3] 独立开发的。作者将此方法集成到 Hugging Face Transformers 中,现在 `_generate()` 调用具有可选的 `assistant_model` 参数以启用此方法。

用法与实验

辅助生成的使用非常简单。此处提供了一个示例:https://github.com/huggingface/optimum-habana/tree/main/examples/text-generation#run-speculative-sampling-on-gaudi。不出所料,参数 `_assistant_model` 用于指定草稿模型。草稿模型用于生成 K 个 token,然后这些 token 在目标模型中进行评估。当草稿模型被拒绝时,目标模型用于生成下一个 token。草稿模型用于生成接下来的 K 个 token,然后重复这个过程。草稿模型的接受率部分取决于输入文本。通常,我们看到大型基于 Transformer 的模型有大约 2 倍的加速。

结论

使用 Gaudi 加速文本生成现在得到支持且易于使用。这可以用于提高 Intel Gaudi 处理器上的性能。该方法基于推测采样,该方法已被证明在提高大型基于 Transformer 的模型性能方面是有效的。

参考文献

[1] N. Shazeer, “Fast Transformer Decoding: One Write-Head is All You Need,” Nov. 2019. arXiv:1911.02150。

[2] C. Chen, S. Borgeaud, G. Irving, J.B. Lespiau, L. Sifre, and J. Jumper, “Accelerating Large Language Model Decoding with Speculative Sampling,” Feb. 2023. arXiv:2302.01318。

[3] J. Gante, “Assisted Generation: a new direction toward low-latency text generation,” May 2023, https://huggingface.co/blog/assisted-generation

社区

注册登录 发表评论