（Beta）PyTorch 在 AWS Graviton 处理器上的推理性能调优

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

作者：Sunita Nadampalli

AWS Graviton 是 AWS 设计的一系列基于 ARM 的处理器。AWS Graviton3 处理器针对机器学习 (ML) 工作负载进行了优化，包括支持 bfloat16、可伸缩矢量扩展 (SVE) 和相比 Graviton2 两倍的单指令多数据 (SIMD) 带宽。

PyTorch 为机器学习运算符（如卷积、matmul、relu 等）提供原生参考 ATen 内核。这些运算符可以通过来自基本线性代数 (BLAS) 库的平台特定内核实现来加速。在 AWS Graviton CPU 上，带有 Arm Compute Library (ACL) 和 OpenBLAS 库的 MKLDNN 为一部分运算符提供了优化的实现。这两个库都已集成到 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 版本开始，PyTorch 原生支持 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 配置的性能分析器输出

名称	Self CPU %	Self CPU	CPU total %	CPU total	CPU time avg	# of Calls
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) 内核，并通过 MKLDNN 后端集成到 PyTorch 2.0 中。可以使用快速数学 GEMM 内核优化推理性能。默认情况下未启用快速数学模式，因为这些内核以 bfloat16 精度而不是 float 精度执行 GEMM，因此会导致模型推理精度略有下降。但是，精度下降在 torchbench 测试套件中为 bfloat16 后端定义的 cosine similarity 阈值范围内，因此对于大多数应用程序来说是可以接受的。要启用快速数学 GEMM 内核，请设置以下环境变量

$ export DNNL_DEFAULT_FPMATH_MODE=BF16

当您运行上述推理脚本时，您应该看到以下启用了 MKLDNN 快速数学模式的性能分析器输出

名称	Self CPU %	Self CPU	CPU total %	CPU total	CPU time avg	# of Calls
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

使用 bfloat16 快速数学内核，性能提升约 2 倍（7.262 秒 vs 16.201 秒）。接下来，让我们看看较小批次维度的情况。

场景 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	# of Calls
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	# of Calls
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.47 倍（4.123 秒 vs 6.094 秒） 的性能提升。虽然这种改进值得注意，但整体性能仍有改进空间。这是因为来自 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	# of Calls
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.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 快速数学模式的性能分析器输出

名称	Self CPU %	Self CPU	CPU total %	CPU total	CPU time avg	# of Calls
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	# of Calls
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.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 内存分配的性能分析器输出

名称	Self CPU %	Self CPU	CPU total %	CPU total	CPU time avg	# of Calls
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.697 秒 vs 7.262 秒）。

结论

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