利用英特尔 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 个节点上分别运行的两个训练进程。