Hugging Face 提出辅助生成，用小型助手模型降低自回归文本生成延迟

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

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

理解文本生成延迟

现代文本生成的核心理解起来很直接。让我们看看核心部分，即机器学习模型。它的输入包含一个文本序列，其中包括到目前为止生成的文本，以及可能的其他模型特定组件（例如，Whisper 还有音频输入）。模型接收输入并运行一次前向传播：输入被馈送到模型，并依次通过其各层，直到预测出下一个标记的未归一化对数概率（也称为 logits）。根据模型的不同，一个标记可能由整个单词、子词甚至单个字符组成。如果你想深入了解文本生成的这一部分，图解 GPT-2 是一个很好的参考。

一次模型前向传播会得到下一个标记的 logits，你可以自由地操纵它们（例如，将不希望的单词或序列的概率设为0）。文本生成的下一步是从这些 logits 中选择下一个标记。常见的策略包括选择最可能的标记，即贪婪解码，或从其分布中采样，也称为多项式采样。将模型前向传播与下一个标记选择迭代链接起来，就实现了文本生成。这一解释只是解码方法的冰山一角；请参阅我们关于文本生成的博客文章以进行深入探索。

从上述描述中，文本生成中的延迟瓶颈显而易见：对大型模型运行一次前向传播很慢，而你可能需要按顺序运行数百次。但让我们更深入地探讨：为什么前向传播很慢？前向传播通常由矩阵乘法主导，快速查看相应的维基百科部分后，你可以看出内存带宽是此操作的限制因素（例如，从 GPU 内存到 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 Parallelism（张量并行）来分摊工作负载，从而获得更低的延迟。使用张量并行，你将内存带宽负担分散到多台设备上，但此时除了运行多台设备的金钱成本外，你还必须考虑设备间通信的瓶颈。收益在很大程度上取决于模型规模：能够轻松装进单台消费级设备的模型收益非常有限。从这篇 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 解码和重新编码步骤。此外，这些额外步骤必须在 CPU 上执行，这又可能需要缓慢的设备间数据传输。助手的快速使用对于辅助生成展现其优势至关重要。

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

综上所述，这是我们的辅助生成循环的原始实现（代码）：

使用贪心解码来生成一定数量的候选标记，由辅助模型产生candidates。首次调用辅助生成时，生成的候选标记数量初始化为5。
利用我们的模型，对candidates进行一次前向传递，获得logits。
使用标记选择方法（贪心搜索为.argmax()，采样为.multinomial()）从logits中得到next_tokens。
将next_tokens与candidates对比，得到匹配标记的数量。请注意，此比较必须遵循从左到右的因果性：一旦出现首个不匹配，后续所有候选均视为无效。
根据匹配数量进行切片处理，并丢弃与未确认候选标记相关的变量。本质上，在next_tokens中，保留匹配的标记以及第一个分歧标记（我们的模型从有效的候选子序列生成该标记）。
调整下一次迭代中将要生成的候选标记数量——我们的原始启发式策略是：如果所有标记都匹配，则增加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，未做任何额外优化，因此您应该能在自己的环境中复现它们。

浏览收集到的数据，我们可以看出，辅助生成能在多种场景下显著减少延迟，但它并非万能灵药——在应用到您的具体用例之前，最好先进行基准测试。我们可以得出结论，辅助生成：

🤏 需要访问一个至少比您的模型小一个数量级的辅助模型（差距越大越好）；
🚀 当模型能容纳于GPU内存时，若存在INT8，可获得最高3倍的加速，否则最高2倍；
🤯 如果您的模型超出GPU内存并依赖内存卸载，可看到最高10倍的加速；
📄 在基于输入的任务中表现出色，如自动语音识别或摘要生成。

辅助生成下的采样

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

从下一个词元的概率分布中采样会导致我们的贪婪助手更频繁地失败，从而减少其延迟优势。然而，我们可以借助大多数基于采样的应用中存在的温度系数来控制下一个词元概率分布的锐利程度。在极端情况下，当温度接近 0 时，采样将近似于贪婪解码，偏向于最可能的词元。在另一个极端，将温度设置为远大于 1 的值时，采样会变得混沌，从均匀分布中抽取。因此，较低的温度更有利于你的助手模型，保留了辅助生成带来的大部分延迟优势，如下所示。

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

未来方向

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

最初在我们的 🤗 Transformers 库中发布，用于 .generate() 函数，我们期望在整个 HF Mirror 宇宙中提供它。其实现也是完全开源的，因此，如果你正在从事文本生成工作而未使用我们的工具，请随意将其作为参考。

最后，辅助生成重新提出了文本生成中的一个关键问题。该领域的发展一直受限于这样一个约束：对于给定的模型，所有新词元都是固定计算量的结果。在纯自回归方式下，每个均匀前向传递产生一个词元。这篇博客文章强化了这样一种观点：情况不应如此：生成输出的大部分也可以由规模仅为其几分之一的模型同样生成。为此，我们需要新的模型架构和解码方法——我们很期待看到未来会带来什么！

引用

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

致谢

我要感谢 Sylvain Gugger、Nicolas Patry 和 Lewis Tunstall 分享了许多宝贵的建议来改进这篇博客文章。最后，感谢 Chunte Lee 设计了您在我们网页上看到的华丽封面。