(Beta) PyTorch 在 AWS Graviton 处理器上的推理性能调优 ¶

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

作者：Sunita Nadampalli

AWS Graviton 是 AWS 设计的一系列基于 ARM 的处理器。AWS Graviton3 处理器针对机器学习（ML）工作负载进行了优化，包括对 bfloat16 ，可伸缩向量扩展（SVE）的支持，以及与 Graviton2 相比，单指令多数据（SIMD）带宽提高了一倍。

PyTorch 提供了针对机器学习运算符（如卷积、矩阵乘法、ReLU 等）的本地引用 ATen 内核。这些运算符可以通过基本线性代数（BLAS）库的平台特定内核实现进行加速。在 AWS Graviton CPU 上，MKLDNN 与 Arm Compute Library（ACL）和 OpenBLAS 库为部分运算符提供了优化的实现。这两个库都集成到了 PyTorch 2.0 版本中。

在本教程中，我们将介绍如何在 AWS Graviton3 CPU（AWS c7g 实例）上使用 bfloa16 内核以及正确的后端选择来实现线性层神经网络的最佳推理性能。

目录 ¶

1. 基本用法

2. 使用 Bfloat16 快速数学内核加速推理

3. 使用 OpenBLAS 提高小批量维度下的推理性能

4. 使用 Linux 透明大页优化内存分配开销

5. 结论

备注

要成功运行本教程并重现下面显示的速度提升数值，您需要一个来自 Graviton3 系列的实例（ c7g/r7g/m7g ）。对于本教程，我们使用了 c7g.xl（4vcpu）实例。

基本用法 ¶

PyTorch 从 2.0 版本开始原生支持 AWS Graviton3 优化。请参阅此博客以获取有关优化的更多详细信息。

1. 通过运行以下命令安装 PyTorch：

   python3 -m pip install torch
   

2. 我们将首先导入所需的依赖项并定义运行设备：

import torch
import torch.nn as nn
from torch.profiler import profile, record_function, ProfilerActivity

# AWS Graviton3 cpu
device = ("cpu")
print(f"Using {device} device")

3. 给定的线性层是多个神经网络的核心，包括 Transformer，因此我们在此演示中采用线性层。我们通过子类化 nn.Module 并初始化 __init__ 中的层来定义我们的神经网络。我们使用典型的大型语言模型参数构建网络，以匹配真实世界场景：

class MyNeuralNetwork(nn.Module):
  def __init__(self):
      super().__init__()
      self.flatten = nn.Flatten()
      self.linear_relu_stack = nn.Sequential(
          nn.Linear(4096, 4096),
          nn.ReLU(),
          nn.Linear(4096, 11008),
          nn.ReLU(),
          nn.Linear(11008, 10),
      )

  def forward(self, x):
      x = self.flatten(x)
      logits = self.linear_relu_stack(x)
      return logits

4. 让我们创建一个实例 MyNeuralNetwork ，并将其移动到设备上：

model = MyNeuralNetwork().to(device)
print(model)

接下来，让我们通过一个 nn.Softmax 模块的实例来获取预测概率：

X = torch.rand(1, 64, 64, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")

输出：

Predicted class: tensor([2])

我们的网络功能已验证。接下来，我们将分析性能。让我们检查两种不同的场景：小批量和大批量维度。

场景 1：更大的批量维度，例如 256：

# warm it up first and loop over multiple times to have enough execution time

X = torch.rand(256, 64, 64, device=device)

with torch.set_grad_enabled(False):
    for _ in range(50):
        model(X) #Warmup
    with profile(activities=[ProfilerActivity.CPU]) as prof:
        with record_function("mymodel_inference"):
            for _ in range(100):
                model(X)

print(prof.key_averages().table(sort_by="self_cpu_time_total"))

以下为默认 PyTorch 配置的分析器输出：

名称	自用 CPU 百分比	自用 CPU	CPU 总百分比	CPU 总	CPU 平均时间	调用次数
aten::addmm	97.61%	15.813 秒	98.61%	15.977 秒	53.255 毫秒	300
aten::clamp_min	1.09%	177.032 毫秒	1.09%	177.032 毫秒	885.160 微秒	200
aten::copy	1.00%	162.054 毫秒	1.00%	162.054 毫秒	540.180 微秒	300
mymodel 推理	0.22%	35.738 毫秒	100.00%	16.201 秒	16.201 秒	1
aten::线性	0.02%	2.955 毫秒	98.66%	15.985 秒	53.282 毫秒	300
aten::t	0.01%	2.421 毫秒	0.03%	5.043 毫秒	16.810 微秒	300
aten::relu	0.01%	2.356 毫秒	1.11%	179.388 毫秒	896.940us	200

自身 CPU 总时间：16.201s

使用 bfloat16 快速数学内核加速推理 ¶

AWS Graviton3 处理器支持 bfloat16 MMLA 指令。Arm 计算库（ACL）为 AWS Graviton 处理器提供优化的 bfloat16 通用矩阵乘法（GEMM）内核，并从 PyTorch 2.0 开始通过 MKLDNN 后端集成到 PyTorch 中。可以使用快速数学 GEMM 内核优化推理性能。默认情况下不启用快速数学模式，因为这些内核在 bfloat16 精度下执行 GEMM，而不是 float ，因此会导致模型推理精度略有下降。然而，精度下降在 cosine similarity 阈值内，该阈值是为 bfloat16 后端在 torchbench 测试套件中定义的，因此对于大多数应用来说是可接受的。要启用快速数学 GEMM 内核，请设置以下环境变量：

$ export DNNL_DEFAULT_FPMATH_MODE=BF16

当您运行上述推理脚本时，应看到以下配置文件输出，其中启用了 MKLDNN 快速数学模式：

名称	自用 CPU 百分比	自用 CPU	CPU 总百分比	CPU 总	CPU 平均时间	调用次数
aten::addmm	95.61%	6.943 秒	97.10%	7.052 秒	23.507 毫秒	300
aten::clamp_min	2.31%	167.653ms	2.31%	167.653ms	838.265us	200
aten::copy	1.48%	107.593ms	1.48%	107.593ms	358.643us	300
mymodel 推理	0.43%	31.167 毫秒	100.00%	7.262 秒	7.262 秒	1
aten::线性	0.04%	2.911ms	97.21%	7.060s	23.533ms	300
aten::t	0.03%	2.414ms	0.07%	4.892ms	16.307us	300
aten::relu	0.03%	2.281ms	2.34%	169.934ms	849.670us	200

自身 CPU 总时间：7.262 秒

这大约是 2x (7.262s vs 16.201s) 性能提升，使用了 bfloat16 快速数学内核。接下来，让我们看看较小的批量维度场景。

场景 2：较小的批量维度，例如，32：

X = torch.rand(32, 64, 64, device=device)
with torch.set_grad_enabled(False):
    for _ in range(50):
        model(X) #Warmup
    with profile(activities=[ProfilerActivity.CPU]) as prof:
        with record_function("mymodel_inference"):
            for _ in range(100):
                model(X)

print(prof.key_averages().table(sort_by="self_cpu_time_total"))

当使用 PyTorch 默认配置运行上述脚本时，你应该看到以下分析器输出：

名称	自用 CPU 百分比	自用 CPU	CPU 总百分比 %	CPU 总	CPU 平均时间	调用次数
aten::addmm	95.51%	5.821 秒	97.04%	5.914 秒	19.713 毫秒	300
aten::clamp_min	2.33%	142.244 毫秒	2.33%	142.244 毫秒	711.220 微秒	200
aten::copy	1.51%	92.322 毫秒	1.51%	92.322 毫秒	307.740 微秒	300
mymodel 推理	0.45%	27.713 毫秒	100.00%	6.094 秒	6.094 秒	1
aten::线性	0.04%	2.495 毫秒	97.16%	5.921 秒	19.736 毫秒	300
aten::t	0.03%	2.131 毫秒	0.07%	4.441 毫秒	14.803 微秒	300
aten::relu	0.03%	1.942 毫秒	2.37%	144.186 毫秒	720.930 微秒	200

自身 CPU 总时间：6.094 秒

启用 MKLDNN 快速数学模式时运行的配置文件输出如下：

$ export DNNL_DEFAULT_FPMATH_MODE=BF16

名称	自用 CPU 百分比	自用 CPU	CPU 总百分比	CPU 总	CPU 平均时间	调用次数
aten::addmm	93.31%	3.848 秒	95.66%	3.944 秒	13.148 毫秒	300
aten::clamp_min	3.43%	141.309ms	3.43%	141.309ms	706.545us	200
aten::copy	2.33%	95.916ms	2.33%	95.916ms	319.720us	300
mymodel 推理	0.67%	27.431 毫秒	100.00%	4.123 秒	4.123 秒	1
aten::线性	0.06%	2.471ms	95.83%	3.951s	13.170ms	300
aten::t	0.05%	2.027 毫秒	0.10%	4.243 毫秒	14.143 微秒	300
aten::relu	0.05%	1.928ms	3.47%	143.237ms	716.185us	200

自身 CPU 总时间：4.123 秒

MKLDNN 快速数学模式在较小的批量维度上大约提供了 1.47 倍（4.123 秒比 6.094 秒）的性能提升。尽管这种提升值得关注，但整体性能仍有提升空间。这是因为来自 oneDNN 和 ACL 后端的运行时开销（权重重排和内核启动时间）超过了 ACL GEMM 内核在较小批量计算中的计算优势。

提升 OpenBLAS 在较小批量维度上的推理性能

通过将较小的形状从 MKLDNN 卸载到 OpenBLAS 后端，可以提高较小批量维度的推理性能。我们正在努力使后端选择自动进行，并具有鲁棒的启发式算法，以供未来的版本使用。在启发式算法实现之前，可以通过增加 MKLDNN 后端选择的阈值，将较小的形状卸载到 OpenBLAS。在下面的示例中，我们使用 64 作为阈值，因此具有 batch dimension of 32 的输入不会被发送到 MKLDNN。相反，它会被发送到 OpenBLAS。

$ export TORCH_MKLDNN_MATMUL_MIN_DIM=64

这里是使用 OpenBLAS 后端的分析器输出：

名称	自用 CPU 百分比	自用 CPU	CPU 总百分比	CPU 总	CPU 平均时间	调用次数
aten::addmm	96.25%	1.958 秒	97.51%	1.984 秒	6.612 毫秒	300
aten::clamp_min	1.28%	26.124ms	1.28%	26.124ms	130.620us	200
aten::copy	1.23%	24.951ms	1.23%	24.951ms	83.170us	300
mymodel 推理	0.86%	17.423 毫秒	100.00%	2.034 秒	2.034 秒	1
aten::线性	0.08%	1.691ms	97.74%	1.988s	6.628ms	300
aten::t	0.07%	1.520ms	0.14%	2.945ms	9.817us	300
aten::relu	0.06%	1.258ms	1.35%	27.382ms	136.910us	200

自身 CPU 总时间：2.034 秒

如上图所示，切换到 OpenBLAS 将性能翻倍（2.034 秒 vs 4.123 秒），与默认的 MKLDNN 后端配置相比。这对于更小的批量维度来说也很重要，例如，对于批量维度为 10 的情况：

X = torch.rand(10, 64, 64, device=device)
with torch.set_grad_enabled(False):
    for _ in range(50):
        model(X) #Warmup
    with profile(activities=[ProfilerActivity.CPU]) as prof:
        with record_function("mymodel_inference"):
            for _ in range(100):
                model(X)

print(prof.key_averages().table(sort_by="self_cpu_time_total"))

以下是在 MKLDNN 快速数学模式下生成的分析器输出：

名称	自用 CPU 百分比	自用 CPU	CPU 总百分比	CPU 总	CPU 平均时间	调用次数
aten::addmm	87.81%	3.613 秒	91.90%	3.781 秒	12.604 毫秒	300
aten::clamp_min	7.18%	295.437 毫秒	7.18%	295.437 毫秒	1.477 毫秒	200
aten::copy	4.07%	167.516 毫秒	4.07%	167.516 毫秒	558.387 微秒	300
mymodel 推理	0.67%	27.708 毫秒	100.00%	4.115 秒	4.115 秒	1
aten::线性	0.06%	2.499 毫秒	92.06%	3.788 秒	12.627 毫秒	300
aten::t	0.05%	1.982 毫秒	0.11%	4.385 毫秒	14.617 微秒	300
aten::relu	0.05%	1.932 毫秒	7.23%	297.369 毫秒	1.487 毫秒	200

自身 CPU 总时间：4.115 秒

以下为使用 OpenBLAS 后端的分析器输出：

$ export TORCH_MKLDNN_MATMUL_MIN_DIM=64

名称	自用 CPU 百分比	自用 CPU	CPU 总百分比	CPU 总	CPU 平均时间	调用次数
aten::addmm	92.66%	1.179s	95.23%	1.211s	4.038ms	300
aten::clamp_min	2.83%	36.060ms	2.83%	36.060ms	180.300us	200
aten::copy	2.52%	32.013ms	2.52%	32.013ms	106.710us	300
mymodel 推理	1.38%	17.521 毫秒	100.00%	1.272 秒	1.272 秒	1
aten::线性	0.14%	1.750ms	95.60%	1.216s	4.054ms	300
aten::t	0.12%	1.475ms	0.24%	3.033ms	10.110us	300
aten::relu	0.10%	1.285ms	2.94%	37.345ms	186.725us	200

自身 CPU 总时间：1.272 秒

我们观察到通过适当调整后端阈值，性能提升了 3.2 倍（1.272 秒比 4.115 秒）

使用 Linux 透明大页面（THP）优化内存分配开销 ¶

我们还观察到，对于这些较大的网络，张量内存分配占据了推理延迟的很大一部分。这可以通过启用 PyTorch C10 内存分配器的 Linux 透明大页分配来优化。目前该功能默认未启用，因为它会略微增加内存占用。要启用它，请设置以下环境变量：

$ export THP_MEM_ALLOC_ENABLE=1

对于 256 的批次维度以及启用 MKLDNN 快速数学模式：

X = torch.rand(256, 64, 64, device=device)
with torch.set_grad_enabled(False):
    for _ in range(50):
        model(X) #Warmup
    with profile(activities=[ProfilerActivity.CPU]) as prof:
        with record_function("mymodel_inference"):
            for _ in range(100):
                model(X)

print(prof.key_averages().table(sort_by="self_cpu_time_total"))

以下是在启用 THP 内存分配的情况下，分析器的输出：

名称	自用 CPU 百分比	自用 CPU	CPU 总百分比	CPU 总	CPU 平均时间	调用次数
aten::addmm	91.31%	6.115 秒	94.39%	6.321 秒	21.069 毫秒	300
aten::clamp_min	4.82%	322.568 毫秒	4.82%	322.568 毫秒	1.613 毫秒	200
aten::copy	3.06%	204.602 毫秒	3.06%	204.602 毫秒	682.007 微秒	300
mymodel 推理	0.61%	40.777 毫秒	100.00%	6.697 秒	6.697 秒	1
aten::线性	0.05%	3.082 毫秒	94.51%	6.329 秒	21.097 毫秒	300
aten::relu	0.04%	2.547 毫秒	4.85%	325.115 毫秒	1.626 毫秒	200

自身 CPU 总时间：6.697 秒

在上述已优化的 MKLDNN 快速数学模式下，这又是一个额外的 1.08 倍或 8%（6.697 秒比 7.262 秒）的提升。

结论 ¶

在本教程中，我们通过介绍基本用法、演示使用快速数学内核的加速、比较不同批处理尺寸的不同后端以及如何使用 Linux 透明大页优化张量内存分配延迟，涵盖了在 AWS Graviton3 实例上使用 PyTorch 推理。建议对于较大的张量形状使用 MKLDNN 后端和 Bfloat16 快速数学模式以及 THP 内存分配，对于较小的张量形状使用 OpenBLAS 后端。希望您能尝试一下！