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

由于专门的内置指令，新一代 CPU 在机器学习 (ML) 推理方面提供了显著的性能改进。结合其灵活性、高开发速度和低运营成本，这些通用处理器为其他现有硬件解决方案提供了替代的 ML 推理解决方案。

AWS、Arm、Meta 等公司帮助优化了 PyTorch 2.0 推理在基于 Arm 的处理器上的性能。因此，我们很高兴地宣布，对于 PyTorch 2.0，基于 Arm 的 AWS Graviton 实例推理性能对于 ResNet-50 而言，比以前的 PyTorch 版本速度提高了 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：在不同 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”）支持 Arm® 架构计算库 (ACL) GEMM 内核，适用于 AArch64 平台。优化主要针对 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() 函数以获取总体推理时间。如图所示，在基于 Graviton3 的 c7g 实例上，PyTorch 2.0 中的 ResNet-50 模型推理时间比 PyTorch 1.13 快约 3 倍。

图 6：分析器跟踪视图：PyTorch 1.13 和 PyTorch 2.0 上的前向传播墙钟持续时间

下一个图表是“操作符”视图，它显示了 PyTorch 操作符列表及其执行时间。与前面的跟踪视图类似，“操作符”视图显示，在基于 Graviton3 的 c7g 实例上，PyTorch 2.0 中的 ResNet-50 模型操作符主机持续时间比 PyTorch 1.13 快约 3 倍。

图 7：分析器操作符视图：PyTorch 1.13 和 PyTorch 2.0 上的前向操作符主机持续时间

基准测试 Hugging Face 模型

您可以使用 Amazon SageMaker 推理推荐器 实用程序自动在不同实例之间进行性能基准测试。通过推理推荐器，您可以找到为给定 ML 模型提供最佳性能和最低成本的实时推理端点。我们通过将模型部署到生产端点，使用推理推荐器笔记本收集了上述数据。有关推理推荐器的更多详细信息，请参阅 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（应用人工智能工程主管）对本博客的指导。

关于作者

Sunita Nadampalli 是 AWS 的机器学习工程师和软件开发经理。

Ankith Gunapal 是 Meta (PyTorch) 的 AI 合作伙伴工程师。