(Beta) AWS Graviton 处理器上的 PyTorch 推理性能调优¶
AWS Graviton 是一系列由 AWS 设计的基于 ARM 的处理器。AWS Graviton3 处理器针对机器学习 (ML) 工作负载进行了优化，包括对 bfloat16、可扩展向量扩展 (SVE) 和与 Graviton2 相比两倍的单指令多数据 (SIMD) 带宽的支持。
PyTorch 为机器学习操作符（如卷积、矩阵乘法、ReLU 等）提供原生参考 ATen 内核。这些操作符可以通过来自基本线性代数 (BLAS) 库的平台特定内核实现来加速。在 AWS Graviton CPU 上，带有 Arm 计算库 (ACL) 的 MKLDNN 和 OpenBLAS 库为部分操作符提供了优化的实现。这两个库都已集成到 PyTorch 2.0 版本中。
在本教程中，我们将介绍如何在 AWS Graviton3 CPU (AWS c7g 实例) 上使用 bfloa16 内核和正确的后端选择，为线性层神经网络获得最佳推理性能。
内容¶
基本用法
使用 Bfloat16 快速数学内核加速推理
对于较小的批次维度，使用 OpenBLAS 提高推理性能
使用 Linux 透明巨页优化内存分配开销
结论
注意
为了成功运行本教程并复现以下所示的加速数值，您需要一个来自 Graviton3 系列（c7g/r7g/m7g）的硬件实例。在本教程中，我们使用了 c7g.xl (4vcpu) 实例。
基本用法¶
从 PyTorch 2.0 版本开始，PyTorch 原生支持 AWS Graviton3 优化。请参阅此 博客 以了解有关优化的更多详细信息。
运行以下命令安装 PyTorch
python3 -m pip install torch
我们将从导入所需的依赖项和定义将要运行的设备开始
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")
鉴于线性层是包括 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
让我们创建一个
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.813s |
98.61% |
15.977s |
53.255ms |
300 |
aten::clamp_min |
1.09% |
177.032ms |
1.09% |
177.032ms |
885.160us |
200 |
aten::copy |
1.00% |
162.054ms |
1.00% |
162.054ms |
540.180us |
300 |
mymodel_inference |
0.22% |
35.738ms |
100.00% |
16.201s |
16.201s |
1 |
aten::linear |
0.02% |
2.955ms |
98.66% |
15.985s |
53.282ms |
300 |
aten::t |
0.01% |
2.421ms |
0.03% |
5.043ms |
16.810us |
300 |
aten::relu |
0.01% |
2.356ms |
1.11% |
179.388ms |
896.940us |
200 |
自身 CPU 时间总计： 16.201s
使用 bfloat16 快速数学内核加速推理¶
AWS Graviton3 处理器支持 bfloat16 MMLA 指令。Arm Compute Library (ACL) 为 AWS Graviton 处理器提供了优化的 bfloat16 通用矩阵乘法 (GEMM) 内核，并通过 MKLDNN 后端集成到 PyTorch 中，从 PyTorch 2.0 开始。推理性能可以通过快速数学 GEMM 内核进行优化。默认情况下不会启用快速数学模式，因为这些内核以 bfloat16 精度而不是 float 精度执行 GEMM，因此会导致模型推理精度略有下降。但是，精度下降在 torchbench 测试套件中为 bfloat16 后端定义的 cosine similarity 阈值范围内，因此对于大多数应用程序来说是可以接受的。要启用快速数学 GEMM 内核，请设置以下环境变量
$ export DNNL_DEFAULT_FPMATH_MODE=BF16
运行上述推理脚本时，您应该会看到启用 MKLDNN 快速数学模式后的以下分析器输出
名称 |
自身 CPU % |
自身 CPU |
CPU 总 % |
CPU 总计 |
CPU 时间平均值 |
调用次数 |
|---|---|---|---|---|---|---|
aten::addmm |
95.61% |
6.943s |
97.10% |
7.052s |
23.507ms |
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_inference |
0.43% |
31.167ms |
100.00% |
7.262s |
7.262s |
1 |
aten::linear |
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.262s
使用 bfloat16 快速数学内核可以提高大约 2x (7.262s vs 16.201s) 的性能。接下来，让我们看看较小的批次维度场景。
场景 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.821s |
97.04% |
5.914s |
19.713ms |
300 |
aten::clamp_min |
2.33% |
142.244ms |
2.33% |
142.244ms |
711.220us |
200 |
aten::copy |
1.51% |
92.322ms |
1.51% |
92.322ms |
307.740us |
300 |
mymodel_inference |
0.45% |
27.713ms |
100.00% |
6.094s |
6.094s |
1 |
aten::linear |
0.04% |
2.495ms |
97.16% |
5.921s |
19.736ms |
300 |
aten::t |
0.03% |
2.131ms |
0.07% |
4.441ms |
14.803us |
300 |
aten::relu |
0.03% |
1.942ms |
2.37% |
144.186ms |
720.930us |
200 |
自身 CPU 时间总计： 6.094s
以下是启用 MKLDNN 快速数学模式时的分析器输出
$ export DNNL_DEFAULT_FPMATH_MODE=BF16
名称 |
自身 CPU % |
自身 CPU |
CPU 总 % |
CPU 总计 |
CPU 时间平均值 |
调用次数 |
|---|---|---|---|---|---|---|
aten::addmm |
93.31% |
3.848s |
95.66% |
3.944s |
13.148ms |
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_inference |
0.67% |
27.431ms |
100.00% |
4.123s |
4.123s |
1 |
aten::linear |
0.06% |
2.471ms |
95.83% |
3.951s |
13.170ms |
300 |
aten::t |
0.05% |
2.027ms |
0.10% |
4.243ms |
14.143us |
300 |
aten::relu |
0.05% |
1.928ms |
3.47% |
143.237ms |
716.185us |
200 |
自身 CPU 时间总计： 4.123s
MKLDNN 快速数学模式为较小的批次维度带来了大约 1.47x (4.123s vs 6.094s) 的性能提升。虽然此改进值得注意，但整体性能仍有改进空间。这是因为对于较小的批次计算，来自 oneDNN 和 ACL 后端的运行时开销（权重重新排序和内核启动时间）超过了 ACL GEMM 内核带来的计算优势。
使用 OpenBLAS 改善较小批次维度的推理性能¶
可以通过将较小的形状从 MKLDNN 卸载到 OpenBLAS 后端来改善较小批次维度的推理性能。我们正在努力在未来的版本中使用稳健的启发式方法来自动进行后端选择。在实现启发式方法之前，可以通过提高 MKLDNN 后端选择的阈值将较小的形状卸载到 OpenBLAS。在以下示例中，我们使用 64 作为阈值，以便批次维度为 32 的输入不会被调度到 MKLDNN。相反，它将被调度到 OpenBLAS。
$ export TORCH_MKLDNN_MATMUL_MIN_DIM=64
以下是使用 OpenBLAS 后端的分析器输出
名称 |
自身 CPU % |
自身 CPU |
CPU 总 % |
CPU 总计 |
CPU 时间平均值 |
调用次数 |
|---|---|---|---|---|---|---|
aten::addmm |
96.25% |
1.958s |
97.51% |
1.984s |
6.612ms |
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_inference |
0.86% |
17.423ms |
100.00% |
2.034s |
2.034s |
1 |
aten::linear |
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.034s
如上所示，与默认的 MKLDNN 后端配置相比，切换到 OpenBLAS 将性能提高了一倍 (2.034s vs 4.123s)。这对于更小的批次维度也变得很重要，例如，对于批次维度为 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.613s |
91.90% |
3.781s |
12.604ms |
300 |
aten::clamp_min |
7.18% |
295.437ms |
7.18% |
295.437ms |
1.477ms |
200 |
aten::copy |
4.07% |
167.516ms |
4.07% |
167.516ms |
558.387us |
300 |
mymodel_inference |
0.67% |
27.708ms |
100.00% |
4.115s |
4.115s |
1 |
aten::linear |
0.06% |
2.499ms |
92.06% |
3.788s |
12.627ms |
300 |
aten::t |
0.05% |
1.982ms |
0.11% |
4.385ms |
14.617us |
300 |
aten::relu |
0.05% |
1.932ms |
7.23% |
297.369ms |
1.487ms |
200 |
自身 CPU 时间总计： 4.115s
以下是使用 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_inference |
1.38% |
17.521ms |
100.00% |
1.272s |
1.272s |
1 |
aten::linear |
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.272s
在这里，我们通过适当地调整后端阈值观察到 3.2x (1.272s vs 4.115s) 的性能提升。
使用 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.115s |
94.39% |
6.321s |
21.069ms |
300 |
aten::clamp_min |
4.82% |
322.568ms |
4.82% |
322.568ms |
1.613ms |
200 |
aten::copy |
3.06% |
204.602ms |
3.06% |
204.602ms |
682.007us |
300 |
mymodel_inference |
0.61% |
40.777ms |
100.00% |
6.697s |
6.697s |
1 |
aten::linear |
0.05% |
3.082ms |
94.51% |
6.329s |
21.097ms |
300 |
aten::relu |
0.04% |
2.547ms |
4.85% |
325.115ms |
1.626ms |
200 |
自身 CPU 时间总计： 6.697s
这在上面测量的已优化的 MKLDNN 快速数学模式的基础上额外提高了 1.08x 或 8% (6.697s vs 7.262s) 的性能。
结论¶
在本教程中，我们涵盖了 AWS Graviton3 实例上的 PyTorch 推理，内容包括基本用法、演示快速数学内核的加速、比较不同批次维度下的不同后端，以及如何使用 Linux 透明大页优化张量内存分配延迟。建议对于较大的张量形状，使用启用 Bfloat16 快速数学模式和 THP 内存分配的 MKLDNN 后端，而对于较小的张量形状，使用 OpenBLAS 后端。我们希望您能尝试一下！
