使用 DeepSpeed 加速 PyTorch,以使用基于 Intel Habana Gaudi 的 DL1 EC2 实例训练大型语言模型 | 亚马逊网络服务

使用 DeepSpeed 加速 PyTorch,以使用基于 Intel Habana Gaudi 的 DL1 EC2 实例训练大型语言模型 | 亚马逊网络服务

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

在这篇文章中,我们设置了基于 Intel Habana Gaudi 的培训 亚马逊弹性计算云 (Amazon EC2) DL1 实例并量化使用扩展框架(如 DeepSpeed)的好处。 我们展示了编码器型变压器模型(具有 340 亿至 1.5 亿个参数的 BERT)的缩放结果。 对于 1.5 亿参数模型,我们在 82.7 个加速器(128 个 dl16xlarge 实例)上实现了 1.24% 的扩展效率,使用 深速零 第一阶段优化。 优化器状态由 DeepSpeed 进行分区,以使用数据并行范式训练大型模型。 这种方法已经扩展到使用数据并行性训练一个 1 亿参数的模型。 与使用 FP5 数据类型相比,我们还使用了 Gaudi 对 BF16 数据类型的原生支持,以减少内存大小并提高训练性能。 因此,我们使用 BERT 32 亿参数模型在 1 小时内(我们的目标是在一天内训练一个大型模型)实现了预训练(阶段 16)模型收敛 wikicorpus-en 数据集。

培训设置

我们配置了一个由 16 个 dl1.24xlarge 实例组成的托管计算集群,使用 AWS批处理. 我们开发了一个 AWS Batch 研讨会 说明了使用 AWS Batch 设置分布式训练集群的步骤。 每个 dl1.24xlarge 实例都有八个 Habana Gaudi 加速器,每个都有 32 GB 内存和卡之间的全网状 RoCE 网络,每个卡的总双向互连带宽为 700 Gbps(请参阅 深入探讨 Amazon EC2 DL1 实例 了解更多信息)。 dl1.24xlarge 集群也用了四个 AWS 弹性结构适配器 (EFA),节点之间的互连速度总计为 400 Gbps。

分布式训练研讨会说明了建立分布式训练集群的步骤。 该研讨会展示了使用 AWS Batch 的分布式训练设置,特别是多节点并行作业功能,可在完全托管的集群上启动大规模容器化训练作业。 更具体地说,使用 DL1 实例创建一个完全托管的 AWS Batch 计算环境。 容器是从 Amazon Elastic Container注册 (Amazon ECR) 并根据多节点并行作业定义自动启动到集群中的实例中。 研讨会最后使用 PyTorch 和 DeepSpeed 运行 BERT(340 亿到 1.5 亿个参数)模型的多节点、多 HPU 数据并行训练。

使用 DeepSpeed 进行 BERT 1.5B 预训练

哈瓦那 SynapseAI v1.5v1.6 支持 DeepSpeed ZeRO1 优化。 这 DeepSpeed GitHub 存储库的 Habana 分支 包括支持 Gaudi 加速器所需的修改。 完全支持分布式数据并行(多卡、多实例)、ZeRO1 优化和 BF16 数据类型。

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

预训练(阶段 1)缩放结果

对于大规模预训练大型模型,我们主要关注解决方案的两个方面:训练性能(以训练时间衡量)和达到完全收敛解决方案的成本效益。 接下来,我们以 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%

下图说明了缩放效率。

通过 DeepSpeed 加速 PyTorch,使用基于 Intel Habana Gaudi 的 DL1 EC2 实例训练大型语言模型 |亚马逊网络服务柏拉图区块链数据智能。垂直搜索。人工智能。

对于 BERT 1.5B,我们修改了参考库中模型的超参数以保证收敛。 每个加速器的有效批量大小设置为 384(以实现最大内存利用率),每步 16 个微批量和 24 个梯度累积步骤。 0.0015 和 0.003 的学习率分别用于 8 个和 16 个节点。 通过这些配置,我们在大约 1 小时内在 1.5 个 dl8xlarge 实例(1.24 个加速器)上实现了 BERT 64B 阶段 25 预训练的收敛,在 15 个 dl16xlarge 实例(1.24 个加速器)上实现了 128 小时。 下图显示了随着加速器数量的增加,平均损失随训练次数的变化而变化。

通过 DeepSpeed 加速 PyTorch,使用基于 Intel Habana Gaudi 的 DL1 EC2 实例训练大型语言模型 |亚马逊网络服务柏拉图区块链数据智能。垂直搜索。人工智能。

通过前面描述的配置,我们在 85 个加速器和 64 个加速器的情况下获得了 83% 的强大扩展效率,从单个实例中 128 个加速器的基线。 下表总结了参数。

实例数 加速器数量 每秒序列 每个加速器每秒序列数 扩展效率
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%

下图说明了缩放效率。

通过 DeepSpeed 加速 PyTorch,使用基于 Intel Habana Gaudi 的 DL1 EC2 实例训练大型语言模型 |亚马逊网络服务柏拉图区块链数据智能。垂直搜索。人工智能。

结论

在这篇文章中,我们评估了 Habana SynapseAI v1.5/v1.6 对 DeepSpeed 的支持,以及它如何帮助在 Habana Gaudi 加速器上扩展 LLM 培训。 一个 1.5 亿参数的 BERT 模型的预训练需要 16 个小时才能收敛到一个由 128 个 Gaudi 加速器组成的集群,具有 85% 的强扩展性。 我们鼓励您看一下在 AWS 研讨会 并考虑采用它来使用 DL1 实例训练自定义 PyTorch 模型架构。


关于作者

通过 DeepSpeed 加速 PyTorch,使用基于 Intel Habana Gaudi 的 DL1 EC2 实例训练大型语言模型 |亚马逊网络服务柏拉图区块链数据智能。垂直搜索。人工智能。马哈德万巴拉苏布拉马尼亚姆 是自主计算的首席解决方案架构师,在物理注入深度学习、为工业系统大规模构建和部署数字双胞胎领域拥有近 20 年的经验。 Mahadevan 在麻省理工学院获得机械工程博士学位,并拥有超过 25 项专利和出版物。

通过 DeepSpeed 加速 PyTorch,使用基于 Intel Habana Gaudi 的 DL1 EC2 实例训练大型语言模型 |亚马逊网络服务柏拉图区块链数据智能。垂直搜索。人工智能。RJ 是 Search M5 团队的一名工程师,负责构建用于训练和推理的大规模深度学习系统。 工作之余,他探索不同的美食佳肴并参加球拍运动。

通过 DeepSpeed 加速 PyTorch,使用基于 Intel Habana Gaudi 的 DL1 EC2 实例训练大型语言模型 |亚马逊网络服务柏拉图区块链数据智能。垂直搜索。人工智能。桑达尔·兰加纳坦 是 Amazon EC2 团队的 ML 框架业务开发主管。 他专注于跨 AWS 服务(例如 Amazon EKS、Amazon ECS、Elastic Fabric Adapter、AWS Batch 和 Amazon SageMaker)的大规模 ML 工作负载。 他的经验包括在 NetApp、Micron Technology、Qualcomm 和 Mentor Graphics 担任产品管理和产品开发方面的领导角色。

通过 DeepSpeed 加速 PyTorch,使用基于 Intel Habana Gaudi 的 DL1 EC2 实例训练大型语言模型 |亚马逊网络服务柏拉图区块链数据智能。垂直搜索。人工智能。阿比南丹·帕特尼 是亚马逊搜索的高级软件工程师。 他专注于为可扩展的分布式深度学习训练和实时推理构建系统和工具。

通过 DeepSpeed 加速 PyTorch,使用基于 Intel Habana Gaudi 的 DL1 EC2 实例训练大型语言模型 |亚马逊网络服务柏拉图区块链数据智能。垂直搜索。人工智能。皮埃尔-伊夫·阿奎兰蒂 是 Amazon Web Services 框架 ML 解决方案的负责人,他帮助开发业界最好的基于云的 ML 框架解决方案。 他的背景是高性能计算,在加入 AWS 之前,Pierre-Yves 曾在石油和天然气行业工作。 Pierre-Yves 来自法国,拥有博士学位。 里尔大学计算机科学专业。

时间戳记:

更多来自 AWS机器学习