(Beta) 在 AWS Graviton 處理器上進行 PyTorch 推論效能調整

建立於：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) 的 MKLDNN 和 OpenBLAS 函式庫為運算子的一個子集提供了最佳化的實作。這兩個函式庫都已透過 PyTorch 2.0 版本整合到 PyTorch 中。

在本教學中，我們將介紹如何在 AWS Graviton3 CPU (AWS c7g 執行個體) 上，透過 bfloa16 核心以及正確的後端選擇，為線性層神經網路實現最佳的推論效能。

目錄

1. 基本用法

2. 透過 Bfloat16 快速數學核心加速推論

3. 透過 OpenBLAS 提高較小批次維度的推論效能

4. 使用 Linux Transparent Huge Pages 最佳化記憶體配置開銷

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")

3. 鑑於線性層是多個神經網路（包括 Transformer）的核心，我們採用線性層進行此示範。我們透過子類別化 nn.Module 來定義我們的神經網路，並在 __init__ 中初始化層。我們使用典型的 LLM 參數建構網路，以符合真實世界的場景

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

Self CPU time total: 16.201 秒

透過 bfloat16 快速數學核心加速推論

AWS Graviton3 處理器支援 bfloat16 MMLA 指令。Arm Compute Library (ACL) 為 AWS Graviton 處理器提供最佳化的 bfloat16 General Matrix Multiplication (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 fastmath 核心，效能提升約 2x (7.262s vs 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 預設配置執行上述腳本時，您應該會看到以下效能分析器輸出：

名稱	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.47x (4.123s vs 6.094s) 的效能提升。雖然此改進值得注意，但整體效能仍有改進空間。這是因為來自 oneDNN 和 ACL 後端的執行時間開銷（權重重新排序和核心啟動時間）超過了較小批次運算的 ACL GEMM 核心帶來的運算優勢。

使用 OpenBLAS 提升較小批次維度的推論效能

對於較小的批次維度，可以透過將較小的形狀從 MKLDNN 分流到 OpenBLAS 後端來提高推論效能。我們正在努力使後端選擇自動化，並在未來的版本中使用穩健的啟發式方法。在啟發式方法實施之前，可以透過增加 MKLDNN 後端選擇的閾值，將較小的形狀分流到 OpenBLAS。在以下範例中，我們使用 64 作為閾值，以便 batch dimension of 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.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	# 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.2x (1.272s vs 4.115s)。

使用 Linux Transparent Huge Pages (THP) 最佳化記憶體配置開銷

我們還觀察到，對於這些較大的網路，張量記憶體配置佔據了推論延遲的很大一部分。可以透過從 PyTorch C10 記憶體分配器啟用 Linux transparent huge page 分配來最佳化此過程。目前，預設情況下未啟用此功能，因為它會稍微增加記憶體佔用量。設定以下環境變數以啟用它

$ 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.08x 或 8% (6.697s vs 7.262s) 的效能。

結論

在本教學課程中，我們介紹了在 AWS Graviton3 執行個體上的 PyTorch 推論，涵蓋了基本用法，展示了快速數學核心帶來的加速，比較了不同批次維度的不同後端，以及如何使用 Linux transparent huge pages 最佳化張量記憶體配置延遲。建議對於較大的張量形狀，使用帶有 Bfloat16 fastmath 模式和 THP 記憶體配置的 MKLDNN 後端，對於較小的張量形狀，則使用 OpenBLAS 後端。我們希望您能試試看！