训练拥有数十亿个参数的大型语言模型（LLMs）可能会很具有挑战性。除了设计模型架构外，研究人员需要设置最先进的分布式训练技术，如混合精度支持、梯度累积和检查点。对于大型模型，训练设置更具有挑战性，因为单个加速器设备中可用的内存会限制使用仅数据并行性训练的模型大小，并且使用模型并行训练需要对训练代码进行额外的修改。诸如 DeepSpeed（一种针对 PyTorch 的开源深度学习优化库）之类的库解决了其中一些挑战，并可以帮助加速模型开发和训练。
在本文中，我们在基于 Intel Habana Gaudi 的 Amazon Elastic Compute Cloud（Amazon EC2）DL1 实例上设置了训练，并量化了使用 DeepSpeed 等缩放框架的好处。我们为编码器类型的 transformer 模型（340 million 到 1.5 billion 参数的 BERT）呈现了扩展结果。对于 15 亿参数模型，我们使用 DeepSpeed ZeRO 阶段 1 优化在 128 个加速器（16 个 dl1.24xlarge 实例）上实现了 82.7% 的缩放效率。DeepSpeed 对优化器状态进行了分区，以使用数据并行范例训练大型模型。这种方法已经扩展到使用数据并行训练 50 亿参数模型。我们还使用 Gaudi 对 BF16 数据类型的本地支持，以减少内存大小并提高训练性能，与使用 FP32 数据类型相比。因此，我们使用 wikicorpus-en 数据集在 16 小时内实现了预训练（第一阶段）模型收敛（我们的目标是在一天内训练大型模型）。
训练设置
我们使用 AWS Batch 配置了一个由 16 个 dl1.24xlarge 实例组成的托管计算集群。我们开发了一个 AWS Batch 工作坊，演示了设置使用 AWS Batch 的分布式训练集群的步骤。每个 dl1.24xlarge 实例都有八个 Habana Gaudi 加速器，每个加速器的内存为 32 GB，并且卡之间具有全网 RoCE 网络，每个加速器的总双向互连带宽为 700 Gbps（有关更多信息，请参见 Amazon EC2 DL1 实例深入介绍）。dl1.24xlarge 集群还使用了四个 AWS 弹性面料适配器（EFA），节点之间总共有 400 Gbps 的互连。
分布式训练工作坊演示了设置分布式训练集群的步骤。该工作坊展示了使用 AWS Batch 进行分布式训练设置，特别是多节点并行作业功能，以在完全托管的集群上启动大规模容器化训练作业。更具体地说，创建了一个完全托管的 AWS Batch 计算环境，其中包括 DL1 实例。容器从 Amazon Elastic Container Registry（Amazon ECR）中提取，并根据多节点并行作业定义自动启动到群集中的实例中。工作坊通过使用 PyTorch 和 DeepSpeed 运行一个多节点、多 HPU 数据并行训练的 BERT（340 million 到 1.5 billion 参数）模型来结束。
使用 DeepSpeed 进行 BERT 1.5B 预训练
Habana SynapseAI v1.5 和 v1.6 支持 DeepSpeed ZeRO1 优化。DeepSpeed GitHub 存储库的 Habana 分支包括支持 Gaudi 加速器所必需的修改。支持分布式数据并行（多卡、多实例）、ZeRO1 优化和 BF16 数据类型。
所有这些功能都在 BERT 1.5B 模型参考存储库上启用，该存储库介绍了一个 48 层、1600 隐藏维度和 25 个双向编码器模型，派生自 BERT 实现。该存储库还包含基线 BERT Large 模型实现：24 层、1024 隐藏、16 个头、340 million 参数神经网络架构。预训练建模脚本源自 NVIDIA 深度学习示例存储库，用于下载 wikicorpus_en 数据，将原始数据预处理为标记，并将数据分成较小的 h5 数据集以进行分布式数据并行训练。您可以采用此通用方法，使用 DL1 实例使用您的数据集来训练自定义 PyTorch 模型架构。
预训练（第一阶段）扩展结果
对于大规模预训练大型模型，我们主要关注解决方案的两个方面：训练性能（按训练时间测量）和成本效益。下面，我们以 BERT 1.5B 预训练为例，深入介绍这两个指标。
扩展性能和训练时间
我们首先通过衡量BERT Large实现的性能作为可扩展性的基线。下表列出了从1-8 dl1.24xlarge实例（每个实例有八个加速器设备）的每秒序列量的测量值。使用单实例吞吐量作为基准，我们测量了跨多个实例的扩展效率，这是理解价格性能训练指标的重要杠杆。
| 实例数 | 加速器数 | 每秒序列数 | 每个加速器的每秒序列数 | 扩展效率 |
| 1 | 8 | 1,379.76 | 172.47 | 100.0% |
| 2 | 16 | 2,705.57 | 169.10 | 98.04% |
| 4 | 32 | 5,291.58 | 165.36 | 95.88% |
| 8 | 64 | 9,977.54 | 155.90 | 90.39% |
以下图表显示了扩展效率。
对于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个加速器的基线中获得了64个加速器的85％的强扩展效率和128个加速器的83％。以下表格总结了参数。
| 实例数 | 加速器数 | 每秒序列数 | 每个加速器的每秒序列数 | 扩展效率 |
| 1 | 8 | 276.66 | 34.58 | 100.0% |
| 8 | 64 | 1,883.63 | 29.43 | 85.1% |
| 16 | 128 | 3,659.15 | 28.59 | 82.7% |
以下是展示缩放效率的图示：
结论
在本文中，我们评估了Habana SynapseAI v1.5/v1.6对DeepSpeed的支持以及它如何帮助在Habana Gaudi加速器上扩展LLM训练。使用128个Gaudi加速器集群进行15亿参数BERT模型的预训练，收敛时间为16小时，强缩放达到85%。我们鼓励您查看AWS讲习班中展示的架构，并考虑采用它来使用DL1实例训练自定义PyTorch模型架构。