辅助生成: 低延迟文本生成的新方向

大型语言模型如今风靡一时，许多公司投入大量资源来扩展它们规模并解锁新功能。然而，作为注意力持续时间不断缩短的人类，我们并不喜欢大模型缓慢的响应时间。由于延迟对于良好的用户体验至关重要，人们通常使用较小的模型来完成任务，尽管它们的质量较低 (例如 代码补全任务)。

为什么文本生成这么慢？是什么阻止你在不破产的情况下部署低延迟大型语言模型？在这篇博文中，我们将重新审视自回归文本生成的瓶颈，并介绍一种新的解码方法来解决延迟问题。你会发现，通过使用我们的新的辅助生成方法，你可以将硬件中的延迟降低多达 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 = tokenizer(["Hello world"], return_tensors="pt").to("cuda")
 
def print_tokens_per_second(batch_size):
    new_tokens = 100
    cumulative_time = 0
 
    # warmup
    model.generate(
        **inputs, do_sample=True, max_new_tokens=new_tokens, num_return_sequences=batch_size
    )
 
    for _ in range(10):
        start = time.time()
        model.generate(
            **inputs, do_sample=True, max_new_tokens=new_tokens, num_return_sequences=batch_size
        )
        cumulative_time += time.time() - start
    print(f"Tokens per second: {new_tokens * batch_size * 10 / cumulative_time:.1f}")
 
print_tokens_per_second(1) # Tokens per second: 418.3
print_tokens_per_second(64) # Tokens per second: 16266.2 (~39x more tokens per second)

最后，如果你有多个可用设备，你可以使用 Tensor 并行 分配工作负载并获得更低的延迟。使用 Tensor 并行，你可以将内存带宽负担分摊到多个设备上，但除了在多个设备运行计算的成本之外，你还需要考虑设备间的通信瓶颈。该方法的收益在很大程度上取决于模型大小: 对于可以轻松在单个消费级设备上运行的模型，通常效果并不显著。根据这篇 DeepSpeed 博客，你会发现你可以将大小为 17B 的模型分布在 4 个 GPU 上，从而将延迟减少 1.5 倍 (图 7)。

这三种类型的改进可以串联使用，从而产生 高通量解决方案。然而，在应用特定于硬件的优化后，降低延迟的方法有限——并且现有的方法很昂贵。让我们接下来解决这个问题！

重新回顾语言模型解码器的正向传播

上文我们讲到，每个模型前向传递都会产生下一个 token 的概率，但这实际上是一个不完整的描述。在文本生成期间，典型的迭代包括模型接收最新生成的 token 作为输入，加上所有其他先前输入的缓存内部计算，再返回下一个 token 得概率。缓存用于避免冗余计算，从而实现更快的前向传递，但它不是强制性的 (并且可以设置部分使用)。禁用缓存时，输入包含到目前为止生成的整个 token 序列，输出包含 _所有位置_的下一个 token 对应的概率分布！如果输入由前 N 个 token 组成，则第 N 个位置的输出对应于其下一个 token 的概率分布，并且该概率分布忽略了序列中的所有后续 token。在贪心解码的特殊情况下，如果你将生成的序列作为输入传递并将 argmax 运算符应用于生成的概率，你将获得生成的序列。

from transformers import AutoModelForCausalLM, AutoTokenizer
 
tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
 
inputs = tok(["The"], return_tensors="pt")
generated = model.generate(**inputs, do_sample=False, max_new_tokens=10)
forward_confirmation = model(generated).logits.argmax(-1)
 
# We exclude the opposing tips from each sequence: the forward pass returns
# the logits for the next token, so it is shifted by one position.
print(generated[:-1].tolist() == forward_confirmation[1:].tolist()) # True

这意味着你可以将模型前向传递用于不同的目的: 除了提供一些 token 来预测下一个标记外，你还可以将序列传递给模型并检查模型是否会生成相同的序列 (或部分相同序列)。

(请访问阅读原文查看动态演示)

让我们想象，你可以访问一个神奇的无延迟的预测辅助模型，该模型针对任何给定输入生成与你的模型相同的序列。顺便说一句，这个模型不能直接用，只能辅助你的生成程序。使用上述属性，你可以使用此辅助模型获取候选输出 token，然后使用你的模型进行前向传递以确认它们的正确性。在这个乌托邦式的场景中，文本生成的延迟将从 O(n) 减少到 O(1)，其中生成的 token 数量为 n。对于需要多次迭代生成的过程，我们谈论的是其数量级。

向现实迈出一步，我们假设辅助模型失去了它的预测属性。根据你的模型，现在它是一个无延迟模型，但它会弄错一些候选 token。由于任务的自回归性质，一旦辅助模型得到一个错误的 token，所有后续候选 token 都必须无效。但是，你可以使用模型更正错误 token 并反复重复此过程后再次查询辅助模型。即使辅助模型失败了几个 token，文本生成的延迟也会比原始形式小得多。

显然，世界上没有无延迟的辅助模型。然而，找到一个近似于模型的文本生成输出的其它模型相对容易，例如经过类似训练的相同架构的较小版本模型通常符合此需求。当模型大小的差异变得显著时，使用较小的模型作为辅助模型的成本在跳过几个前向传递后就显得无关紧要了！现在，你了解了 _ 辅助生成 _ 的核心。

使用辅助模型的贪心解码

辅助生成是一种平衡行为。你希望辅助模型快速生成候选序列，同时尽可能准确。如果辅助模型的质量很差，你将承担使用辅助模型的成本，而收益却很少甚至没有。另一方面，优化候选序列的质量可能意味着使用更慢的辅助模型，从而导致网络减速。虽然我们无法为你自动选择辅助模型，但我们包含了一个额外的要求和一个启发式方法，以确保模型与辅助模型一起花费的时间保持在可控范围内。

首先，我们要求辅助模型必须具有与你的模型完全相同的分词器。如果没有此要求，则必须添加昂贵的 token 解码和重新编码步骤。此外，这些额外的步骤必须在 CPU 上进行，这反过来可能增加了设备间数据传输。能够快速地使用辅助模型对于辅助生成的好处是至关重要的。

最后，启发式。至此，你可能已经注意到电影盗梦空间和辅助生成之间的相似之处——毕竟你是在文本生成中运行文本生成。每个候选 token 有一个辅助模型前向传播，我们知道前向传播是昂贵的。虽然你无法提前知道辅助模型将获得的 token 数量，但你可以跟踪此信息并使用它来限制向辅助模型请求的候选 token 数量——输出的某些部分比其它一些部分更容易被预计。

总结一下，这是我们最初实现的辅助生成的循环 (代码):

1. 使用贪心解码与辅助模型生成一定数量的候选 token。当第一次调用辅助生成时，生成的候选 token 的数量被初始化为 5。

2. 使用我们的模型，对候选 token 进行前向计算，获得每个 token 对应的概率。

3. 使用 token 选择方法 (使用.argmax() 进行贪心搜索或使用 .multinomial() 用于采样方法) 来从概率中选取 next_tokens。

4. 比较步骤 3 中选择的 next_tokens 和  候选 token 中相同的 token 数量。请注意，我们需要从左到右进行比较， 在第一次不匹配后，后续所有 候选 token都无效。5. 使用步骤 4 得到的匹配数量将候选 token 分割。也就是，将输入 tokens 加上刚刚验证得到的正确的 tokens。

5. 调整下一次迭代中生成的候选 token 的数量 —— 使用启发式方法，如果步骤 3 中所有 token 都匹配，则候选 token 的长度增加 2，否则减少 1。

(请访问阅读原文查看动态演示)

我们在