评价此页

(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 为卷积、matmul、relu 等机器学习运算符提供了原生的 ATen 内核。这些运算符可以通过基础线性代数 (BLAS) 库中的平台特定内核实现来加速。在 AWS Graviton CPU 上,MKLDNN 结合 Arm Compute Library (ACL) 和 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 从 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")
  1. 鉴于线性层是许多神经网络(包括 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
  1. 让我们创建 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 后端。我们希望您会尝试一下!