利用英特爾 Sapphire Rapids 加速 PyTorch Transformer，第 1 部分

釋出於 2023 年 1 月 2 日

大約一年前，我們向你展示了如何在第三代英特爾至強可擴充套件 CPU（即 Ice Lake）叢集上分散式訓練 Hugging Face Transformer 模型。最近，英特爾推出了代號為 Sapphire Rapids 的第四代至強 CPU，其激動人心的新指令集可以加速深度學習模型中常見的運算。

在這篇文章中，你將學習如何使用在 AWS 上執行的 Sapphire Rapids 伺服器叢集來加速 PyTorch 訓練任務。我們將使用英特爾 oneAPI 集體通訊庫（CCL）來分發任務，並使用英特爾 PyTorch 擴充套件庫（IPEX）來自動啟用新的 CPU 指令。由於這兩個庫都已與 Hugging Face Transformer 庫整合，我們將能夠直接執行示例指令碼，無需更改任何一行程式碼。

在後續的文章中，我們將探討在 Sapphire Rapids CPU 上的推理及其帶來的效能提升。

為什麼你應該考慮在 CPU 上進行訓練

在英特爾至強 CPU 上訓練深度學習（DL）模型可以是一種經濟高效且可擴充套件的方法，尤其是在使用分散式訓練和對中小型資料集進行微調等技術時。

至強 CPU 支援高階向量擴充套件（AVX-512）和超執行緒等高階功能，有助於提高深度學習模型的並行度和效率。這使得訓練時間更短，硬體資源利用率更高。

此外，與 GPU 等通常用於訓練大型深度學習模型的專用硬體相比，至強 CPU 通常更實惠且更易獲得。至強 CPU 還可以輕鬆地用於其他生產任務，從 Web 伺服器到資料庫，使其成為你 IT 基礎設施中一個多功能且靈活的選擇。

最後，雲使用者可以透過 Spot 例項進一步降低在至強 CPU 上的訓練成本。Spot 例項由閒置的計算能力構成，並以折扣價出售。與使用按需例項相比，它們可以節省大量成本，有時高達 90%。最後但同樣重要的是，CPU Spot 例項通常也比 GPU 例項更容易獲取。

現在，讓我們來看看 Sapphire Rapids 架構中的新指令集。

高階矩陣擴充套件：深度學習的新指令集

Sapphire Rapids 架構引入了英特爾高階矩陣擴充套件（AMX）來加速深度學習工作負載。使用它們就像安裝最新版本的 IPEX 一樣簡單。無需對你的 Hugging Face 程式碼做任何更改。

AMX 指令加速了矩陣乘法，這是在資料批次上訓練深度學習模型的核心操作。它們支援 Brain Floating Point（BF16）和 8 位整數（INT8）值，為不同的訓練場景提供了加速。

AMX 引入了新的二維 CPU 暫存器，稱為 Tile 暫存器。由於這些暫存器需要在上下文切換期間儲存和恢復，它們需要核心支援：在 Linux 上，你需要 v5.16 或更新的版本。

現在，讓我們看看如何構建一個用於分散式訓練的 Sapphire Rapids CPU 叢集。

構建 Sapphire Rapids CPU 叢集

在撰寫本文時，體驗 Sapphire Rapids 伺服器最簡單的方法是使用新的 Amazon EC2 R7iz 例項系列。由於它仍處於預覽階段，你必須註冊才能獲得訪問許可權。此外，虛擬伺服器尚不支援 AMX，因此我們將使用裸金屬例項（r7iz.metal-16xl，64 個 vCPU，512GB 記憶體）。

為避免手動設定叢集中的每個節點，我們將首先設定主節點，並從中建立一個新的亞馬遜機器映象（AMI）。然後，我們將使用此 AMI 來啟動其他節點。

從網路角度看，我們需要進行以下設定：

在所有例項上開放 22 埠用於 SSH 訪問，以便進行設定和除錯。
配置從主例項（你將從中啟動訓練）到所有其他例項（包括主例項本身）的免密 SSH 登入。換句話說，主節點的 SSH 公鑰必須在所有節點上被授權。
允許叢集內部的所有網路流量，以便分散式訓練能夠無障礙執行。AWS 提供了一種安全便捷的方式來實現這一點，即使用安全組。我們只需建立一個安全組，允許來自配置了相同安全組的例項的所有流量，並確保將其附加到叢集中的所有例項。這是我的設定示例。

讓我們開始動手，構建叢集的主節點。

設定主節點

我們首先透過啟動一個 r7iz.metal-16xl 例項來建立主節點，該例項使用 Ubunutu 20.04 AMI（ami-07cd3e6c4915b2d18）和我們之前建立的安全組。這個 AMI 包含 Linux v5.15.0，但幸運的是，英特爾和 AWS 已經為核心打了補丁以新增 AMX 支援。因此，我們不需要將核心升級到 v5.16。

例項執行後，我們透過 SSH 連線到它，並使用 lscpu 命令檢查是否確實支援 AMX。你應該在 flags 部分看到以下內容：

amx_bf16 amx_tile amx_int8

然後，我們安裝原生和 Python 依賴。

sudo apt-get update 

# Install tcmalloc for extra performance (https://github.com/google/tcmalloc)
sudo apt install libgoogle-perftools-dev -y

# Create a virtual environment
sudo apt-get install python3-pip -y
pip install pip --upgrade
export PATH=/home/ubuntu/.local/bin:$PATH
pip install virtualenv

# Activate the virtual environment
virtualenv cluster_env
source cluster_env/bin/activate

# Install PyTorch, IPEX, CCL and Transformers
pip3 install torch==1.13.0 -f https://download.pytorch.org/whl/cpu
pip3 install intel_extension_for_pytorch==1.13.0 -f https://developer.intel.com/ipex-whl-stable-cpu
pip3 install oneccl_bind_pt==1.13 -f https://developer.intel.com/ipex-whl-stable-cpu
pip3 install transformers==4.24.0

# Clone the transformers repository for its example scripts
git clone https://github.com/huggingface/transformers.git
cd transformers
git checkout v4.24.0

接下來，我們使用 ssh-keygen 建立一個名為“cluster”的新 SSH 金鑰對，並將其儲存在預設位置（~/.ssh）。

最後，我們從此例項建立一個新的 AMI。

設定叢集

AMI 準備好後，我們用它來啟動另外 3 個 r7iz.16xlarge-metal 例項，別忘了附加之前建立的安全組。

在這些例項啟動的同時，我們透過 SSH 連線到主節點以完成網路設定。首先，我們編輯位於 ~/.ssh/config 的 SSH 配置檔案，以啟用從主節點到所有其他節點的免密連線，使用它們的私有 IP 地址和之前建立的金鑰對。這是我的檔案示例：

Host 172.31.*.*
   StrictHostKeyChecking no

Host node1
    HostName 172.31.10.251
    User ubuntu
    IdentityFile ~/.ssh/cluster

Host node2
    HostName 172.31.10.189
    User ubuntu
    IdentityFile ~/.ssh/cluster

Host node3
    HostName 172.31.6.15
    User ubuntu
    IdentityFile ~/.ssh/cluster

此時，我們可以使用 ssh node[1-3] 連線到任何節點，無需任何提示。

仍在主節點上，我們建立一個 ~/hosts 檔案，其中包含叢集中所有節點的名稱，這些名稱已在上面的 SSH 配置中定義。我們將主節點命名為 localhost，因為我們將在那裡啟動訓練指令碼。這是我的檔案示例：

localhost
node1
node2
node3

叢集現在已經準備就緒。讓我們開始訓練吧！

啟動分散式訓練任務

在這個例子中，我們將在 SQUAD 資料集上微調一個 DistilBERT 模型用於問答任務。如果你願意，也可以嘗試其他例子。

source ~/cluster_env/bin/activate
cd ~/transformers/examples/pytorch/question-answering
pip3 install -r requirements.txt

作為健全性檢查，我們首先啟動一個本地訓練任務。請注意幾個重要的標誌：

no_cuda 確保任務忽略此機器上的任何 GPU，
use_ipex 啟用 IPEX 庫，從而啟用 AVX 和 AMX 指令，
bf16 啟用 BF16 訓練。

export LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc.so"
python run_qa.py --model_name_or_path distilbert-base-uncased \
--dataset_name squad --do_train --do_eval --per_device_train_batch_size 32 \
--num_train_epochs 1  --output_dir /tmp/debug_squad/ \
--use_ipex --bf16 --no_cuda

無需讓任務執行到完成，我們只需執行一分鐘以確保所有依賴項都已正確安裝。這也為我們提供了單例項訓練的基準：1 個 epoch 大約需要 26 分鐘。作為參考，我們在具有相同軟體設定的可比 Ice Lake 例項（c6i.16xlarge）上計時了相同的任務，每個 epoch 需要 3 小時 30 分鐘。這是一個 8 倍的加速。我們已經可以看到新指令集帶來的好處了！

現在，讓我們在四個例項上分散式訓練任務。一個 r7iz.16xlarge 例項有 32 個物理 CPU 核心，我們更喜歡直接使用它們而不是使用 vCPU（KMP_HW_SUBSET=1T）。我們決定為訓練分配 24 個核心（OMP_NUM_THREADS），為 CCL 通訊分配 2 個核心（CCL_WORKER_COUNT），將最後 6 個執行緒留給核心和其他程序。這 24 個訓練執行緒支援 2 個 Python 程序（NUM_PROCESSES_PER_NODE）。因此，在 4 節點叢集上執行的 Python 任務總數為 8（NUM_PROCESSES）。

# Set up environment variables for CCL
oneccl_bindings_for_pytorch_path=$(python -c "from oneccl_bindings_for_pytorch import cwd; print(cwd)")
source $oneccl_bindings_for_pytorch_path/env/setvars.sh

export MASTER_ADDR=172.31.3.190
export NUM_PROCESSES=8
export NUM_PROCESSES_PER_NODE=2
export CCL_WORKER_COUNT=2
export CCL_WORKER_AFFINITY=auto
export KMP_HW_SUBSET=1T

現在，我們啟動分散式訓練任務。

# Launch distributed training
mpirun -f ~/hosts \
 -n $NUM_PROCESSES -ppn $NUM_PROCESSES_PER_NODE  \
 -genv OMP_NUM_THREADS=24 \
 -genv LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc.so" \
 python3 run_qa.py \
 --model_name_or_path distilbert-base-uncased \
 --dataset_name squad \
 --do_train \
 --do_eval \
 --per_device_train_batch_size 32  \
 --num_train_epochs 1  \
 --output_dir /tmp/debug_squad/ \
 --overwrite_output_dir \
 --no_cuda \
 --xpu_backend ccl \
 --bf16

一個 epoch 現在需要 7 分 30 秒。

這是任務執行時的樣子。主節點在頂部，你可以看到在其他 3 個節點上分別執行的兩個訓練程序。