使用英特爾 Sapphire Rapids 加速 PyTorch Transformers 模型

語言: CN / TW / HK

時間 2023-01-10 00:02:11 HuggingFace

主題:

大約一年以前，我們展示^[1] 瞭如何在第三代英特爾至強可擴充套件^[2] CPU (即 Ice Lake) 叢集上分散式訓練 Hugging Face transformers 模型。最近，英特爾釋出了代號為 Sapphire Rapids 的第四代至強可擴充套件 CPU，該 CPU 包含了令人興奮的深度學習加速新指令。

通過本文，你將會學到如何在一個 AWS Sapphire Rapids 叢集上加速一個 PyTorch 訓練任務。我們會使用英特爾 oneAPI 集合通訊庫^[3] (oneAPI Collective Communications Library, oneCCL) 來分散式化訓練任務，並使用英特爾 PyTorch 擴充套件庫^[4] (Intel Extension for PyTorch，IPEX) 來自動使用新指令進行效能優化。因為這兩個庫均已整合入 Hugging Face transformers 庫，我們能夠做到在不修改一行程式碼的前提下開箱即用地執行我們的示例程式碼。

在隨後的另一篇文章裡，我們還會探討如何使用 Sapphire Rapids CPU 進行推理及其效能提升。

為何你應該考慮在 CPU 上訓練

在英特爾至強 CPU 上訓練一個深度學習模型是一個性價比高且可擴充套件的方案，在使用分散式訓練或者在小資料集或中等資料集上微調模型時尤其如此。

至強 CPU 支援一些先進的特性，如 512 位先進向量擴充套件 (Advanced Vector Extensions，AVX-512) 以及超執行緒 (Hyper-Threading) ，這些特性提高了深度學習模型的並行性和效率，使得我們可以在得到更好的硬體資源使用率的同時訓練得更快。

另外，一般而言，相比用於訓練大型深度學習模型的專門硬體如 GPU 等而言，至強 CPU 更便宜和易得。至強 CPU 還更容易用於其他生產任務，從網路服務到資料庫不一而足，這使得它們成為 IT 基礎設施的一個萬用且靈活的選擇。

最後，雲使用者還可以通過使用 spot 例項的方式進一步降低在至強 CPU 上的訓練成本。Spot 例項使用空閒計算資源，因此以折扣價售賣。與按需例項相比，spot 例項提供了高至 90% 的顯著的成本節約。最後同樣重要的是，CPU spot 例項一般來講比 GPU 例項更容易獲得。

現在，讓我們看一下 Sapphire Rapids 架構引入的新指令。

先進矩陣擴充套件 (AMX)：深度學習新指令

Sapphire Rapids 架構引入了英特爾先進矩陣擴充套件 (Advanced Matrix Extensions, AMX) 用於加速深度學習工作負載。使用者只需安裝最新版本的 IPEX 即可受益於新指令，無需更改任何 Hugging Face 程式碼。

AMX 指令用於加速矩陣乘法，該操作是深度學習批量訓練的核心操作。AMX 指令支援 Brain 浮點 (BF16) 和 8 位元整型 (INT8) 資料型別，覆蓋不同訓練場景的加速需求。

AMX 指令引入了新的 2 維 CPU 暫存器，稱作 tile 暫存器。因為這些暫存器在上下文切換時需要儲存和恢復，所以需要核心相關支援。在 Linux 上，核心版本需要在 v5.16 及以上方可支援。

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

構建一個 Sapphire Rapids CPU 叢集

截至本文撰寫之時，使用 Sapphire Rapids 伺服器的最簡單的方式是使用新的亞馬遜 EC2 R7iz^[5] 例項家族。由於它尚在預覽期，你必須登記註冊^[6] 以獲得訪問許可權。另外，虛擬機器尚未支援 AMX，因此，我們將使用裸金屬例項 (r7iz.metal-16xl, 64 vCPU, 512GB RAM) 。

為避免手動設定叢集中的每個節點，我們首先建立一個主節點並依此建立一個新的亞馬遜機器映象 (Amazon Machine Image，AMI^[7]) ，然後，我們用這個 AMI 啟動其他節點。

從網路的角度，我們需要如下設定：

開啟 22 埠，用於所有例項上的 ssh 訪問建立和除錯
配置從主例項 (你啟動訓練的那個例項) 到所有其他例項 (包含主例項本身) 的免密 SSH 訪問。換句話說，主節點的 ssh 公鑰必須在所有階段上被授權
允許叢集內的所有網路通訊，使得分散式訓練可以不受阻礙地執行。AWS 提供了安全組這一安全便捷的方式支援這個功能。我們只需建立一個安全組，確保所有叢集內的例項屬於同一安全組，並允許同一安全組內的所有網路通訊即可，以下是我使用的設定:

讓我們開始建立叢集的主節點。

設定主節點

我們首先啟動一個安裝了 Ubuntu 20.04 AMI (ami-07cd3e6c4915b2d18) 並加入了我們之前建立的安全組的 r7iz.metal-16xl 例項，用於建立主節點。該 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 建立一個新的 ssh 金鑰對，命名為 cluster，並儲存在預設位置 (~/.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^[8] 資料集上微調一個用於問答的 DistilBERT^[9] 模型。如果你想試試別的示例的話，儘管去做吧。

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

我們首先冒個煙，啟動一個單例項訓練任務。請注意如下幾個重要的標誌變數：

no_cuda 確保使用 CPU 進行訓練，忽略 GPU
use_ipex 使能 IPEX 庫，確保 AMX 和 AVX 指令的使用
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 分鐘用於確保所有的依賴已被正常安裝。同時，這也給了我們一個單例項訓練的基線效能：1 個 epoch 花費大約 26 分鐘。供參考，我們測量了同樣的任務在一個相當的 Ice Lake 例項 (c6i.16xlarge) 上的效能，基於相同的軟體設定，每個 epoch 需要 3 小時 30 分鐘。加速比達到 8 倍。我們已經能看到新指令帶來的好處！

現在，讓我們把訓練任務分散式部署到 4 個例項上。一個 r7iz.16xlarge 例項有 32 個物理 CPU 核，我們傾向於直接使用物理核而不是虛擬核 (vCPUs) (KMP_HW_SUBSET=1T)。我們決定分配 24 個核用於訓練 (OMP_NUM_THREADS) ，2 個核用於集合通訊 (CCL_WORKER_COUNT) ，剩下的 6 個核給核心和其他程序使用。這 24 個訓練執行緒分配給 2 個 Python 程序使用 (NUM_PROCESSES_PER_NODE) 。因此，一個 4 節點的叢集上共有 8 (NUM_PROCESSES) 個 Python 程序。

# 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 個節點每個均有 2 個訓練程序在執行。

4 節點的完美線性擴充套件需要 6 分 30 秒的訓練時間 (26 分鐘除以 4) 。我們非常接近於這個理想值，這充分展現了該方法很高的擴充套件性。

結論

如你所見，在一個英特爾至強叢集上訓練 Hugging Face transformers 模型是一個靈活，可擴充套件且價效比高的解決方案，特別是在你的資料集和模型是小尺寸或者中等尺寸情況下。

以下列出了一些其他可幫助你起步的資源：

Intel IPEX ^[10] GitHub 程式碼倉庫
Hugging Face 文件: Efficient training on CPU ^[11] 及 Efficient training on many CPUs ^[12]

如你有任何問題或反饋，請通過留言方式告訴我們。

感謝閱讀！

英文原文: https://huggingface.co/blog/intel-sapphire-rapids

譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

頭圖: 茶葉蛋蛋

文內連結

[1]

Hugging Face 部落格文章: 使用英特爾技術加速 PyTorch 分散式微調: https://huggingface.co/blog/accelerating-pytorch

[2]

英特爾® 至強® 可擴充套件處理器: https://www.intel.cn/content/www/cn/zh/products/details/processors/xeon/scalable.html

[3]

英特爾 oneAPI 集合通訊庫: https://www.intel.cn/content/www/cn/zh/developer/tools/oneapi/oneccl.html

[4]

英特爾 PyTorch 擴充套件庫: https://github.com/intel/intel-extension-for-pytorch

[5]

Amazon EC2 R7iz 例項: https://aws.amazon.com/cn/ec2/instance-types/r7iz/?nc1=h_ls

[6]

登記註冊 Amazon EC2 R7iz 例項 (預覽版): https://pages.awscloud.com/R7iz-Preview.html

[7]

瞭解亞馬遜機器映象 (AMI): https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AMIs.html

[8]

Hugging Face 平臺上的資料集頁面: SQUAD: https://huggingface.co/datasets/squad

[9]

Hugging Face 平臺上的模型頁面: DistilBERT: https://huggingface.co/distilbert-base-uncased

[10]

Intel® Extension for PyTorch 的 GitHub 頁面: https://github.com/intel/intel-extension-for-pytorch

[11]

Hugging Face 文件頁面: CPU 高效訓練: https://huggingface.co/docs/transformers/perf_train_cpu

[12]

Hugging Face 文件頁面: 多 CPU 上的高效訓練: https://huggingface.co/docs/transformers/perf_train_cpu_many

「其他文章」