使用 DeepSpeed 加速 PyTorch，使用基于英特尔哈瓦那高迪的 DL1 EC2 实例训练大型语言模型

训练具有数十亿个参数的大型语言模型 (LLM) 可能具有挑战性。除了设计模型架构外，研究人员还需要为分布式训练设置最先进的训练技术，例如混合精度支持、梯度累积和检查点。对于大型模型，训练设置更具挑战性，因为单个加速器设备中的可用内存限制了仅使用数据并行度训练的模型的大小，并且使用模型并行训练需要对训练代码进行额外的修改。诸如 DeepSpeed （PyTorch 的开源深度学习优化库）之类的库可以解决其中的一些挑战，可以帮助加速模型开发和训练。

在这篇文章中，我们对基于英特尔哈瓦那高迪的 亚马逊弹性计算云 （Amazon EC2） DL1 实例进行了培训，并量化了使用DeepSpeed等扩展框架的好处。我们提供了编码器型变压器模型（具有3.4亿至15亿个参数的BERT）的缩放结果。 对于具有 15 亿个参数的模型，我们使用 DeepSpeed Zero 第 1 阶段优化，在 128 个加速器（16 个 dl1.24xlarge 实例）上实现了 82.7% 的扩展效率。 优化器状态由 DeepSpeed 进行分区，以使用数据并行模式训练大型模型。这种方法已扩展到使用数据并行性训练一个 50 亿个参数的模型。与使用 FP32 数据类型相比，我们还使用了 Gaudi 对 BF16 数据类型的原生支持，从而减小了内存大小并提高了训练性能。 结果，我们使用 wikicorpus-en 数据集在 16 小时内实现了 BERT 15 亿参数模型的预训练（第 1 阶段）模型收敛（我们的目标是在一天之内训练一个大型模型）。

训练设置

我们使用 亚马逊云科技 Batch 预置了一个由 16 个 dl1.24xlarge 实例组成的托管计算集群。 我们开发了一 个 亚马逊云科技 Batch 研讨会 ，演示了使用 亚马逊云科技 Batch 设置分布式训练集群的步骤。 每个 dl1.24xlarge 实例都有八个 Habana Gaudi 加速器，每个 Habana Gaudi 加速器具有 32 GB 的内存，每个卡之间有一个全网状 RoCE 网络，每个卡的双向互连总带宽为 700 Gbps（有关更多信息，请参阅 Amazon EC2 DL1 实例 Deep Dl1）。 dl1.24xlarge 集群还使用了四个 亚马逊云科技 弹性结构适配器 (EFA)，节点之间的互连速度总共为 400 Gbps。

分布式培训研讨会说明了设置分布式训练集群的步骤。研讨会展示了使用 亚马逊云科技 Batch 的分布式训练设置，尤其是在完全托管的集群上启动大规模容器化训练任务的多节点并行任务功能。更具体地说，使用 DL1 实例创建了一个完全托管的 亚马逊云科技 Batch 计算环境。容器从 Amazon 弹性容器注册表 (Amazon ECR) 中提取，并根据多节点并行任务定义自动启动到集群中的实例中。研讨会最后使用 PyTorch 和 DeepSpeed 对 BERT（3.4 亿到 15 亿个参数）模型进行多节点、多 HPU 数据并行训练。

使用 DeepSpeed 进行预训练

Habana SynapseAI v1.5 和 v 1.6 支持 Deep Speed Zero1 优化。 DeepSpeed GitHub 存储库的 哈瓦那分支 包括支持高迪加速 器所需的修改。完全支持分布式数据并行（多卡、多实例）、Zero1 优化和 BF16 数据类型。

所有这些功能都在 B ERT 1.5B 模型参考库中启用 ，该库 引入了 48 层、1600 个隐藏尺寸和 25 头双向编码器模型，该模型源自 BERT 实现。该存储库还包含基准的 BERT 大型模型实现：一个 24 层、1024 个隐藏、16 头、3.4 亿个参数的神经网络架构。预训练建模脚本源自 NVIDIA 深度学习示例存储库 ， 用于下载 wikicorpus_en 数据，将原始数据预处理成令牌，并将数据分片成较小的 h5 数据集以进行分布式数据并行训练。您可以采用这种通用方法，使用 DL1 实例使用您的数据集来训练自定义 PyTorch 模型架构。

预训练（第 1 阶段）扩展结果

对于大规模预训练大型模型，我们主要关注解决方案的两个方面：训练性能（以训练时间衡量）和达成完全融合解决方案的成本效益。接下来，我们以 BERT 1.5B 预训练为例，深入研究这两个指标。

扩展性能和训练时间

我们首先衡量 BERT Large 实现的性能，以此作为可扩展性的基准。下表列出了从 1-8 个 dl1.24xlarge 实例（每个实例有八个加速器设备）测得的每秒序列吞吐量。我们使用单实例吞吐量作为基准，测量了跨多个实例扩展的效率，这是理解性价比训练指标的重要杠杆。

下图说明了扩展效率。

对于 BERT 1.5B，我们在参考存储库中修改了模型的超参数以保证收敛。每个加速器的有效批次大小设置为 384（以实现最大内存利用率），微批次为每步 16 个，梯度累积为 24 步。8 个和 16 个节点分别使用了 0.0015 和 0.003 的学习率。通过这些配置，我们在大约 25 小时内在 8 个 dl1.24xlarge 实例（64 个加速器）上实现了 BERT 1.5B 第 1 阶段预训练的趋同，在 16 个 dl1.24xlarge 实例（128 个加速器）上实现了 15 小时的收敛。下图显示了当我们扩大加速器数量时，平均损失与训练周期数的函数。

通过前面描述的配置，我们在单个实例中使用 8 个加速器获得了 85% 的强扩展效率，使用 128 个加速器获得了 83% 的强扩展效率。下表汇总了这些参数。

下图说明了扩展效率。

结论

在这篇文章中，我们评估了哈瓦那 SynapseAI v1.5/v1.6 对 DeepSpeed 的支持，以及它如何帮助在哈瓦那高迪加速器上扩展 LLM 训练。15亿个参数的BERT模型的预训练花了16个小时才汇聚到由128个高迪加速器组成的集群上，扩展能力强度为85％。我们鼓励您看看 A WS 研讨 会上演示的架构， 并考虑采用它来训练使用 DL1 实例的自定义 PyTorch 模型架构。

作者简介

Mahadevan Balasubramaniam 是自主计算首席解决方案架构师，在注入物理学的深度学习、大规模构建和部署工业系统数字双胞胎领域拥有近 20 年的经验。Mahadevan 在麻省理工学院获得机械工程博士学位，拥有超过 25 项专利和出版物。

RJ 是 Search M5 团队的工程师，负责构建用于训练和推理的大规模深度学习系统。工作之余，他探索不同的食物并玩球拍运动。

桑达尔·兰加纳森 是亚马逊 EC2 团队机器学习框架业务开发主管。他专注于亚马逊 EKS、亚马逊 ECS、Elastic Fabric Adapter、亚马逊云科技 Batch 和亚马逊 SageMaker 等亚马逊服务中的大规模机器学习工作负载。他的经验包括在NetApp、美光科技、高通和Mentor Graphics的产品管理和产品开发方面担任领导职务。

Abhinandan Patni 是亚马逊搜索的高级软件工程师。他专注于为可扩展的分布式深度学习训练和实时推理构建系统和工具。

Pierre-Yves Aquilanti 是亚马逊网络服务框架机器学习解决方案主管，他帮助开发业界最好的基于云的机器学习框架解决方案。他的背景是高性能计算，在加入 亚马逊云科技 之前，Pierre-Yves 曾在石油和天然气行业工作。Pierre-Yves 来自法国，拥有里尔大学计算机科学博士学位。