使用 AWS Graviton 处理器优化 PyTorch 2.0 推理

与前面的推理成本比较图类似，下面的图表显示了相同五种实例类型的模型 p90 延迟。我们将延迟结果归一化到 c5.4xlarge 实例，这是图表 Y 轴上“1”的基准度量。c7g.4xlarge（AWS Graviton3）模型推理延迟比在 c5.4xlarge、c6i.4xlarge 和 c6a.4xlarge 上测量的延迟提高了高达 50%。

优化细节

PyTorch 通过 oneDNN 后端（之前称为“MKL-DNN”）支持 Arm® 架构的 Compute Library (ACL) GEMM 内核，用于 AArch64 平台。优化主要针对 PyTorch ATen CPU BLAS、ACL 内核的 fp32 和 bfloat16，以及 oneDNN 原语缓存。没有前端 API 变更，因此无需在应用程序级别进行任何更改即可在基于 Graviton3 的实例上启用这些优化。

PyTorch 级别优化

我们扩展了 ATen CPU BLAS 接口，通过 oneDNN 后端加速更多算子和张量配置，以加速 aarch64 平台。以下图表（橙色部分）突出了优化组件，这些组件提高了 aarch64 平台上的 PyTorch 推理性能。

图像 4：PyTorch 软件栈，突出显示（橙色）针对 AArch64 平台推理性能优化而优化的组件

ACL 内核和 BFloat16 FPmath 模式

ACL 库为 Neon 和 SVE 优化了 GEMM 内核，支持 fp32 和 bfloat16 格式：这些内核提高了 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®架构的 Compute Library（ACL）GEMM 内核如何集成到 oneDNN 后端

如何利用优化

安装来自官方仓库的 PyTorch 2.0 轮子，并设置环境变量以启用额外的优化。

# 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 分析基于 Graviton3 的 c7g 实例上 ResNet-50 的推理性能。我们使用 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

在浏览器中启动以下内容以查看分析器输出。分析器支持“概述”、“操作”、“跟踪”和“模块”视图，以深入了解推理执行。

http://localhost:6006/#pytorch_profiler

以下图表是分析器的“跟踪”视图，显示了函数的调用栈以及每个函数的执行时间。在分析器中，我们选择了 forward() 函数以获取整体推理时间。如图所示，在 Graviton3-based 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 笔记本收集了前面的数据，通过在生产端点部署模型来实现。有关 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 人工智能合作伙伴工程师。