在 AWS Graviton 处理器上通过 ONNX Runtime 加速 NLP 推理

关键要点

在机器学习ML领域，推理速度是应用性能的关键。使用开放源代码框架 ONNX 提供了跨多种框架、操作系统和硬件平台的互操作性。ONNX Runtime 是用于推理和训练的运行时引擎。

AWS 的 Graviton3 处理器专为 ML 负载而优化，支持 bfloat16、可扩展矢量扩展SVE和矩阵乘法MMLA指令。通过使用 bfloat16 加速的 SGEMM 内核以及 int8 MMLA 加速的量化 GEMMQGEMM内核，ONNX 在 AWS Graviton3 的多种自然语言处理模型上提高了推理性能，fp32 推理性能提升了高达 65，而 int8 量化推理则提高了高达 30。自版本 v1170 起，ONNX Runtime 开始支持这些优化内核。

优化的 GEMM 内核

ONNX Runtime 默认使用 Microsoft 线性代数子例程MLAS作为深度学习运算符的执行提供者。AWS Graviton3 基础的 EC2 实例如 c7g、m7g、r7g、c7gn 和 Hpc7g 实例支持 bfloat16 格式和 MMLA 指令，从而加快深度学习运算符的执行。这些指令可提高 SIMD 硬件利用率，并使端到端推理延迟减少多达 165 倍，相对于基于 armv8 DOT 产品指令的内核。

AWS 团队实现了针对 bfloat16 快速数学和 int8 量化的 GEMM 的 MLAS 内核，利用 BFMMLA、SMMLA 和 UMMLA 指令，这些指令的矩阵乘法吞吐量高于 DOT 指令。bfloat16 的支持使得使用 bfloat16、fp32 以及自动混合精度AMP训练的模型得以高效部署，而无需量化。以下图示展示了优化的 GEMM 内核已集成到 ONNX Runtime CPU 执行提供者作为 MLAS 内核。

ONNX 软件栈

下图展示了 ONNX 软件栈的架构，突出显示了橙色部分在 AWS Graviton3 平台上优化推理性能的各个组件。

ONNX Runtime 流程图

下图显示了 ONNX Runtime 执行提供者的流程，突出显示了橙色部分在 AWS Graviton3 平台上优化推理性能的组件。

启用优化

这些优化是 ONNX Runtime 在 1170 版本中的一部分，并自双击 onnxruntime1170 python wheels 和 conda1170 packages 开始可用。优化的 int8 内核默认启用，可以自动应用于 AWS Graviton3 处理器。而 bfloat16 快速数学内核默认未启用，需在 ONNX Runtime 中设置以下会话选项以启用它们：

对于 C 应用程序

cppSessionOptions so soconfigoptionsAddConfigEntry(kOrtSessionOptionsMlasGemmFastMathArm64Bfloat16 1)

对于 Python 应用程序

pythonsessoptions = onnxruntimeSessionOptions()sessoptionsaddsessionconfigentry(mlasenablegemmfastmatharm64bfloat16 1)

基准测试结果

我们首先测量了未使用优化的 fp32 模型的推理吞吐量每秒查询数，以 ONNX Runtime 1160 为基准。随后我们将内核在 ONNX Runtime 1171 中对 bfloat16 快速数学内核的改进进行对比。下图显示了归一化的结果。您可以看到，对于 BERT、RoBERTa 和 GPT2 模型，吞吐量提升高达65。推理延迟也表现出类似的改善。

我们对未使用优化的 int8 量化模型也进行了相同的吞吐量测量使用 ONNX Runtime 1160，并将其与优化后的 MMLA 内核在同一模型推理中的改进进行了比较。下图显示了归一化的结果。同样，在 BERT、RoBERTa 和 GPT2 模型上，吞吐量提升可达到 30。推理延迟表现出类似的进步。

基准测试设置

我们使用一台 AWS Graviton3 基础的 c7g4xl EC2 实例以及 Ubuntu 2204 作为 AMI，展示 ONNX Runtime 中优化的 GEMM 内核带来的性能提升。实例和 AMI 详情如下：

plaintext实例 c7g4xl区域 uswest2AMI ami0a24e6e101933d294 (Ubuntu 2204/Jammy with 6501014aws kernel)

ONNX Runtime 仓库提供了基于变换器模型的推理基准测试脚本。该脚本支持多种模型、框架和格式。我们选择了基于 PyTorch 的 BERT、RoBERTa 和 GPT 模型来涵盖常见的语言任务，如文本分类、情感分析和预测遮罩单词。模型涵盖编码器和解码器变换器架构。

下面的代码列出了使用 ONNX Runtime 基准测试脚本运行 fp32 模型的步骤，启用了 bfloat16 快速数学模式和 int8 量化模式。该脚本会下载模型，将其导出到 ONNX 格式，量化为 int8 以进行 int8 推理，并在不同的序列长度和批次大小上进行推理。脚本成功完成时，会打印出推理吞吐量每秒查询数和延迟毫秒，以及系统配置。有关更多详情，请参考 ONNX Runtime 基准测试脚本。

bash

安装 Python

sudo aptget updatesudo aptget install y python3 python3pip

将 pip3 升级到最新版本

python3 m pip install upgrade pip

安装 onnx 和 onnx runtime

注意：在收集本文数据时，我们使用了 1171 版本而不是 1170

python3 m pip install onnx==1150 onnxruntime==1171

安装依赖项

python3 m pip install transformers==4381 torch==221 psutil==598

克隆 onnxruntime 仓库以获取基准测试脚本

git clone recursive https//githubcom/microsoft/onnxruntimegitcd onnxruntimegit checkout 430a086f22684ad0020819dc3e7712f36fe9f016cd onnxruntime/python/tools/transformers

运行 bertlarge fp32 推理，并启用 bfloat16 快速数学模式

python3 benchmarkpy m bertlargeuncased p fp32 enablearm64bfloat16fastmathmlasgemm

运行 bertbase fp32 推理，并启用 bfloat16 快速数学模式

python3 benchmarkpy m bertbasecased p fp32 enablearm64bfloat16fastmathmlasgemm

运行 robertabase fp32 推理，并启用 bfloat16 快速数学模式

python3 benchmarkpy m robertabase p fp32 enablearm64bfloat16fastmathmlasgemm

运行 gpt2 fp32 推理，并启用 bfloat16 快速数学模式

python3 benchmarkpy m gpt2 p fp32 enablearm64bfloat16fastmathmlasgemm

运行 bertlarge int8 量化推理

python3 benchmarkpy m bertlargeuncased p int8

运行 bertbase int8 量化推理

python3 benchmarkpy m bertbasecased p int8

运行 robertabase int8 量化推理

python3 benchmarkpy m robertabase p int8

运行 gpt2 int8 量化推理

python3 benchmarkpy m gpt2 p int8

结论

在本文中，我们介绍了如何在 AWS Graviton3 基础的 EC2 实例上运行 ONNX Runtime 推理，以及如何配置实例使用优化的 GEMM 内核。我们还展示了由此产生的加速效果。希望您能试一试！

如果您发现类似性能提升未在 AWS Graviton 上得到体现的使用案例，请在 AWS Graviton 技术指南的 GitHub 上开具问题，让我们知道。

关于作者

Sunita Nadampalli 是 AWS 的软件开发经理。她负责机器学习和高性能计算HPC负载的 Graviton 软件性能优化。她热衷于开源软件开发，并致力于提供高性能和可持续的软件解决方案，通过 Arm SoC 实现。

标签 机器学习 自然语言处理