更快的 LLM 通过推测解码和 AWS Inferentia2 机器学习博客

• 2026-01-27 14:25:14

加速大型语言模型的推理过程

关键要点

最近几年，我们见证了用于解决诸如问答和文本摘要等自然语言处理NLP任务的大型语言模型的规模大幅增加。目前，像 Llama370B 这样的模型参数达到了数百亿，它们通常能产生更好的效果。例如，Llama370B 在阅读理解SQuAD 856 比 764等指标上表现优于较小的 8B 参数版本。因此，客户通常会尝试更大和更新的模型，以构建能够增值的机器学习产品。

然而，模型越大，对计算资源的需求越高，部署成本也相应上升。例如，在 AWS Trainium 上，Llama370B 的每个标记的中位延迟为 214 毫秒，而 Llama38B 则为 47 毫秒。同样，Llama270B 的每个标记的中位延迟为 206 毫秒，Llama27B 则为 37 毫秒。客户必须考虑性能，以确保满足用户需求。在本文中，我们将探讨如何使用 推测抽样 技术使大型语言模型的推理在 AWS Inferentia 和 Trainium 上更加计算高效和具有成本效益。这项技术 提升 LLM 的推理吞吐量和输出标记延迟TPOT。

引言

现代语言模型基于 变换器架构。输入提示首先通过称为上下文编码的技术进行处理，由于其可并行化，运行速度较快。接下来，我们执行自回归标记生成，输出标记是顺序生成的。请注意，在已知前一个标记之前，我们无法生成下一个标记，如图 1 所示。因此，要生成 N 个输出标记，我们需要 N 次通过解码器的串行运行。通过较大的模型如 Llama370B 进行的运行比通过较小模型如 Llama38B 要耗时更长。

图 1 LLM 中的顺序标记生成

从计算角度来看，LLM 中的标记生成是一个受内存带宽限制的过程。模型越大，我们在内存传输上等待的概率就越高。这导致计算单元的低利用率，而无法充分利用可用的浮点操作FLOPS。

推测抽样

推测抽样是一种在保持准确性的同时提高 LLM 推理计算效率的技术。该方法使用较小、更快的草稿模型生成多个标记，然后由较大、更慢的目标模型进行验证。这一步骤以单次通过处理多个标记而非顺序处理，有效地提高了计算效率。增加并行处理的标记数量提升了计算强度，因为可以用相同的权重张量对更多的标记进行乘法运算。这种方式的性能优于通常受内存带宽限制的非推测运行，从而更好地利用硬件资源。

推测过程涉及一个可调窗口 k，目标模型提供一个保证正确的标记，而草稿模型则推测下 k1 个标记。如果草稿模型的标记被接受，处理速度就会加快。如果未被接受，目标模型将接手以确保准确性。

图 2 所有推测标记被接受的情况

图 2 展示了所有推测标记都被接受的情况，从而加快了处理速度。目标模型提供一个保证给出的输出标记，草稿模型多次运行以生成一系列可能的输出标记。这些标记通过目标模型进行验证，并且随后通过概率方法被接受。

图 3 某些推测标记被拒绝的情况

另一方面，图 3 显示了某些标记被拒绝的情况。进行这一推测抽样循环所花费的时间与图 2 相同，但我们得到的输出标记更少。这意味着我们将重复该过程更多次以完成响应，导致整体处理速度减慢。

通过调整窗口大小 k，并了解何时草稿与目标模型可能产生相似结果，我们可以最大化推测抽样的好处。

Llama270B/7B 演示

我们将在 AWS EC2 的 Inferentia2 驱动的 Inf2 实例和 Trainium 驱动的 EC2 Trn1 实例上展示推测抽样的工作原理。我们将使用一个 示例，通过使用 Llama27B 模型作为草稿模型更快速地生成 Llama270B 的文本。示例的演示基于 Llama2 模型，但您也可以对 Llama3 模型采用类似的过程。

加载模型

您可以使用数据类型 bfloat16 加载 Llama2 模型。草稿模型需要按照标准方式加载，如下例所示。参数 npositions 可调，表示您希望允许生成的最大序列长度。当前支持的唯一 batchsize 为 1。我们将在本节稍后解释 tpdegree。

pythondraftmodel = LlamaForSamplingfrompretrained(Llama27b npositions=128 batchsize=1 tpdegree=32 amp=bf16)

目标模型应以类似的方式加载，但需启用推测抽样功能。前面提到的值 k 可以在此处设置。

pythontargetmodel = LlamaForSamplingfrompretrained(Llama270b npositions=128 batchsize=1 tpdegree=32 amp=bf16)targetmodelenablespeculativedecoder(k)

结合这两个模型，权重几乎需要 200 GB 的设备内存，此外还需要数 GB 的键值KV缓存。如果您希望使用 float32 参数的模型，则它们将需要大约 360 GB 的设备内存。请注意，KV 缓存随着序列长度输入标记 尚未生成的标记线性增长。使用 neurontop 可以实时查看内存利用情况。为了满足这些内存要求，我们需要使用最大的 Inf2 实例inf248xlarge或最大的 Trn1 实例trn132xlarge。

由于模型的大小，它们的权重需要使用一种称为 张量并行性 的技术分布在 NeuronCores 上。在提供的示例中，使用 tpdegree 指定每个模型应该使用多少个 NeuronCores。这会影响内存带宽的利用率，这对于标记生成的性能至关重要。较高的 tpdegree 会导致更好的带宽利用和吞吐量。Trn1 的拓扑要求 tpdegree 设置为 1、2、8、16 或 32 的倍数。对于 Inf2，则需设置为 1 或 2 的倍数。

加载模型的顺序也很重要。在为一个模型初始化和分配了一组 NeuronCores 后，您不能对其他模型使用同一组 NeuronCores，除非它们是完全相同的设置。如果您尝试只使用之前初始化的一些 NeuronCores，您将会遇到 nrtloadcollectives global neccomm is already initd 错误。

让我们通过在 trn132xlarge32 个 NeuronCores上进行两个示例来更好地理解这一点。我们将计算每个模型需要多少 NeuronCores。使用的公式是利用 neurontop 观察的模型在内存中的大小除以 16GB每个 NeuronCore 的设备内存。

如果我们使用 bfloat16 运行模型，Llama270B 需要超过 10 个 NeuronCores，Llama27B 需要超过 2 个 NeuronCores。由于拓扑约束，这意味着 Llama270B 至少需要 tpdegree=16。我们可以使用剩余的 16 个 NeuronCores 来服务 Llama27B。然而，由于这两个模型可以在 32 个 NeuronCores 的内存中共同使用，因此我们应为两个模型都设置 tpdegree=32，从而加快每个模型的推理速度。

如果我们使用 float32 运行模型，Llama270B 需要超过 18 个 NeuronCores，Llama27B 需要超过 3 个 NeuronCores。由于拓扑约束，我们必须为 Llama270B 设置 tpdegree=32。这意味着 Llama27B 需要重用同一组 NeuronCores，因此您也需要为 Llama27B 设置 tpdegree=32。

演示步骤

我们将在 transformersneuronx 中使用的解码器是 LlamaForSampling，它适合加载和运行 Llama 模型。您也可以使用 NeuronAutoModelForCausalLM，它将尝试自动检测使用哪个解码器。为了执行推测抽样，我们需要首先创建一个推测生成器，其中包含两个模型和前面描述的值 k。

pythonspecgen = SpeculativeGenerator(draftmodel targetmodel k)

通过调用以下函数来启动推理过程：

pythonspecgensample(inputids=inputtokenids sequencelength=totaloutputlength)

在抽样过程中，有几个超参数例如：temperature、topp 和 topk会影响多个运行之间的输出是否是确定性的。在撰写时，推测抽样的实现会为这些超参数设置默认值。使用这些值，预期在多次运行模型时结果会有随机性，即使使用相同的提示。这是大型语言模型的正常预期行为，因为这能改善它们的质量反应。

当您执行抽样时，将使用基于 DeepMind 提出的推测抽样论文 的默认标记接受器，该方法使用概率方法来接受标记。不过，您也可以实现自定义标记接受器，并在初始化 SpeculativeGenerator 时将其作为 acceptor 参数传入。如果您想要更确定的响应，可以采用这种方法。有关如何编写自己的接受器，请查看 DefaultTokenAcceptor 类的实现。

结论

随着越来越多的开发者希望将 LLMs 集成到他们的应用中，他们面临着选择：使用更大、更昂贵且更慢的模型以获得更高的质量结果，还是选择更小、更经济和更快速的模型，尽管可能降低了回答质量。如今，通过采用 AWS 人工智能AI芯片和推测抽样，开发者不再需要做这样的选择。他们可以利用大型模型的高质量输出以及小型模型的速度和响应能力。

在本文中，我们展示了如何利用新特性推测抽样加速大型模型如 Llama270B的推理过程。

想要亲身体验，请查阅 推测抽样示例，并调整输入提示和 k 参数以观察结果。对于更高级的用例，您可以自行开发标记接受器实现。有关在 Inferentia 和 Trainium 实例上运行模型的更多信息，请参见 AWS Neuron 文档。您还可以访问 repostaws AWS Neuron 频道，与 AWS Neuron 社区讨论您的实验并分享想法。

作者介绍

Syl Taylor 是高效计算领域的专职解决方案架构师。她负责在 EMEA 向客户提供亚马逊 EC2 成本优化和应用性能提升的建议。Syl 曾在 AWS 专业服务从事软件开发和 AI/ML，设计并实施云原生解决方案。她居住在英国，喜欢在大自然中度过时光。

Emir Ayar 是 AWS 原型团队的高级技术解决方案架构师。他专注于协助客户构建机器学习和生成AI解决方案，并实施架构最佳实践。他支持客户实验解决方案架构以实现他们的业务目标，强调敏捷创新和原型设计。他居住在卢森堡，喜欢玩合成器。

标签

AI/ML，亚马逊机器学习，AWS Inferentia，生成式 AI，机器学习