新一代 CPU 由于其内置的专用指令，在机器学习 (ML) 推理方面提供了显著的性能改进。结合其灵活性、快速开发速度和低运营成本，这些通用处理器为现有的其他硬件解决方案提供了一种替代的 ML 推理解决方案。
AWS、Arm、Meta 等公司帮助优化了 PyTorch 2.0 在基于 Arm 的处理器上的推理性能。因此，我们很高兴地宣布，基于 Arm 的 AWS Graviton 实例在 PyTorch 2.0 上的推理性能，对于 ResNet-50 而言，比之前的 PyTorch 版本快了 3.5 倍，对于 BERT 而言，快了 1.4 倍，这使得基于 Graviton 的实例成为 AWS 上这些模型最快的计算优化实例（见下图）。
图 1：从 PyTorch 1.13 版升级到 2.0 版所实现的相对速度提升（越高越好）。性能在 c7g.4xlarge 实例上测量。
如下图所示，我们测量了使用基于 Graviton3 的 c7g 实例进行 PyTorch 推理，在 Torch Hub ResNet-50 和多个 Hugging Face 模型上，与可比较的基于 x86 的计算优化 Amazon EC2 实例相比，可节省高达 50% 的成本。对于该图，我们首先测量了五种实例类型每百万次推理的成本。然后，我们将每百万次推理的成本结果标准化为 c5.4xlarge 实例，该实例在图表的 Y 轴上是“1”的基线测量值。
图 2：在不同 AWS 实例上运行 PyTorch 推理的相对成本（越低越好）。
来源：AWS ML 博客关于 Graviton PyTorch2.0 推理性能。
与前面的推理成本比较图类似，下图显示了相同五种实例类型的模型 p90 延迟。我们将延迟结果标准化为 c5.4xlarge 实例，该实例在图表的 Y 轴上是“1”的基线测量值。c7g.4xlarge (AWS Graviton3) 模型推理延迟比在 c5.4xlarge、c6i.4xlarge 和 c6a.4xlarge 上测量的延迟提高了高达 50%。\
图 3：在不同 AWS 实例上运行 PyTorch 推理的相对延迟（p90）（越低越好）。
来源：AWS ML 博客关于 Graviton PyTorch2.0 推理性能。
优化详情
PyTorch 通过 oneDNN 后端（以前称为“MKL-DNN”）支持适用于 AArch64 平台的 Arm® 架构计算库 (ACL) GEMM 内核。优化主要针对 PyTorch ATen CPU BLAS、用于 fp32 和 bfloat16 的 ACL 内核以及 oneDNN 原始缓存。没有前端 API 更改，因此在应用程序级别无需进行任何更改即可在基于 Graviton3 的实例上启用这些优化。
PyTorch 级别优化
我们扩展了 ATen CPU BLAS 接口，以通过 oneDNN 后端加速适用于 aarch64 平台的更多操作符和张量配置。下图突出显示（橙色）了在 aarch64 平台上改进 PyTorch 推理性能的优化组件。
图 4：PyTorch 软件堆栈，突出显示（橙色）为提高 AArch64 平台上的推理性能而优化的组件
ACL 内核和 BFloat16 FPmath 模式
ACL 库为 fp32 和 bfloat16 格式提供 Neon 和 SVE 优化的 GEMM 内核：这些内核提高了 SIMD 硬件利用率并减少了端到端推理延迟。Graviton3 中的 bfloat16 支持允许高效部署使用 bfloat16、fp32 和自动混合精度 (AMP) 训练的模型。标准 fp32 模型通过 oneDNN FPmath 模式使用 bfloat16 内核，无需模型量化。与没有 bfloat16 FPmath 支持的现有 fp32 模型推理相比，它们提供了高达两倍的性能。有关 ACL GEMM 内核支持的更多详细信息，请参阅 Arm Compute Library github。
原始缓存
下图显示了 ACL 操作符如何集成到 oneDNN 后端。如图所示，ACL 对象被视为 oneDNN 资源，而不是原始对象。这是因为 ACL 对象是有状态且可变的。由于 ACL 对象被视为资源对象，因此它们无法使用 oneDNN 中支持的默认原始缓存功能进行缓存。我们在 ideep 操作符级别为“卷积”、“矩阵乘法”和“内积”操作符实现了原始缓存，以避免冗余的 GEMM 内核初始化和张量分配开销。
图 5：调用序列图显示了 Arm® 架构计算库 (ACL) GEMM 内核如何集成到 oneDNN 后端
如何利用这些优化
从官方存储库安装 PyTorch 2.0 wheel 并设置环境变量以启用额外的优化。
# Install Python
sudo apt-get update
sudo apt-get install -y python3 python3-pip
# Upgrade pip3 to the latest version
python3 -m pip install --upgrade pip
# Install PyTorch and extensions
python3 -m pip install torch
python3 -m pip install torchvision torchaudio torchtext
# Turn on Graviton3 optimization
export DNNL_DEFAULT_FPMATH_MODE=BF16
export LRU_CACHE_CAPACITY=1024
运行推理
您可以使用 PyTorch torchbench 来衡量 CPU 推理性能改进，或比较不同的实例类型。
# Pre-requisite:
# pip install PyTorch2.0 wheels and set the above mentioned environment variables
# Clone PyTorch benchmark repo
git clone https://github.com/pytorch/benchmark.git
# Setup ResNet-50 benchmark
cd benchmark
python3 install.py resnet50
# Install the dependent wheels
python3 -m pip install numba
# Run ResNet-50 inference in jit mode. On successful completion of the inference runs,
# the script prints the inference latency and accuracy results
python3 run.py resnet50 -d cpu -m jit -t eval --use_cosine_similarity
性能分析
现在，我们将使用 PyTorch 性能分析器分析 ResNet-50 在基于 Graviton3 的 c7g 实例上的推理性能。我们使用 PyTorch 1.13 和 PyTorch 2.0 运行以下代码，并在测量性能之前运行推理进行几次迭代作为预热。
# Turn on Graviton3 optimization
export DNNL_DEFAULT_FPMATH_MODE=BF16
export LRU_CACHE_CAPACITY=1024
import torch
from torchvision import models
sample_input = [torch.rand(1, 3, 224, 224)]
eager_model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model = torch.jit.script(eager_model, example_inputs=[sample_input, ])
model = model.eval()
model = torch.jit.optimize_for_inference(model)
with torch.no_grad():
# warmup runs
for i in range(10):
model(*sample_input)
prof = torch.profiler.profile(
on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'), record_shapes=True, with_stack=True)
# profile after warmup
prof.start()
model(*sample_input)
prof.stop()
我们使用 tensorboard 来查看性能分析器的结果并分析模型性能。
如下安装 PyTorch Profiler Tensorboard 插件：
pip install torch_tb_profiler
使用以下命令启动 tensorboard：
tensorboard --logdir=./logs
在浏览器中启动以下内容以查看性能分析器输出。性能分析器支持“概述”、“操作符”、“跟踪”和“模块”视图，以深入了解推理执行。
https://:6006/#pytorch_profiler
下图是性能分析器的“跟踪”视图，它显示了调用堆栈以及每个函数的执行时间。在性能分析器中，我们选择了 forward() 函数来获取总推理时间。如图所示，在 PyTorch 2.0 中，ResNet-50 模型在基于 Graviton3 的 c7g 实例上的推理时间比 PyTorch 1.13 快了约 3 倍。
图 6：性能分析器跟踪视图：PyTorch 1.13 和 PyTorch 2.0 上的前向传播总时长
下图是“操作符”视图，它显示了 PyTorch 操作符列表及其执行时间。与前面的跟踪视图类似，操作符视图显示，在 PyTorch 2.0 中，ResNet-50 模型在基于 Graviton3 的 c7g 实例上的操作符主机持续时间比 PyTorch 1.13 快了约 3 倍。
图 7：性能分析器操作符视图：PyTorch 1.13 和 PyTorch 2.0 上的前向操作符主机持续时间
基准测试 Hugging Face 模型
您可以使用 Amazon SageMaker Inference Recommender 工具来自动化跨不同实例的性能基准测试。借助 Inference Recommender，您可以为给定的 ML 模型找到以最低成本提供最佳性能的实时推理端点。我们通过将模型部署到生产端点，使用 Inference Recommender 笔记本收集了上述数据。有关 Inference Recommender 的更多详细信息，请参阅 amazon-sagemaker-examples GitHub 存储库。我们为本文基准测试了以下模型：ResNet50 图像分类、DistilBERT 情感分析、RoBERTa 填充掩码和 RoBERTa 情感分析。
总结
对于 PyTorch 2.0，基于 Graviton3 的 C7g 实例是用于推理的最具成本效益的计算优化 Amazon EC2 实例。这些实例可在 SageMaker 和 Amazon EC2 上获得。AWS Graviton 技术指南提供了优化库列表和最佳实践，可帮助您在不同工作负载下利用 Graviton 实例实现成本效益。
如果您发现 Graviton 上未观察到类似性能提升的用例，请在 aws-graviton-getting-started github 上提出问题，告知我们。我们将继续添加更多性能改进，使基于 AWS Graviton 的实例成为使用 PyTorch 进行推理最具成本效益和最高效的通用处理器。
致谢
我们谨此感谢 AWS 的 Ali Saidi（高级首席工程师）和 Csaba Csoma（高级软件开发经理），以及 Arm 的 Ashok Bhat（高级产品经理）、Nathan Sircombe（高级工程经理）和 Milos Puzovic（首席软件工程师）在 Graviton PyTorch 推理优化工作中的支持。我们还要感谢 Meta 的 Geeta Chauhan（应用 AI 工程负责人）对本博客的指导。
关于作者
Sunita Nadampalli 是 AWS 的机器学习工程师和软件开发经理。
Ankith Gunapal 是 Meta (PyTorch) 的 AI 合作伙伴工程师。