Intel Sapphire Rapids を使用した PyTorch Transformers の高速化 – パート 1

TL;DR

Hugging Face は、Intel Sapphire Rapids CPU の 4 ノードクラスタ上で Intel Extension for PyTorch (IPEX) と oneCCL を使用して PyTorch Transformers をトレーニングすると、同等の Ice Lake インスタンスに比べて 8 倍の高速化が得られ、ほぼ線形にスケールすることを示しています。これにより、CPU ベースのトレーニングは GPU のコスト効果の高い代替手段となります。

なぜ CPU でトレーニングするのか？

Intel Xeon CPU でのトレーニングは、特に小規模から中規模のモデルやデータセットに対して、GPU ベースのトレーニングよりもコストが低く、スケーラビリティが高くなることがあります。Xeon CPU は AVX‑512、ハイパースレッディング、そして現在は Advanced Matrix Extensions (AMX) を提供し、行列乗算を高速化します。CPU は広く入手可能で、他のワークロードに再利用でき、クラウドのスポットインスタンスはオンデマンドインスタンスに比べて最大 90% コスト削減が可能です。

Advanced Matrix Extensions (AMX)

Sapphire Rapids は AMX を導入しました。これは行列乗算（ディープラーニングトレーニングの核心的な演算）を高速化する新しい命令セットです。AMX は BF16 と INT8 データ型をサポートし、2 次元タイルレジスタを追加します。AMX を使用するには Linux カーネルがバージョン 5.16 以上である必要がありますが、本ガイドで使用するベアメタル r7iz.metal-16xl インスタンス向けに、Intel と AWS が必要なサポートをカーネル 5.15 にバックポートしています。

AWS で Sapphire Rapids クラスタを構築する

Sapphire Rapids ハードウェアにアクセスする最も簡単な方法は、Amazon EC2 の R7iz ベアメタルインスタンス（例：r7iz.metal-16xl）を利用することです。プレビュー段階では仮想マシンでの AMX がまだサポートされていないため、本ガイドでは 64 vCPU と 512 GB RAM を持つベアメタルインスタンスを使用します。

ネットワーク設定

すべてのノードで SSH（ポート 22）を開放する。
マスターノードから全ノード（マスター自身を含む）へのパスワードなし SSH を設定する。
クラスタ内部のトラフィックを無制限に許可するセキュリティグループを作成し、すべてのインスタンスにアタッチする。

マスターノードのセットアップ

Ubuntu 20.04 AMI ami-07cd3e6c4915b2d18 を使用して r7iz.metal-16xl インスタンスを起動する。
lscpu で AMX のサポートを確認する（フラグ amx_bf16 amx_tile amx_int8）。

依存関係をインストールする:

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

SSH キーペア（ssh-keygen）を生成し、~/.ssh/cluster として保存する。
設定済みインスタンスから AMI を作成し、再利用できるようにする。

残りのノードのセットアップ

上記で作成した AMI を使用して、追加の r7iz.metal-16xl インスタンスを 3 台起動する。
マスター上で ~/.ssh/config を編集し、ホスト node1、node2、node3 をそれぞれのプライベート IP と cluster キーで定義する。
ssh node[1-3] でパスワードなしアクセスを確認する。
~/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

単一エポックは ≈26 分 で完了し、Ice Lake の c6i.16xlarge インスタンスでの ≈3 時間 30 分 と比較して 8 倍の高速化 となります。

4 ノードでの分散実行

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

エポック時間は 7 分 30 秒 に短縮され、理想的な線形スケーリング目標である 6 分 30 秒（26 分 ÷ 4）よりわずか 1 分だけ上回っています。元記事のスクリーンショットでは、各ノードで 2 つのトレーニングプロセスが実行され、マスタープロセスが実行を調整している様子が示されています。

考察と要点

コスト効率 – CPU スポットインスタンスは GPU インスタンスに比べて劇的に安価であり、多くのワークロードで同等のトレーニング速度を提供できます。
スケーラビリティ – 4 つの Sapphire Rapids ノード間でのほぼ線形スケーリングは、IPEX と oneCCL の組み合わせが計算と通信の両方を効果的に分散できることを示しています。
使いやすさ – コードの変更は不要で、--use_ipex と --bf16 を有効にするだけで新しい AMX 命令が自動的に有効化されます。
今後の課題 – 著者は、Sapphire Rapids 上での推論性能に関するフォローアップ記事を予定しており、エンドツーエンドの CPU 加速の全体像が完成します。

追加リソース

Intel Extension for PyTorch (IPEX): https://github.com/intel/intel-extension-for-pytorch
Hugging Face の CPU トレーニングガイド:
- Efficient training on CPU – https://huggingface.co/docs/transformers/perf_train_cpu
- Efficient training on many CPUs – https://huggingface.co/docs/transformers/perf_train_cpu_many
ディスカッションフォーラム: https://discuss.huggingface.co/