SageMaker 的分布式数据并行库简介 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅中国的 Amazon Web Services 服务入门

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

SageMaker 的分布式数据并行库简介

为什么要使用 SageMaker 分布式数据并行库?

SageMaker 的分布式数据并行库通过两种方式解决通信开销问题:

  1. 库执行AllReduce,这是分布式培训期间的一项关键操作,负责大部分通信开销。

  2. 该库通过充分利用Amazon的网络基础设施和 Amazon EC2 实例拓扑。

使用此数据并行库可将训练模型(如 BERT)中的速度提高高达 25%。虽然像 Horovod 这样的实现提供了大规模的亚线性能,但该库可以大规模提供近线性的性能。这意味着您可以获得更快的训练时间和更低的成本来训练模型。

注意

SageMaker 分布式培训库只能通过Amazon深度学习容器,用于 SageMaker 培训平台中的 TensorFlow、PyTorch 和 HugingFace 框架。要使用这些库,您必须使用 SageMaker Python 软件开发工具包或通过 SDK for Python (Boto3)(Bto3)或Amazon Command Line Interface. 在整个文档中,说明和示例侧重于如何将分布式培训库与 SageMaker Python SDK 结合使用。

训练基准

使用 SageMaker 的数据并行库进行 PyTorch

使用实例类型p3dn.24xl和 2、4 和 8 个节点群集:

  • 伯特:当与 PyTorch 一起使用时,SageMaker 库比火炬 DDP 快 41%、52% 和 13%。

  • 马斯克林:当与 PyTorch 一起使用时,SageMaker 库比火炬 DDP 快 4%、19% 和 15%。

这些基准测试在 PyTorch v1.6 上运行时使用ml.p3dn.24xlarge实例。您可以在SageMaker 示例网站. 示例网站还有这些模型的基准训练代码使用 TensorFlow 2.3.

平衡融合缓冲器的最佳带宽使用

SageMaker 的分布式数据并行库使用类似于参数服务器的通信模式,以减少传输的数据量和平均来自多个 GPU 的梯度所涉及的步数。它还使用称为平衡融合缓冲区的新技术,以最佳利用群集中所有节点的可用带宽。

传统参数服务器的一个关键缺点是它们对可用网络带宽的利用不佳。参数服务器将变量视为原子单位,并将每个变量放置在一台服务器上。由于渐变在向后传递期间顺序可用,因此在任何给定的时刻,从不同服务器发送和接收的数据量都存在不平衡。有些服务器正在接收和发送更多的数据,有些则更少,有些则没有。随着参数服务器数量的增加,此问题变得更糟。

该库通过引入均衡融合缓冲区. 平衡融合缓冲区是 GPU 中的缓冲区,它保存渐变,直到缓冲区的大小超过阈值。在使用 N 个参数服务器的设置中,当缓冲区超过阈值时,平衡融合缓冲区将复制到 CPU 内存中,分片为 N 个部分,并将第 i 个部分发送到第 i 个参数服务器。每个服务器从平衡融合缓冲区接收完全相同的字节数。第 i 个服务器从所有工作线程接收平衡融合缓冲区的第 i 个分区,对它们进行总和,然后将结果发送回所有工作线程。由于所有服务器均参与平均每个平衡融合缓冲区的平均值,因此可以有效地利用服务器带宽。

高效利用 GPU 的最佳使用AllReduce与向后过程重叠

SageMaker 的分布式数据并行库可实现AllReduce操作,显著提高了 GPU 利用率,并通过优化 CPU 和 GPU 之间的任务实现近线性扩展效率和缩短训练时间。库执行AllReduce,而 GPU 正在计算渐变,而不会占用额外的 GPU 周期,这使得库更快。

  • 利用 CPU:库使用 CPU 来完成所有减少渐变,从 GPU 卸载此任务。

  • 改进的 GPU 利用率:集群的 GPU 专注于计算梯度,在整个培训过程中提高了它们的利用率。

SageMaker 分布式数据并行架构

该库支持每个节点具有 8 个 GPU 的较大计算实例:ml.p3.16xlargeml.p3dn.24xlarge, 和ml.p4d.24xlarge. SageMaker 分布式数据并行库的高级工作流程如下所示:

  1. 库将等级分配给 GPU(工作人员)。

  2. 在每次迭代中,库将每个全局批次除以工作人员总数(世界大小),并将小批次(批处理分片)分配给工作线程。

    • 全局批处理的大小为(number of nodes in a cluster) * (number of GPUs per node) * (per batch shard).

    • 批处理分片(小批处理)是每次迭代分配给每个 GPU(worker)的数据集的子集。

  3. 库在每个工作人员上启动一个训练脚本。

  4. 该库在每次迭代结束时管理来自工作线程的模型权重和渐变副本。

  5. 该库同步工作线程中的模型权重和梯度,以聚合单个训练过的模型。

下面的体系结构图显示了库如何为 3 个节点的集群设置数据并行性的示例。


        SageMaker 分布式数据并行架构图

要开始使用 SageMaker 分布式数据并行库,请参阅运行 SageMaker 分布式数据并行培训 Job来设置一个 SageMaker 估计器,通过Amazon SageMaker Python 开发工具包, 和使用 SageMaker 数据并行库修改训练脚本使用 SageMaker 分布式数据并行库调整训练脚本。