由于内置了专门指令，新一代 CPU 在机器学习 (ML) 推理方面提供了显著的性能提升。凭借其灵活性、高开发速度和低运营成本，这些通用处理器为其他现有硬件解决方案提供了替代的 ML 推理解决方案。
AWS、Arm、Meta 和其他公司帮助优化了 PyTorch 2.0 在 Arm 处理器上的推理性能。因此，我们很高兴地宣布，与之前的 PyTorch 版本相比，PyTorch 2.0 在基于 Arm 的 AWS Graviton 实例上的推理性能提升显著，其中 ResNet-50 速度提升高达 3.5 倍，BERT 速度提升高达 1.4 倍，这使得基于 Graviton 的实例成为 AWS 上运行这些模型时速度最快的计算优化型实例（参见下图）。
图 1：从 PyTorch 1.13 版升级到 2.0 版所实现的相对速度提升（越高越好）。性能在 c7g.4xlarge 实例上测量。
如下图所示，与同类基于 x86 的计算优化型 Amazon EC2 实例相比，使用基于 Graviton3 的 c7g 实例在 Torch Hub ResNet-50 和多个 Hugging Face 模型上进行 PyTorch 推理可节省高达 50% 的成本。对于该图，我们首先测量了五种实例类型的每百万次推理成本。然后，我们将每百万次推理成本结果归一化到 c5.4xlarge 实例，这是图表 Y 轴上“1”的基准测量值。
图 2：PyTorch 推理在不同 AWS 实例上运行的相对成本（越低越好）。
来源：关于 Graviton PyTorch 2.0 推理性能的 AWS ML 博客。
与前面的推理成本比较图类似，下图显示了相同五种实例类型的模型 p90 延迟。我们将延迟结果归一化到 c5.4xlarge 实例，这是图表 Y 轴上“1”的基准测量值。基于 Graviton3 的 c7g.4xlarge 的模型推理延迟比在 c5.4xlarge、c6i.4xlarge 和 c6a.4xlarge 上测得的延迟好高达 50%。
图 3：PyTorch 推理在不同 AWS 实例上运行的相对延迟 (p90)（越低越好）。
来源：关于 Graviton PyTorch 2.0 推理性能的 AWS ML 博客。
优化细节
PyTorch 通过 oneDNN 后端（之前称为“MKL-DNN”）支持适用于 AArch64 平台的 Compute Library for the Arm® Architecture (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：调用序列图，显示 Compute Library for the Arm® Architecture (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 profiler 分析 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 来查看 profiler 的结果并分析模型性能。
按如下方式安装 PyTorch Profiler Tensorboard 插件
pip install torch_tb_profiler
使用以下命令启动 tensorboard
tensorboard --logdir=./logs
在浏览器中打开以下内容以查看 profiler 输出。profiler 支持“概览”、“算子”、“跟踪”和“模块”视图，以深入了解推理执行情况。
http://localhost:6006/#pytorch_profiler
下图是 profiler 的“跟踪”视图，它显示了调用堆栈以及每个函数的执行时间。在 profiler 中，我们选择 forward() 函数以获取总推理时间。如图所示，在基于 Graviton3 的 c7g 实例上，ResNet-50 模型在 PyTorch 2.0 中的推理时间比 PyTorch 1.13 中快约 3 倍。
图 6：Profiler 跟踪视图：PyTorch 1.13 和 PyTorch 2.0 上的前向传播墙钟时间
下图是“算子”视图，显示了 PyTorch 算子列表及其执行时间。与前面的跟踪视图类似，“算子”视图显示，在基于 Graviton3 的 c7g 实例上，ResNet-50 模型在 PyTorch 2.0 中的算子主机持续时间比 PyTorch 1.13 中快约 3 倍。
图 7：Profiler 算子视图：PyTorch 1.13 和 PyTorch 2.0 上的前向算子主机持续时间
Hugging Face 模型基准测试
您可以使用 Amazon SageMaker Inference Recommender 实用工具来自动化跨不同实例的性能基准测试。借助 Inference Recommender，您可以找到以最低成本为给定 ML 模型提供最佳性能的实时推理端点。我们通过将模型部署到生产端点，使用 Inference Recommender notebook 收集了上述数据。有关 Inference Recommender 的更多详细信息，请参阅 amazon-sagemaker-examples GitHub 仓库。我们为本文档测试了以下模型：ResNet50 图像分类、DistilBERT 情感分析、RoBERTa 填充掩码和RoBERTa 情感分析。
结论
对于 PyTorch 2.0，基于 Graviton3 的 C7g 实例是用于推理的最具成本效益的计算优化型 Amazon EC2 实例。这些实例可在 SageMaker 和 Amazon EC2 上使用。AWS 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（应用人工智能工程负责人）在本博客方面的指导。
关于作者
Sunita Nadampalli 是 AWS 的 ML 工程师和软件开发经理。
Ankith Gunapal 是 Meta (PyTorch) 的 AI 合作伙伴工程师。
