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

创建于: 2024 年 1 月 24 日 | 最后更新: 2024 年 1 月 24 日 | 最后验证: 2024 年 11 月 5 日

作者: Sunita Nadampalli

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 内核并选择合适的后端，为线性层神经网络实现最佳推理性能。

目录

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.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

使用 bfloat16 快速数学内核，性能提升约 2 倍 (7.262s 对比 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.47 倍 (4.123s 对比 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

如上所示，切换到 OpenBLAS 后，性能比默认的 MKLDNN 后端配置提升了 一倍 (2.034s 对比 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.2 倍 (1.272s 对比 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.08 倍或 8% (6.697s 对比 7.262s) 的提升。

结论

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