辅助生成：一种低延迟文本生成的新方向

大型语言模型如今炙手可热，许多公司投入大量资源来扩展它们并解锁新功能。然而，作为注意力持续时间不断缩短的人类，我们也不喜欢它们缓慢的响应时间。延迟对于良好的用户体验至关重要，因此尽管小型模型的质量较低（例如在代码补全中），也常常被使用。

为什么文本生成如此缓慢？是什么阻止了您在不破产的情况下部署低延迟大型语言模型？在这篇博文中，我们将重新审视自回归文本生成的瓶颈，并引入一种新的解码方法来解决延迟问题。您将看到，通过使用我们的新方法——辅助生成，您可以在商用硬件上将延迟降低多达 10 倍！

理解文本生成延迟

现代文本生成的核心很容易理解。让我们看看其核心部分，即机器学习模型。它的输入包含一个文本序列，其中包括迄今为止生成的文本，以及可能存在的其他模型特定组件（例如，Whisper 也有音频输入）。模型接收输入并执行前向传播：输入被馈送到模型并沿其层顺序传递，直到预测出下一个 token 的未归一化对数概率（也称为 logits）。一个 token 可以是整个单词、子词，甚至单个字符，这取决于模型。GPT-2 图解是您深入了解文本生成这部分的绝佳参考。

模型的前向传播会为您提供下一个 token 的 logits，您可以自由操作它们（例如，将不想要的单词或序列的概率设置为 0）。文本生成的下一步是从这些 logits 中选择下一个 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)

最后，如果您有多个设备可用，您可以使用张量并行来分发工作负载并获得更低的延迟。使用张量并行，您可以将内存带宽负担分散到多个设备上，但现在除了运行多个设备的金钱成本之外，您还必须考虑设备间通信瓶颈。其好处在很大程度上取决于模型大小：对于可以轻松放入单个消费级设备的模型，好处非常有限。根据这篇DeepSpeed 博客文章的结果，您可以看到可以将一个 17B 参数模型分散到 4 个 GPU 上，以将延迟降低 1.5 倍（图 7）。

这三种类型的改进可以协同使用，从而产生高吞吐量解决方案。然而，在应用特定硬件优化之后，减少延迟的选择有限——而且现有选择成本高昂。让我们解决这个问题！

语言解码器前向传播，再探

您上面读到，每次模型前向传播都会生成下一个 token 的 logits，但这实际上是一个不完整的描述。在文本生成过程中，典型的迭代是模型接收最新生成的 token 作为输入，加上所有其他先前输入的缓存内部计算，返回下一个 token 的 logits。缓存用于避免冗余计算，从而加快前向传播速度，但这不是强制性的（并且可以部分使用）。当禁用缓存时，输入包含迄今为止生成的整个 token 序列，输出包含对应于序列中所有位置的下一个 token 的 logits！位置 N 的 logits 对应于如果输入由前 N 个 token 组成（忽略序列中所有后续 token）时下一个 token 的分布。在贪婪解码的特定情况下，如果您将生成的序列作为输入并对结果 logits 应用 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[0, 1:].tolist() == forward_confirmation[0, :-1].tolist())  # True

这意味着您可以将模型前向传播用于不同的目的：除了输入一些 token 来预测下一个 token，您还可以将一个序列传递给模型，并再次检查模型是否会生成相同的序列（或其一部分）。

让我们考虑一下，您可以使用一个神奇的无延迟预言模型，该模型可以为任何给定输入生成与您的模型相同的序列。为了便于讨论，它不能直接使用，它仅限于作为您生成过程的助手。使用上面描述的属性，您可以使用这个助手模型来获取候选输出 token，然后通过您的模型进行前向传播以确认它们确实是正确的。在这种乌托邦式的情况下，文本生成的延迟将从 O(n) 减少到 O(1)，其中 n 是生成的 token 数量。对于长时间生成，我们谈论的是几个数量级。

更接近现实一步，假设助手模型失去了其预言属性。现在它是一个无延迟模型，根据您的模型，它会错误地预测一些候选 token。由于任务的自回归性质，一旦助手错误地预测了一个 token，所有后续的候选 token 都必须失效。然而，这并不能阻止您在用您的模型更正错误的 token 后再次查询助手，并迭代地重复此过程。即使助手预测错误了几个 token，文本生成的延迟也比其原始形式低一个数量级。

显然，不存在无延迟的辅助模型。然而，相对容易找到一个模型，其文本生成输出近似于另一个模型——相同架构的较小版本，以类似方式训练，通常符合此属性。此外，当模型大小差异显着时，使用较小模型作为助手的成本在考虑到跳过一些前向传播的好处后，变得微不足道！您现在了解了辅助生成的核心。

使用辅助生成进行贪婪解码

辅助生成是一种平衡行为。您希望助手快速生成候选序列，同时尽可能准确。如果助手质量差，您将承担使用助手模型的成本，而几乎没有任何好处。另一方面，优化候选序列的质量可能意味着使用缓慢的助手，从而导致总体速度变慢。虽然我们无法为您自动选择助手模型，但我们添加了一个额外要求和启发式方法，以确保与助手在一起的时间保持在可控范围内。

首先，要求是——助手必须与您的模型具有完全相同的分词器。如果未满足此要求，则必须添加昂贵的 token 解码和重新编码步骤。此外，这些额外步骤必须在 CPU 上进行，这反过来可能需要缓慢的设备间数据传输。助手的快速使用对于辅助生成的好处显现至关重要。

最后，启发式方法。到目前为止，您可能已经注意到电影《盗梦空间》与辅助生成之间的相似之处——毕竟，您正在文本生成内部运行文本生成。每个候选 token 都会有一个辅助模型的前向传播，我们知道前向传播是昂贵的。虽然您无法提前知道辅助模型将正确识别多少个 token，但您可以跟踪此信息并使用它来限制向辅助模型请求的候选 token 数量——输出的某些部分比其他部分更容易预测。

总而言之，以下是我们辅助生成循环的原始实现（代码）

1. 使用贪婪解码通过辅助模型生成一定数量的候选 token，生成 candidates。第一次调用辅助生成时，生成的候选 token 数量初始化为 5。
2. 使用我们的模型，对 candidates 进行前向传播，获得 logits。
3. 使用 token 选择方法（贪婪搜索的 .argmax() 或采样的 .multinomial()）从 logits 中获取 next_tokens。
4. 将 next_tokens 与 candidates 进行比较，并获取匹配的 token 数量。请记住，此比较必须从左到右进行因果关系：在第一次不匹配之后，所有候选者都将失效。
5. 使用匹配数量来切分内容并丢弃与未确认候选 token 相关的变量。本质上，在 next_tokens 中，保留匹配的 token 以及第一个分歧的 token（我们的模型从有效的候选子序列生成）。
6. 调整下一次迭代中要生成的候选 token 的数量——如果所有 token 都匹配，我们的原始启发式方法将其增加 2，否则减少 1。

我们以无忧的方式设计了 🤗 Transformers 中的 API。您只需在新参数 assistant_model 中传递辅助模型，即可获得延迟收益！在发布此博客文章时，辅助生成仅限于批处理大小为 1。

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

prompt = "Alice and Bob"
checkpoint = "EleutherAI/pythia-1.4b-deduped"
assistant_checkpoint = "EleutherAI/pythia-160m-deduped"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt").to(device)

model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint).to(device)
outputs = model.generate(**inputs, assistant_model=assistant_model)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
# ['Alice and Bob are sitting in a bar. Alice is drinking a beer and Bob is drinking a']

额外的内部复杂性是否值得？让我们看一下贪婪解码情况下的延迟数据（采样结果在下一节中），考虑到批处理大小为 1。这些结果直接来自 🤗 Transformers，没有任何额外优化，因此您应该能够在您的设置中重现它们。

粗略地看一下收集到的数字，我们看到辅助生成可以在不同的设置中显著减少延迟，但它并非万能药——您应该在将其应用于您的用例之前进行基准测试。我们可以得出结论，辅助生成：

1. 🤏 需要访问一个至少比您的模型小一个数量级的辅助模型（差异越大越好）；
2. 🚀 在存在 INT8 的情况下，当模型适合 GPU 内存时，可将速度提高至 3 倍；否则，可提高至 2 倍；
3. 🤯 如果您使用的模型不适合您的 GPU 并且依赖内存卸载，您可以看到高达 10 倍的加速；
4. 📄 在输入型任务中表现出色，例如自动语音识别或摘要。

使用辅助生成进行采样

贪婪解码适用于输入依赖型任务（自动语音识别、翻译、摘要等）或寻求事实知识的任务。需要高度创造性的开放式任务，例如语言模型作为聊天机器人的大多数用途，应改为使用采样。辅助生成自然是为贪婪解码设计的，但这并不意味着您不能将辅助生成与多项式采样一起使用！

从下一个 token 的概率分布中抽取样本将导致我们的贪婪助手更频繁地失败，从而降低其延迟收益。然而，我们可以使用大多数基于采样的应用程序中存在的温度系数来控制下一个 token 的概率分布的尖锐程度。在一个极端，当温度接近 0 时，采样将近似于贪婪解码，偏向最可能的 token。在另一个极端，当温度设置为远大于 1 的值时，采样将是混沌的，从均匀分布中抽取。因此，低温对您的助手模型更有利，保留了辅助生成的大部分延迟收益，如下所示。

为什么不亲自体验一下，感受一下辅助生成呢？

未来方向

辅助生成表明，现代文本生成策略已成熟，可进行优化。理解它目前是内存受限问题，而不是计算受限问题，使我们能够应用简单的启发式方法，以最大限度地利用可用内存带宽，从而缓解瓶颈。我们相信，对辅助模型使用的进一步完善将使我们获得更大的延迟降低——例如，如果我们请求辅助模型生成多个候选续写，我们可能能够跳过更多的前向传播。自然，发布高质量的小型模型以用作辅助模型，对于实现和放大收益至关重要。

最初发布在我们的 🤗 Transformers 库中，与 .generate() 函数一起使用，我们期望在整个 Hugging Face 生态系统中提供它。其实现也完全开源，因此，如果您正在进行文本生成并且不使用我们的工具，请随意将其用作参考。

最后，辅助生成重新提出了文本生成中的一个关键问题。该领域一直在发展，其约束是所有新 token 都是给定模型固定计算量的结果。在纯自回归模式下，每个同构前向传播产生一个 token。这篇博客文章强化了不应如此的观点：生成的输出的大部分子部分也可以由大小仅为一部分的模型同样生成。为此，我们将需要新的模型架构和解码方法——我们很高兴看到未来会怎样！

相关工作

在这篇博客文章最初发布后，我注意到其他工作也探索了相同的核心原理（使用前向传播来验证更长的续写）。特别是，请查看以下工作：

• 分块并行解码，由 Google Brain
• 推测性采样，由 DeepMind

引用

@misc {gante2023assisted,
    author       = { {Joao Gante} },
    title        = { Assisted Generation: a new direction toward low-latency text generation },
    year         = 2023,
    url          = { https://huggingface.co/blog/assisted-generation },
    doi          = { 10.57967/hf/0638 },
    publisher    = { Hugging Face Blog }
}

致谢

我要感谢 Sylvain Gugger、Nicolas Patry 和 Lewis Tunstall 提供了许多宝贵的建议来改进这篇博客文章。最后，感谢 Chunte Lee 为我们网页设计了精美的封面。