使用 ONNX Runtime 加速 LLaMA-2 推理

有兴趣更快地运行 Llama2 吗？让我们探讨一下 ONNX Runtime 如何助力您的 Llama2 变体实现更快的推理！

借助 ONNX Runtime 顶尖的算子融合与内核优化，您现在可以体验到 7B、13B 和 70B 模型显著的推理性能提升——最高可达 3.8 倍。本博客详细介绍了性能 enhancements，深入探讨了 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 提供了多项内置优化，包括图融合和内核优化。与处于 PyTorch 编译模式（针对 CUDA FP16 提示词延迟）下的 Hugging Face (HF) Llama2 变体相比，其推理加速情况如下所述。下面显示的端到端吞吐量或实际耗时吞吐量定义为：batch size * (prompt length + token generation length) / wall-clock latency，其中 wall-clock latency = 端到端运行的延迟，token generation length = 生成的 256 个 token。与 PyTorch 编译模式相比，E2E 吞吐量提高了 2.4 倍（13B）和 1.8 倍（7B）。对于诸如 (16, 2048) 这样更高的批大小和序列长度组合，PyTorch eager 模式会超时，而 ORT 则表现出比编译模式更好的性能。

延迟与吞吐量

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

下方的 Token 生成吞吐量是前 256 个生成 token 的平均吞吐量。与 PyTorch 编译模式相比，我们看到 Token 生成吞吐量提升了约 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 图融合。在 batch size 为 1 时，70B 模型的 token 生成吞吐量约为每秒 30 个 token，并且通过这些优化，对于较短的序列长度，端到端吞吐量从 30 tps 开始。您可以在此处找到更多示例脚本。

ONNX Runtime 优化

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

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

这种方法使得维护和支持未来版本的旋转嵌入计算变得更加容易，因为模式匹配仅取决于算子的输入和输出，而不取决于其内部语义表示。它还允许对类似模型（如 GPT-NeoX、Falcon、Mistral、Zephyr 等）中旋转嵌入的其他现有实现进行模式匹配和融合，且只需进行极少甚至无需修改。

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

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

这些优化适用于 Hugging Face 版本（以 -hf 结尾的模型）和微软版本。您可以从微软的 LLaMA-2 ONNX 仓库下载优化的 HF 版本。敬请期待即将推出的更新的微软版本！

使用 Olive 优化您自己的模型

Olive 是一款硬件感知的模型优化工具，它整合了模型压缩、优化和编译等先进技术。我们已通过 Olive 提供 ONNX Runtime 优化，以便您能够以简单的体验简化针对特定硬件的整个优化过程。

这里有一个使用 Olive 优化 Llama2 的示例，它利用了本博客中强调的 ONNX Runtime 优化。不同的优化流程可满足各种需求。例如，您可以根据自己的准确率容忍度，在 CPU 和 GPU 推理中灵活选择不同的数据类型进行量化。此外，您还可以在客户端 GPU 上使用 Olive-QLoRa 微调您自己的 Llama2 模型，并使用 ONNX Runtime 优化执行推理。

使用示例

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

结论

本博客中讨论的技术进步通过 ONNX Runtime 实现了更快的 Llama2 推理，为 AI 应用和研究提供了令人兴奋的可能性。随着性能和效率的提高，创新的道路无限宽广，我们热切期待其充满活力的开发者社区利用 Llama2 和 ONNX Runtime 构建出新的应用。敬请关注更多更新！