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

大型语言模型目前非常流行，许多公司都在投入大量资源来扩展这些模型并开发新的功能。然而，作为拥有越来越短注意力的人类，我们也不喜欢它们的响应速度慢。延迟对于良好的用户体验至关重要，因此尽管质量较低（例如在代码补全中），人们通常还是使用较小的模型。

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

理解文本生成的延迟

现代文本生成的核心原理很容易理解。让我们来看看其核心部分——机器学习模型。输入包含一个文本序列，其中包括迄今为止生成的文本以及其他可能的模型特定组件（例如，Whisper还具有音频输入）。模型接受输入并运行前向传递：将输入馈送到模型并依次通过其各个层，直到预测出下一个标记的非归一化对数概率（也称为logits）。一个标记可以由整个单词、子单词或甚至单个字符组成，这取决于模型的设计。如果您想深入了解文本生成的这个部分，可以参考图示的GPT-2。

模型的前向传递可以获得下一个标记的logits，您可以自由操作这些logits（例如，将不希望出现的单词或序列的概率设为0）。文本生成的下一步是从这些logits中选择下一个标记。常见的策略包括选择最有可能的标记，即贪婪解码，或从它们的分布中进行抽样，也称为多项分布抽样。通过将模型的前向传递与下一个标记的选择迭代地结合起来，就可以实现文本生成。当涉及到解码方法时，这只是冰山一角，请参考我们关于文本生成的博客文章以进行深入探索。

从上面的描述中，文本生成的延迟瓶颈显而易见：对于大型模型来说，运行模型的前向传递速度较慢，您可能需要按顺序进行数百次前向传递。但让我们深入探讨一下：为什么前向传递速度慢？前向传递通常由矩阵乘法主导，通过快速访问对应的维基百科条目，您可以了解到在这个操作中，内存带宽是限制因素（例如，从GPU内存到GPU计算核心）。换句话说，前向传递的瓶颈来自于将模型层的权重加载到设备的计算核心中，而不是来自于执行计算本身。

目前，有三个主要途径可以提高文本生成的性能，都是解决模型前向传递性能的。首先，您可以进行硬件特定的模型优化。例如，您的设备可能与Flash Attention兼容，通过重新排序操作来加速注意力层，或者使用INT8量化来减小模型权重的大小。

其次，当您知道会同时进行多个文本生成请求时，您可以对输入进行批处理，从而大幅增加吞吐量，但会有一定的延迟惩罚。设备加载的模型层权重现在可以在并行处理多个输入行上使用，这意味着您可以在大致相同的内存带宽负担下获得更多的标记输出。批处理的问题在于需要额外的设备内存（或将内存转移到其他地方）- 在这个范围的末端，您可以看到类似FlexGen的项目，它以牺牲延迟为代价来优化吞吐量。

# 展示批量生成对性能的影响的示例。测量设备：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

    # 预热
    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"每秒生成的标记数：{new_tokens * batch_size * 10 / cumulative_time:.1f}")

print_tokens_per_second(1)   # 每秒生成的标记数：418.3
print_tokens_per_second(64)  # 每秒生成的标记数：16266.2（每秒大约生成39倍的标记数）

最后，如果您有多个可用设备，您可以使用张量并行性分配工作负载，并获得更低的延迟。使用张量并行性，您将内存带宽负担分散到多个设备上，但现在您还必须考虑设备间通信瓶颈，除了运行多个设备的经济成本。收益在很大程度上取决于模型大小：可以轻松适应单个消费设备的模型收益非常有限。从这篇 DeepSpeed 博客文章的结果中可以看到，您可以将一个具有 17B 参数的模型分布到 4 个 GPU 上，将延迟减少 1.5 倍（图 7）。

这三种改进方法可以同时使用，从而实现高吞吐量的解决方案。然而，在应用硬件特定的优化之后，减少延迟的选择有限 – 而且现有的选择都很昂贵。让我们解决这个问题！

语言解码器前向传递，重新审视

您已经了解到，每次模型前向传递都会产生下一个标记的对数几率，但这实际上是不完整的描述。在文本生成过程中，典型的迭代过程是模型接收最新生成的标记作为输入，以及所有其他先前输入的缓存内部计算，返回下一个标记的对数几率。缓存用于避免冗余计算，从而实现更快的前向传递，但它不是强制性的（可以部分使用）。当禁用缓存时，输入包含到目前为止生成的所有标记序列，输出包含序列中所有位置的下一个标记的对数几率！位置 N 处的对数几率对应于如果输入由前 N 个标记组成，则忽略序列中所有后续标记的下一个标记的分布。在贪婪解码的特殊情况下，如果您将生成的序列作为输入，并将 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

这意味着您可以将模型的前向传递用于不同的目的：除了传递一些标记以预测下一个标记外，您还可以将序列传递给模型，并通过您的模型进行确认以验证它们是否正确。在这个理想化的情景中，文本生成的延迟将从 O(n) 减少到 O(1)，其中 n 是生成的标记数。对于长文本生成，我们谈论的是几个数量级的差异。

朝着现实迈出一步，让我们假设助手模型失去了预测性质。现在它是一个无延迟模型，根据您的模型，它会得出一些候选标记的错误。由于任务的自回归性质，一旦助手获得一个错误的标记，所有后续的候选标记都必须无效。然而，这并不妨碍您在纠正了错误的标记后再次查询助手，并反复进行此过程。即使助手失败了一些标记，文本生成的延迟仍然比原始形式的延迟小一个数量级。

显然，没有无延迟的助手模型。然而，相对容易找到一个近似某个其他模型文本生成输出的模型 – 训练方式相似的较小版本通常符合此属性。此外，当模型大小的差异变得显著时，使用较小模型作为助手的成本在考虑跳过一些前向传递的好处后变得微不足道！现在您已经理解了辅助生成的核心。

辅助生成的贪婪解码

辅助生成是一种平衡行为。您希望助手能够快速生成候选序列，同时尽可能准确。如果助手质量较差，使用助手模型的成本几乎没有好处。另一方面，优化候选序列的质量可能意味着使用较慢的助手，导致净减速。虽然我们无法为您自动选择助手模型，但我们已经包含了一个额外的要求和一种启发式方法，以确保与助手一起花费的时间保持在适当范围内。

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

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

总结一下，这是我们辅助生成循环的原始实现（代码）：

1. 使用贪婪解码生成一定数量的候选标记，使用助手模型生成。第一次调用辅助生成时，产生候选标记的数量初始化为。
2. 使用我们的模型，对进行前向传递，获取。
3. 使用标记选择方法（贪婪搜索的或采样的）从中获取。
4. 将与进行比较，并获取匹配标记的数量。请记住，此比较必须按照从左到右的因果关系进行：在第一个不匹配之后，所有候选标记都无效。
5. 使用匹配的标记加上第一个不同的标记（我们的模型从有效的候选子序列生成）来切分和丢弃与未确认候选标记相关的变量。
6. 调整下一次迭代中要生成的候选标记数量-我们的原始启发式方法在所有标记匹配时将其增加，否则将其减小<1。

我们在🤗 Transformers中设计了API，使此过程对您来说非常简便。您只需要将助手模型传递给新的关键字参数，并获得延迟减少的好处！在发布此博文时，辅助生成仅限于批大小为。

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']

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

浏览收集到的数字时，我们可以看到辅助生成在不同的环境中可以带来显著的延迟降低，但它并不是万能解决方法-在将其应用于您的用例之前，您应该对其进行基准测试。我们可以得出结论，辅助生成：

1. 🤏 需要访问至少比您的模型小一个数量级的助手模型（差异越大，越好）；
2. 🚀 在 INT8 存在时，可以获得高达3倍的加速，否则最高为2倍，当模型适合GPU内存时；
3. 🤯 如果您正在处理无法适应GPU的模型，并依赖于内存卸载，可以获得高达10倍的加速；
4. 📄 在输入有依据的任务（如自动语音识别或摘要）中表现出色。

辅助生成的示例

贪婪解码适用于与输入相关的任务（自动语音识别，翻译，摘要，…）或事实知识寻求。对于需要大量创造力的开放式任务，例如将语言模型用作聊天机器人，应改用采样。辅助生成自然设计用于贪婪解码，但这并不意味着您不能使用多项采样进行辅助生成！

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

为什么不亲自试试，感受一下辅助生成的感觉呢？

未来的方向

辅助生成表明，现代文本生成策略可以进行优化。了解当前的问题是一个受到内存限制的问题，而不是计算限制的问题，使我们能够应用简单的启发式方法，以充分利用可用的内存带宽，缓解瓶颈。我们相信，进一步改进助理模型的使用将带来更大的延迟降低 – 例如，如果我们要求助手生成多个候选续写，我们可能能够跳过更多前向传递。当然，发布高质量的小型模型供作为助理使用将对实现和放大效益至关重要。

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

最后，辅助生成重新引发了文本生成中的一个关键问题。该领域一直在以固定数量的计算生成所有新标记的约束下发展，对于给定的模型。每个均匀的前向传递中的一个标记，以纯自回归的方式。本博文强调了这不应该是这种情况：生成的输出的大部分可以由模型大小的一小部分同样生成。为此，我们将需要新的模型架构和解码方法-我们很兴奋看到未来会带来什么。

相关工作

在本博文最初发布后，我注意到其他作品探索了同样的核心原理（使用前向传递验证更长的续写）。特别是，请查看以下作品：

• 谷歌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设计了我们网页上看到的美丽封面。