加速PyTorch Transformers与英特尔Sapphire Rapids – 第1部分

加速PyTorch Transformers与英特尔Sapphire Rapids

大约一年前，我们展示了如何在集群或第三代英特尔至强可扩展CPU（也称为Ice Lake）上分发Hugging Face transformers的训练。最近，英特尔推出了第四代至强CPU，代号为Sapphire Rapids，具有令人兴奋的新指令，可以加速深度学习模型中常见的操作。

在本文中，您将学习如何使用在AWS上运行的一组Sapphire Rapids服务器加速PyTorch训练作业。我们将使用英特尔的oneAPI Collective Communications Library（CCL）来分发作业，并使用英特尔的PyTorch扩展（IPEX）库自动利用新的CPU指令。由于这两个库已经与Hugging Face transformers库集成，我们将能够无需更改代码即可直接运行示例脚本。

在后续的文章中，我们将探讨在Sapphire Rapids CPU上进行推理以及它们带来的性能提升。

为什么应该考虑在CPU上进行训练

在英特尔至强CPU上训练深度学习（DL）模型可以是一种具有成本效益且可扩展的方法，特别是在使用分布式训练和对小型和VoAGI数据集进行微调等技术时。

至强CPU支持诸如先进矢量扩展（AVX-512）和超线程等高级功能，有助于改善DL模型的并行性和效率。这使得训练时间更快，硬件资源利用率更高。

此外，与通常需要用于训练大型深度学习模型的专用硬件（例如GPU）相比，至强CPU通常更经济实惠且易于获取。至强CPU还可以轻松重新用于其他生产任务，从Web服务器到数据库，使其成为您的IT基础设施的多功能和灵活选择。

最后，云用户可以通过使用spot实例进一步降低在至强CPU上进行训练的成本。spot实例是由多余的计算容量构建的，并以折扣价格出售。与使用按需实例相比，它们可以提供显着的成本节省，有时高达90%。最后但并非最不重要的是，CPU spot实例通常比GPU实例更容易获得。

现在，让我们来看看Sapphire Rapids架构中的新指令。

高级矩阵扩展：用于深度学习的新指令

Sapphire Rapids架构引入了英特尔高级矩阵扩展（AMX）以加速DL工作负载。使用它们就像安装最新版本的IPEX一样简单。您无需更改Hugging Face代码中的任何内容。

AMX指令加速矩阵乘法，这是在数据批次上训练DL模型时的核心操作。它们支持脑浮点（BF16）和8位整数（INT8）值，可以为不同的训练场景提供加速。

AMX引入了新的二维CPU寄存器，称为瓦片寄存器。由于这些寄存器在上下文切换期间需要保存和恢复，因此它们需要内核支持：在Linux上，您需要v5.16或更新版本。

现在，让我们看看如何构建一组用于分布式训练的Sapphire Rapids CPU集群。

构建一组Sapphire Rapids CPU集群

在撰写本文时，获得Sapphire Rapids服务器的最简单方法是使用新的Amazon EC2 R7iz实例系列。由于它仍处于预览阶段，您必须注册以获取访问权限。此外，虚拟服务器尚不支持AMX，因此我们将使用裸金属实例（r7iz.metal-16xl，64 vCPU，512GB RAM）。

为了避免手动设置集群中的每个节点，我们将首先设置主节点并从中创建新的Amazon Machine Image（AMI）。然后，我们将使用此AMI启动其他节点。

从网络设置的角度来看，我们需要进行以下设置：

  • 对于设置和调试，打开所有实例上用于ssh访问的22端口。

  • 将主实例（您将从中启动训练的实例）到所有其他实例（包括主实例）的ssh设置为无密码。换句话说，主节点的ssh公钥必须在所有节点上得到授权。

  • 允许集群内的所有网络流量，以便分布式训练顺利运行。AWS提供了一种安全且便捷的方法来使用安全组进行此操作。我们只需要创建一个安全组，允许来自配置了相同安全组的实例的所有流量，并确保将其附加到集群中的所有实例上。下面是我的设置。

让我们开始工作，构建集群的主节点。

设置主节点

我们首先通过启动一个 r7iz.metal-16xl 实例来创建主节点，使用的是 Ubuntu 20.04 AMI ( ami-07cd3e6c4915b2d18 ) 和之前创建的安全组。这个 AMI 包含了 Linux v5.15.0，但是 Intel 和 AWS 已经修补了内核以添加 AMX 支持。因此，我们不需要升级内核到 v5.16。

实例运行后，我们通过 ssh 连接到它，并使用 lscpu 命令检查是否支持 AMX。你应该在 flags 部分看到以下内容：

amx_bf16 amx_tile amx_int8

然后，我们安装本地和 Python 的依赖。

sudo apt-get update 

# 为了提高性能，安装 tcmalloc (https://github.com/google/tcmalloc)
sudo apt install libgoogle-perftools-dev -y

# 创建虚拟环境
sudo apt-get install python3-pip -y
pip install pip --upgrade
export PATH=/home/ubuntu/.local/bin:$PATH
pip install virtualenv

# 激活虚拟环境
virtualenv cluster_env
source cluster_env/bin/activate

# 安装 PyTorch, IPEX, CCL 和 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

# 克隆 transformers 仓库以获取其示例脚本
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个时期大约需要26分钟。作为参考，我们在一台相似的Ice Lake实例（ c6i.16xlarge ）上使用相同的软件设置对同一作业计时，每个时期需要3小时30分钟。这是一个8倍的加速。我们已经可以看到新指令的好处了！

现在，让我们将训练作业分布在四个实例上。一个 r7iz.16xlarge 实例有32个物理CPU核心，我们更喜欢直接使用这些核心而不是使用虚拟CPU（ 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 ）。

# 设置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

现在，我们启动分布式训练作业。

# 启动分布式训练
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秒。

这是作业的外观。主节点在顶部，你可以看到其他3个节点上运行的两个训练进程。

在4个节点上的完美线性扩展将是6分钟30秒（26分钟除以4）。我们非常接近这个理想值，这显示了这种方法的可扩展性。

结论

如您所见，在一组英特尔Xeon CPU的集群上训练Hugging Face transformers是一种灵活、可扩展且具有成本效益的解决方案，特别适用于小型或VoAGI规模的模型和数据集。

以下是一些帮助您入门的附加资源：

  • Intel IPEX在GitHub上的页面
  • Hugging Face文档： ” Efficient training on CPU ” 和 ” Efficient training on many CPUs “。

如果您有任何问题或反馈，我们非常乐意在Hugging Face论坛上阅读。

感谢阅读！

使用Mask2Former和OneFormer进行通用图像分割
什么使对话代理程序有用？
使用 LoRA 进行高效稳定的扩散微调
2D资产生成：游戏开发中的人工智能 #4
Hugging Face在计算机视觉领域的现状 🤗
深入了解视觉语言模型
人工智能

人工智能

探索人工智能的世界及AI的未来潜力