(Beta) 在 AWS Graviton 处理器上进行 PyTorch 推理性能调优#
创建日期:2024 年 1 月 24 日 | 最后更新:2024 年 1 月 24 日 | 最后验证:2024 年 11 月 5 日
AWS Graviton 是 AWS 设计的一系列基于 ARM 的处理器。AWS Graviton3 处理器针对机器学习 (ML) 工作负载进行了优化,包括支持 bfloat16
、可扩展向量扩展 (SVE) 以及是 Graviton2 两倍的单指令多数据 (SIMD) 带宽。
PyTorch 为卷积、matmul、relu 等机器学习运算符提供了原生的 ATen 内核。这些运算符可以通过基础线性代数 (BLAS) 库中的平台特定内核实现来加速。在 AWS Graviton CPU 上,MKLDNN 结合 Arm Compute Library (ACL) 和 OpenBLAS 库为一部分运算符提供了优化实现。这两个库都已集成到 PyTorch 2.0 版本中。
在本教程中,我们将重点介绍如何在 AWS Graviton3 CPU(AWS c7g 实例)上使用 bfloa16
内核和正确的后端选择,以实现线性层神经网络的最佳推理性能。
目录#
基本用法
使用 Bfloat16 快速数学内核加速推理
通过 OpenBLAS 提高小批量尺寸的推理性能
通过 Linux 透明大页优化内存分配开销
结论
注意
要成功运行本教程并重现以下加速数据,您需要一个 Graviton3 系列(c7g/r7g/m7g
)的硬件实例。在本教程中,我们使用了 c7g.xl (4vcpu) 实例。
基本用法#
PyTorch 从 PyTorch 2.0 版本开始原生支持 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 配置的剖析器输出
名称 |
Self CPU % |
Self CPU |
CPU total % |
CPU total |
CPU time avg |
调用次数 |
---|---|---|---|---|---|---|
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 |
Self CPU time total: 16.201s
使用 bfloat16
快速数学内核加速推理#
AWS Graviton3 处理器支持 bfloat16 MMLA 指令。Arm Compute Library (ACL) 为 AWS Graviton 处理器提供了优化的 bfloat16
通用矩阵乘法 (GEMM) 内核,并自 PyTorch 2.0 起通过 MKLDNN 后端集成。可以通过快速数学 GEMM 内核优化推理性能。快速数学模式默认未启用,因为这些内核以 bfloat16
精度执行 GEMM,而不是 float
精度,因此会导致模型推理精度略有下降。但是,精度下降在 torchbench
测试套件中定义的 余弦相似度
阈值内,因此对大多数应用程序来说是可接受的。要启用快速数学 GEMM 内核,请设置以下环境变量
$ export DNNL_DEFAULT_FPMATH_MODE=BF16
当您运行上述推理脚本时,您应该会看到以下剖析器输出,其中启用了 MKLDNN 快速数学模式
名称 |
Self CPU % |
Self CPU |
CPU total % |
CPU total |
CPU time avg |
调用次数 |
---|---|---|---|---|---|---|
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 |
Self CPU time total: 7.262s
这大约是 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 默认配置运行上述脚本时,您应该会看到以下剖析器输出
名称 |
Self CPU % |
Self CPU |
CPU total % |
CPU total |
CPU time avg |
调用次数 |
---|---|---|---|---|---|---|
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 |
Self CPU time total: 6.094s
当启用 MKLDNN 快速数学模式运行时,以下是剖析器输出
$ export DNNL_DEFAULT_FPMATH_MODE=BF16
名称 |
Self CPU % |
Self CPU |
CPU total % |
CPU total |
CPU time avg |
调用次数 |
---|---|---|---|---|---|---|
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 |
Self CPU time total: 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 后端的剖析器输出
名称 |
Self CPU % |
Self CPU |
CPU total % |
CPU total |
CPU time avg |
调用次数 |
---|---|---|---|---|---|---|
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 |
Self CPU time total: 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 快速数学模式的剖析器输出
名称 |
Self CPU % |
Self CPU |
CPU total % |
CPU total |
CPU time avg |
调用次数 |
---|---|---|---|---|---|---|
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 |
Self CPU time total: 4.115s
以及使用 OpenBLAS 后端的剖析器输出
$ export TORCH_MKLDNN_MATMUL_MIN_DIM=64
名称 |
Self CPU % |
Self CPU |
CPU total % |
CPU total |
CPU time avg |
调用次数 |
---|---|---|---|---|---|---|
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 |
Self CPU time total: 1.272s
我们在此观察到,通过适当调整后端阈值,性能提升了 **3.2 倍 (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 内存分配后的剖析器输出
名称 |
Self CPU % |
Self CPU |
CPU total % |
CPU total |
CPU time avg |
调用次数 |
---|---|---|---|---|---|---|
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 |
Self CPU time total: 6.697s
这比上面测量的已优化的 MKLDNN 快速数学模式额外实现了 **1.08 倍或 8% (6.697s vs 7.262s)** 的性能提升。
结论#
在本教程中,我们通过介绍基本用法、演示使用快速数学内核的加速、比较不同批量尺寸的后端以及如何使用 Linux 透明大页优化张量内存分配延迟,全面介绍了 AWS Graviton3 实例上的 PyTorch 推理。建议对于较大的张量形状使用 MKLDNN 后端和 Bfloat16 快速数学模式以及 THP 内存分配,对于较小的张量形状使用 OpenBLAS 后端。我们希望您会尝试一下!