本文探讨了大型语言模型在文本生成中面临的延迟问题，并提出了一种新的解码方法——辅助生成，通过使用小型辅助模型来减少延迟，最多可降低10倍。文章介绍了文本生成的延迟瓶颈、现有优化方法以及如何利用辅助模型来加速生成过程，特别适合输入驱动的任务。此外，还讨论了辅助生成在不同解码策略（如贪心和采样）下的表现，以及未来的优化方向。
大型语言模型如今风靡一时，许多公司投入大量资源来扩展它们规模并解锁新功能。然而，作为注意力持续时间不断缩短的人类，我们并不喜欢大模型缓慢的响应时间。由于延迟对于良好的用户体验至关重要，人们通常使用较小的模型来完成任务，尽管它们的质量较低 (例如 代码补全任务)。
为什么文本生成这么慢？是什么阻止你在不破产的情况下部署低延迟大型语言模型？在这篇博文中，我们将重新审视自回归文本生成的瓶颈，并介绍一种新的解码方法来解决延迟问题。你会发现，通过使用我们的新的辅助生成方法，你可以将硬件中的延迟降低多达 10 倍！
理解文本生成延迟
文本生成的核心很容易理解。让我们看看核心部分 (即 ML 模型)，它的输入包含一个文本序列，其中包括到目前为止生成的文本，以及其他特定于模型的组件 (例如 Whisper 还有一个音频输入)。该模型接受输入并进行前向传递: 输入被喂入模型并一层一层顺序传递，直到预测出下一个 token 的非标准化对数概率 (也称为 logits)。一个 token 可能包含整个词、子词，或者是单个字符，这取决于具体模型。如果你想深入了解文本生成的原理，GPT-2 插图 是一个很好的参考。
(请访问阅读原文查看动态演示)
模型的前向传递提供了下一个 token 的概率，你可以自由操作 (例如，将不需要的单词或序列的概率设置为 0)。文本生成的步骤就是从这些概率中选择下一个 token。常见的策略包括选择最有可能的 token (贪心解码)，或从它们的分布中抽样 (多项式抽样)。在选择了下一个 token 之后，我们将模型前向传递与下一个 token 迭代地连接起来，继续生成文本。这个解释只是解码方法的冰山一角; 请参阅我们 关于文本生成的博客 以进行深入探索。
(请访问阅读原文查看动态演示)
从上面的描述中可以看出，文本生成的延迟瓶颈很明显: 运行大型模型的前向传递很慢，你可能需要依次执行数百次迭代。但让我们深入探讨一下: 为什么前向传递速度慢？前向传递通常以矩阵乘法为主，通过查阅相应的 维基百科，你可以看出内存带宽是此操作的限制 (例如，从 GPU RAM 到 GPU 计算核心)。换句话说， _前向传递的瓶颈来自将模型权重加载到设备的计算核心中，而不是来自执行计算本身_。
目前，你可以探索三个主要途径来充分理解文本生成，所有这些途径都用于解决模型前向传递的性能问题。首先，对于特定硬件的模型优化。例如，如果你的设备可能与 Flash Attention 兼容，你可以使用它通可以过重新排序操作或 INT8 量化 来加速注意力层，其减少了模型权重的大小。
其次，如果你有并发文本生成需求，你可以对输入进行批处理，从而实现较小的延迟损失并大幅增加吞吐量。你可以将模型对于多个输入并行计算，这意味着你将在大致相同的内存带宽负担情况下获得了更多 token。批处理的问题在于你需要额外的设备内存 (或在某处卸载内存)。你可以看到像 FlexGen 这样的项目以延迟为代价来优化吞吐量。
# Example showcasing the impact of batched generation. Measurement device: RTX3090
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2").to("cuda")
inputs =
最低0.47元/天 解锁文章
扫一扫
738
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
