加速 PyTorch Transformers 在 Intel Sapphire Rapids 上的训练 – 第 1 部分

TL;DR

Hugging Face 演示了在四节点 Intel Sapphire Rapids CPU 集群上使用 Intel Extension for PyTorch (IPEX) 和 oneCCL 训练 PyTorch Transformers，可比同等的 Ice Lake 实例提升 8 倍速度，并且几乎线性扩展，使基于 CPU 的训练成为成本更低的 GPU 替代方案。

为什么在 CPU 上训练？

在 Intel Xeon CPU 上训练相较于基于 GPU 的训练成本更低且更具可扩展性，尤其是针对中小规模模型和数据集。Xeon CPU 提供 AVX‑512、超线程以及现在的高级矩阵扩展 (AMX)，可加速矩阵乘法。CPU 资源普遍可得，可在完成训练后重新用于其他工作负载，且云端 Spot 实例的费用相比按需实例可降低高达 90%。

高级矩阵扩展 (AMX)

Sapphire Rapids 引入了 AMX，这是一组加速矩阵乘法（深度学习训练核心操作）的新指令。AMX 支持 BF16 和 INT8 数据类型，并增加了二维瓦片寄存器。要使用 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 RAM 的裸金属实例。

网络设置

  • 在所有节点上打开 SSH（端口 22）。
  • 配置从主节点到每个节点（包括主节点本身）的免密码 SSH。
  • 创建一个安全组，允许集群内部无限制流量，并将其关联到所有实例。

设置主节点

  1. 启动一台 r7iz.metal-16xl 实例，使用 Ubuntu 20.04 AMI ami-07cd3e6c4915b2d18。
  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

启动分布式训练

本示例在 SQuAD 数据集上微调 DistilBERT。

单节点基准

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 倍加速。

四节点分布式运行

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 进程 × 4 节点
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           # 每进程训练线程数

使用 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 分钟。原文截图展示了每节点两个训练进程以及主进程负责调度运行。

意义与要点

  • 成本效益 – CPU Spot 实例的费用远低于 GPU 实例，同时在多数工作负载下仍能提供相当的训练速度。
  • 可扩展性 – 四个 Sapphire Rapids 节点的近线性扩展表明 IPEX 与 oneCCL 的组合能够高效分配计算与通信。
  • 易用性 – 无需修改代码，只需开启 --use_ipex 与 --bf16 即可自动利用新的 AMX 指令。
  • 后续工作 – 作者计划在后续文章中介绍 Sapphire Rapids 上的推理性能，以完整呈现端到端的 CPU 加速方案。

进一步资源

Sources