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

大约一年以前，我们 展示 了如何在第三代 英特尔至强可扩展 CPU (即 Ice Lake) 集群上分布式训练 Hugging Face transformers 模型。最近，英特尔发布了代号为 Sapphire Rapids 的第四代至强可扩展 CPU，该 CPU 包含了令人兴奋的深度学习加速新指令。

通过本文，你将会学到如何在一个 AWS Sapphire Rapids 集群上加速一个 PyTorch 训练任务。我们会使用 英特尔 oneAPI 集合通信库 (oneAPI Collective Communications Library, oneCCL) 来分布式化训练任务，并使用 英特尔 PyTorch 扩展库 (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 实例家族。由于它尚在预览期，你必须 登记注册 以获得访问权限。另外，虚拟机尚未支持 AMX，因此，我们将使用裸金属实例 (r7iz.metal-16xl, 64 vCPU, 512GB RAM) 。