(Beta) 基于 AWS Graviton 处理器的 PyTorch 推理性能优化¶
Created On: Jan 24, 2024 | Last Updated: Jan 24, 2024 | Last Verified: Nov 05, 2024
AWS Graviton 是 AWS 设计的一系列基于 ARM 的处理器。AWS Graviton3 处理器专为机器学习（ML）工作负载优化，包括支持 bfloat16、可扩展矢量扩展（SVE）和比 Graviton2 更高的单指令多数据（SIMD）带宽。
PyTorch 为机器学习操作符（如卷积、矩阵乘法和 ReLU 等）提供了本地的 ATen 内核。这些操作符可以通过来自基础线性代数库（BLAS）的特定平台内核实现来加速。在 AWS Graviton CPU 上，MKLDNN 与 Arm 计算库 (ACL) 和 OpenBLAS 提供了优化的部分操作符实现。这两个库自 PyTorch 2.0 版本起已集成到 PyTorch 中。
在本教程中，我们将讨论如何在 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])
我们的网络功能已验证完成。接下来我们将进行性能分析。让我们检查两个不同的场景：小批量和大批量尺寸。
场景一: 较大的批量尺寸，例如 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 计算库 (ACL) 提供了针对 AWS Graviton 处理器优化的 bfloat16 通用矩阵乘法（GEMM）内核，并从 PyTorch 2.0 开始通过 MKLDNN 后端集成到 PyTorch 中。可以通过快速数学 GEMM 内核优化推理性能。快速数学模式默认未启用，因为这些内核以 bfloat16 精度而非 float 进行 GEMM 操作，因此会导致模型推理精度略有下降。然而，该精度下降在 torchbench 测试套件中 bfloat16 后端定义的余弦相似性阈值范围内，因此对大多数应用程序而言是可以接受的。要启用快速数学 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
这是约 2 倍 (7.262s vs 16.201s) 的性能提升，得益于 bfloat16 快速数学内核。然而，让我们看看小批量尺寸场景。
场景二: 较小的批量尺寸，例如 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.47 倍 (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.958秒 |
97.51% |
1.984秒 |
6.612毫秒 |
300 |
aten::clamp_min |
1.28% |
26.124毫秒 |
1.28% |
26.124毫秒 |
130.620微秒 |
200 |
aten::copy |
1.23% |
24.951毫秒 |
1.23% |
24.951毫秒 |
83.170微秒 |
300 |
mymodel_inference |
0.86% |
17.423毫秒 |
100.00% |
2.034秒 |
2.034秒 |
1 |
aten::linear |
0.08% |
1.691毫秒 |
97.74% |
1.988秒 |
6.628毫秒 |
300 |
aten::t |
0.07% |
1.520毫秒 |
0.14% |
2.945毫秒 |
9.817微秒 |
300 |
aten::relu |
0.06% |
1.258毫秒 |
1.35% |
27.382毫秒 |
136.910微秒 |
200 |
自CPU时间总计： 2.034秒
如上所示，与默认的MKLDNN后端配置相比，切换到OpenBLAS使性能提高了一倍 （2.034秒 vs 4.123秒）。这种提升对更小的批次尺寸会更加显著，例如，批次尺寸为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_inference |
0.67% |
27.708毫秒 |
100.00% |
4.115秒 |
4.115秒 |
1 |
aten::linear |
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.179秒 |
95.23% |
1.211秒 |
4.038毫秒 |
300 |
aten::clamp_min |
2.83% |
36.060毫秒 |
2.83% |
36.060毫秒 |
180.300微秒 |
200 |
aten::copy |
2.52% |
32.013毫秒 |
2.52% |
32.013毫秒 |
106.710微秒 |
300 |
mymodel_inference |
1.38% |
17.521毫秒 |
100.00% |
1.272秒 |
1.272秒 |
1 |
aten::linear |
0.14% |
1.750毫秒 |
95.60% |
1.216秒 |
4.054毫秒 |
300 |
aten::t |
0.12% |
1.475毫秒 |
0.24% |
3.033毫秒 |
10.110微秒 |
300 |
aten::relu |
0.10% |
1.285毫秒 |
2.94% |
37.345毫秒 |
186.725微秒 |
200 |
自CPU时间总计： 1.272秒
通过适当地调整后端阈值，我们观察到**3.2倍（1.272秒 vs 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_inference |
0.61% |
40.777毫秒 |
100.00% |
6.697秒 |
6.697秒 |
1 |
aten::linear |
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秒 vs 7.262秒）**的改进。
结论¶
在本教程中，我们通过介绍AWS Graviton3实例上的PyTorch推理的基本用法、演示快速数学核的加速、对比不同批次尺寸的不同后端以及如何通过Linux透明巨页优化张量内存分配延迟，全面探讨了PyTorch推理。推荐对于更大的张量形状使用Bfloat16快速数学模式和THP内存分配的MKLDNN后端，而对于较小的张量形状使用OpenBLAS后端。希望您能尝试一下！
