使用 ONNX Runtime 加速 LLaMA-2 推理

有兴趣更快地运行 Llama2 吗？ 让我们探索 ONNX Runtime 如何加速您的 Llama2 变体，以实现更快的推理！

现在，借助 ONNX Runtime 的最先进的融合和内核优化，您可以体验到 7B、13B 和 70B 模型推理性能的显著提升，速度提升高达 3.8 倍。 这篇博客详细介绍了性能增强、深入探讨了 ONNX Runtime 融合优化、多 GPU 推理支持，并指导您如何在跨平台推理中利用 ONNX Runtime 的跨平台优势。 这是即将推出的一系列博客中的第一篇，这些博客将涵盖使用 ONNX Runtime 量化更新实现高效内存使用以及跨平台使用场景的其他方面。

背景： Llama2 和微软

Llama2 是 Meta 最先进的开源 LLM，参数规模从 7B 到 70B（7B、13B、70B）。 微软和 Meta 在 2023 年 7 月宣布了他们在 Azure 和 Windows 上的 AI 合作。 作为公告的一部分，Llama2 被添加到 Azure AI 模型目录中，该目录是基础模型的中心，旨在帮助开发人员和机器学习 (ML) 专业人员轻松发现、评估、自定义和大规模部署预构建的大型 AI 模型。

ONNX Runtime 使您可以轻松地将这种生成式 AI 模型的强大功能集成到您的应用程序和服务中，并通过改进的优化来提高推理速度并降低成本。

通过新的 ONNX Runtime 优化实现更快的推理

作为新的 1.16.2 版本的一部分，ONNX Runtime 现在为 Llama2 提供了多项内置优化，包括图融合和内核优化。 下面列出了 CUDA FP16 提示延迟下，与 PyTorch 编译模式下的 Hugging Face (HF) Llama2 变体相比，推理速度的提升。 下面显示的端到端吞吐量或实际吞吐量定义为批大小 *（提示长度 + 令牌生成长度）/ 实际延迟，其中实际延迟 = 从端到端运行的延迟，令牌生成长度 = 256 个生成的令牌。 与 PyTorch 编译相比，E2E 吞吐量提高了 2.4 倍 (13B) 和 1.8 倍 (7B)。 对于更高的批大小、序列长度对（如 (16, 2048)），PyTorch eager 会超时，而 ORT 显示出比编译模式更好的性能。

延迟和吞吐量

下图显示了 CUDA FP16 上 Llama2 7B 模型的 ONNX Runtime 和 PyTorch 变体之间的延迟比较。 此处的延迟定义为完成模型的一次传递以生成 logits 并同步输出所需的时间。

下面的令牌生成吞吐量是前 256 个生成令牌的平均吞吐量。 与 PyTorch 编译模式相比，我们看到令牌生成吞吐量提高了约 1.3 倍 (7B) 和约 1.5 倍 (13B)。

有关这些指标的更多详细信息，请访问此处。

使用多 GPU 推理的 ONNX Runtime

ONNX Runtime 支持多 GPU 推理，以支持服务大型模型。 即使在 FP16 精度下，LLaMA-2 70B 模型也需要 140GB。 加载模型需要多个 GPU 进行推理，即使使用强大的 NVIDIA A100 80GB GPU 也是如此。

ONNX Runtime 在 70B 模型上应用了 Megatron-LM 张量并行性，将原始模型权重拆分到不同的 GPU 上。 70B 模型上的 Megatron 分片将具有 FP16 精度的 PyTorch 模型分片为 4 个分区，将每个分区转换为 ONNX 格式，然后在转换后的 ONNX 模型上应用新的 ONNX Runtime 图融合。 70B 模型在批大小为 1 时，令牌生成吞吐量约为每秒 30 个令牌，并且通过这些优化，端到端吞吐量从小序列长度开始为 30 tps。 您可以在此处找到其他示例脚本。

ONNX Runtime 优化

ONNX Runtime 用于优化的技术（例如图融合）适用于最先进的模型。 随着这些模型变得越来越复杂，用于应用图融合的技术也会进行调整，以适应额外的复杂性。 例如，ONNX Runtime 现在支持自动模式匹配，而不是手动匹配图中的融合模式。 融合机会可以通过将大型子图导出为函数，然后根据函数的规范进行模式匹配来识别，而不是手动检测大型子图并匹配它们形成的许多路径。

作为一个具体的例子，图 6 是构成旋转嵌入计算的节点示例。 由于需要验证的路径数量，针对此子图进行模式匹配非常繁琐。 通过将其导出为函数，图的父视图将仅显示输入和输出，并将所有这些节点表示为单个运算符。

这种方法使维护和支持未来版本的旋转嵌入计算变得更加容易，因为模式匹配仅依赖于运算符的输入和输出，而不是其内部语义表示。 它还允许类似模型（如 GPT-NeoX、Falcon、Mistral、Zephyr 等）中其他现有旋转嵌入实现通过最小或无需更改即可进行模式匹配和融合。

ONNX Runtime 还增加了对 GroupQueryAttention (GQA) 运算符的支持，该运算符利用新的 Flash Attention V2 算法及其优化的内核来高效地计算注意力。 GQA 运算符支持过去-现在缓冲区共享，在过去键/值缓存（past KV cache）和现在键/值缓存（present KV cache）之间共享。 通过将现在的 KV 缓存绑定到过去的 KV 缓存，无需为两个缓存分配单独的设备上内存。 相反，过去的 KV 缓存可以预先分配足够的设备上内存，以便在推理期间无需请求新的设备上内存。 这减少了在计算密集型工作负载期间 KV 缓存变大时的内存使用量，并通过消除设备上内存分配请求来降低延迟。 过去-现在缓冲区共享可以启用或禁用，而无需更改 ONNX 模型，从而为最终用户提供更大的灵活性来决定哪种方法最适合他们。

除了这些融合和内核优化之外，ONNX Runtime 还减少了模型的内存使用量。 除了量化改进（将在以后的文章中介绍）之外，ONNX Runtime 将每个旋转嵌入中使用的余弦和正弦缓存的大小压缩了 50%。 然后，ONNX Runtime 中运行旋转嵌入计算的计算内核可以识别此格式，并使用其并行化实现来更有效地计算旋转嵌入，并减少内存使用量。 旋转嵌入计算内核还支持交错和非交错格式，以分别支持 Microsoft 版本的 LLaMA-2 和 Hugging Face 版本的 LLaMA-2，同时共享相同的计算。

这些优化适用于 Hugging Face 版本（模型以 -hf 结尾）和 Microsoft 版本。 您可以从Microsoft 的 LLaMA-2 ONNX 存储库下载优化的 HF 版本。 请继续关注即将推出的更新的 Microsoft 版本！

使用 Olive 优化您自己的模型

Olive 是一种硬件感知模型优化工具，它结合了先进的技术，例如模型压缩、优化和编译。 我们已通过 Olive 提供 ONNX Runtime 优化，因此您可以通过简单的体验简化给定硬件的整个优化过程。

这是一个使用 Olive 进行 Llama2 优化的示例，它利用了本博客中重点介绍的 ONNX Runtime 优化。 不同的优化流程满足各种需求。 例如，您可以根据您的精度容差，灵活地为 CPU 和 GPU 推理选择不同的量化数据类型。 此外，您可以使用客户端 GPU 上的 Olive-QLoRa 微调您自己的 Llama2 模型，并使用 ONNX Runtime 优化执行推理。

使用示例

这是一个示例笔记本，向您展示了如何在您的应用程序中使用上述 ONNX Runtime 优化的端到端示例。

结论

本博客中讨论的进展提供了使用 ONNX Runtime 更快的 Llama2 推理，为 AI 应用和研究提供了令人兴奋的可能性。 凭借改进的性能和效率，创新的前景广阔，我们热切期待由充满活力的开发者社区使用 Llama2 和 ONNX Runtime 构建的新应用。 请继续关注更多更新！