加速 PyTorch Transformers 使用 Intel Sapphire Rapids – 第 1 部分

TL;DR

Hugging Face 示範，在使用 Intel Extension for PyTorch (IPEX) 與 oneCCL 的四節點 Intel Sapphire Rapids CPU 叢集上訓練 PyTorch Transformers，可比相當的 Ice Lake 實例提升 8 倍速度，且幾乎線性擴展，使基於 CPU 的訓練成為成本效益高的 GPU 替代方案。

為何在 CPU 上訓練？

在 Intel Xeon CPU 上訓練可能比基於 GPU 的訓練更便宜且更具可擴展性，特別是對於小至中型模型與資料集。Xeon CPU 提供 AVX‑512、超執行緒（Hyper‑Threading），以及現在的先進矩陣擴展（AMX），可加速矩陣乘法。CPU 廣泛可得，可重新分配至其他工作負載，且雲端 Spot 實例相較於隨選實例可降低高達 90% 的成本。

先進矩陣擴展 (AMX)

Sapphire Rapids 引入 AMX，一組新指令可加速矩陣乘法——深度學習訓練的核心運算。AMX 支援 BF16 與 INT8 資料型別，並新增 2 維瓦片暫存器。要使用 AMX，Linux 核心必須為 5.16 版或更新；然而，Intel 與 AWS 已將所需支援回移至 5.15 版核心，供本指南使用的裸金屬 r7iz.metal-16xl 實例使用。

在 AWS 上建置 Sapphire Rapids 叢集

存取 Sapphire Rapids 硬體最簡單的方式是透過 Amazon EC2 R7iz 裸金屬實例（例如 r7iz.metal-16xl）。因為預覽版尚未在虛擬機器中支援 AMX，指南使用 64 個 vCPU 與 512 GB 記憶體的裸金屬實例。

網路設定

  • 開放所有節點的 SSH（埠 22）。
  • 從主節點設定免密碼 SSH 至每個節點，包括主節點本身。
  • 建立允許不受限制叢集內部流量的安全群組，並將其附加至所有實例。

設定主節點

  1. 使用 Ubuntu 20.04 AMI ami-07cd3e6c4915b2d18 啟動 r7iz.metal-16xl 實例。
  2. 使用 lscpu 驗證 AMX 支援（旗標 amx_bf16 amx_tile amx_int8）。
  3. 安裝相依套件：
    sudo apt-get update
sudo apt install libgoogle-perftools-dev -y
sudo apt-get install python3-pip -y
pip install --upgrade pip
pip install virtualenv
virtualenv cluster_env
source cluster_env/bin/activate
pip install torch==1.13.0 -f https://download.pytorch.org/whl/cpu
pip install intel_extension_for_pytorch==1.13.0 -f https://developer.intel.com/ipex-whl-stable-cpu
pip install -f https://developer.intel.com/ipex-whl-stable-cpu
pip install transformers==4.24.0
git clone https://github.com/huggingface/transformers.git
cd transformers && git checkout v4.24.0
  4. 產生 SSH 金鑰對（ssh-keygen），並存放於 ~/.ssh/cluster。
  5. 從此已設定好的實例建立 AMI，以便重複使用。

設定其餘節點

  1. 使用上述建立的 AMI 再啟動三個 r7iz.metal-16xl 實例。
  2. 在主節點編輯 ~/.ssh/config，定義主機 node1、node2、node3，並指定其私有 IP 與 cluster 金鑰。
  3. 使用 ssh node[1-3] 驗證免密碼存取。
  4. 建立 ~/hosts 檔案，列出所有節點（包括主節點的 localhost）：
    localhost
node1
node2
node3

啟動分散式訓練

此範例微調 DistilBERT 於 SQuAD 資料集。

單節點基線

source ~/cluster_env/bin/activate
cd ~/transformers/examples/pytorch/question-answering
pip install -r requirements.txt
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

單個 epoch 完成於 ≈26 分鐘，相較於 Ice Lake c6i.16xlarge 實例的 ≈3 小時 30 分鐘，提升了 8×。

四節點分散式執行

oneCCL 與執行緒配置的環境變數：

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               # 2 processes per node × 4 nodes
export NUM_PROCESSES_PER_NODE=2
export CCL_WORKER_COUNT=2
export CCL_WORKER_AFFINITY=auto
export KMP_HW_SUBSET=1T
export OMP_NUM_THREADS=24           # training threads per process

使用 mpirun 啟動：

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 秒，僅比理想的線性擴展目標 6 分 30 秒（26 分 / 4）高出 1 分鐘。原文中的截圖顯示每個節點有兩個訓練程序，且主程序負責協調執行。

含意與要點

  • 成本效益 – CPU Spot 實例相較於 GPU 實例可大幅降低成本，同時在許多工作負載下提供相當的訓練速度。
  • 可擴展性 – 四個 Sapphire Rapids 節點的近線性擴展證明 IPEX 與 oneCCL 的組合能有效分配計算與通訊。
  • 易用性 – 無需修改程式碼；只要啟用 --use_ipex 與 --bf16 即可自動使用新的 AMX 指令。
  • 未來工作 – 作者計畫在後續文章中探討 Sapphire Rapids 的推論效能，完成 CPU 端端到端加速的全貌。

其他資源

Sources