(测试版) 在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 Compute Library (ACL) 和 OpenBLAS 库的 MKLDNN 为一部分操作符提供了优化实现。这两个库都已集成到 PyTorch 2.0 版本中。

在本教程中，我们将介绍如何在AWS Graviton3 CPU（AWS c7g instance）上使用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. 鉴于线性层是包括变压器在内的多种神经网络的核心，我们在这个演示中使用了一个线性层。我们通过子类化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_inference	0.22%	35.738毫秒	100.00%	16.201秒	16.201秒	1
aten::linear	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.940微秒	200

自身CPU时间总计： 16.201秒

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

AWS Graviton3 处理器支持 bfloat16 MMLA 指令。Arm 计算库 (ACL) 为 AWS Graviton 处理器提供了优化的 bfloat16 通用矩阵乘法 (GEMM) 内核，并从 PyTorch 2.0 开始通过 MKLDNN 后端集成到 PyTorch 中。推理性能可以通过快速数学 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.943秒	97.10%	7.052秒	23.507毫秒	300
aten::clamp_min	2.31%	167.653毫秒	2.31%	167.653毫秒	838.265微秒	200
aten::copy	1.48%	107.593毫秒	1.48%	107.593毫秒	358.643微秒	300
mymodel_inference	0.43%	31.167毫秒	100.00%	7.262秒	7.262秒	1
aten::linear	0.04%	2.911毫秒	97.21%	7.060秒	23.533毫秒	300
aten::t	0.03%	2.414毫秒	0.07%	4.892毫秒	16.307微秒	300
aten::relu	0.03%	2.281毫秒	2.34%	169.934毫秒	849.670微秒	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_inference	0.45%	27.713毫秒	100.00%	6.094秒	6.094秒	1
aten::linear	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.309毫秒	3.43%	141.309毫秒	706.545微秒	200
aten::copy	2.33%	95.916毫秒	2.33%	95.916毫秒	319.720微秒	300
mymodel_inference	0.67%	27.431毫秒	100.00%	4.123秒	4.123秒	1
aten::linear	0.06%	2.471毫秒	95.83%	3.951秒	13.170毫秒	300
aten::t	0.05%	2.027毫秒	0.10%	4.243毫秒	14.143微秒	300
aten::relu	0.05%	1.928毫秒	3.47%	143.237毫秒	716.185微秒	200

自身CPU时间总计： 4.123秒

MKLDNN快速数学模式在较小的批量维度上带来了大约1.47倍（4.123秒 vs 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.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透明大页优化张量内存分配延迟。建议对于较大的张量形状使用MKLDNN后端和Bfloat16快速数学模式以及THP内存分配，对于较小的张量形状使用OpenBLAS后端。我们希望您能尝试一下！