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

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

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

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

  1. 该库执行 AllReduce,这是分布式训练期间的一项关键操作,占据了很大一部分通信开销。

  2. 库充分利用 Amazon 的网络基础设施和 Amazon EC2 实例拓扑,来执行优化的节点到节点通信。

使用此数据并行库可将训练模型(例如 BERT)的速度提高多达 25%。像 Horovod 这样的实施可在大规模下提供亚线性性能,而该库可在大规模下提供近线性的性能。这意味着您可以缩短训练时间并降低模型训练的成本。

注意

SageMaker 分布式训练库只能通过 Amazon Deep Learning Containers,供 SageMaker 训练平台中的 TensorFlow、PyTorch 和 HuggingFace 框架使用。要使用这些库,您必须使用 SageMaker Python SDK,或者通过 SDK for Python (Boto3) 或 Amazon Command Line Interface 来使用 SageMaker API。全文中的说明和示例均重点介绍如何通过 SageMaker Python SDK 使用分布式训练库。

通过平衡融合缓冲区实现最佳带宽使用

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

传统参数服务器的一个主要缺点是它们对可用网络带宽的利用率不佳。参数服务器将变量视为原子单元,并将每个变量放在一台服务器上。由于梯度在向后传递期间会依次可用,因此在任何给定时刻,在不同服务器上发送和接收的数据量都存在不平衡。一些服务器接收和发送的数据较多,一些服务器则较少,甚至一些服务器没有流量。随着参数服务器数量的增加,这个问题会更加严重。

该库引入了平衡融合缓冲区来解决这些问题。平衡融合缓冲区是 GPU 中的一个缓冲区,用于存放梯度,直到缓冲区的大小超过阈值。在具有 N 个参数服务器的设置中,当缓冲区超过阈值时,平衡融合缓冲区将复制到 CPU 内存,分片为 N 个部分,并将第 i 部分发送到第 i 个参数服务器。每台服务器从平衡融合缓冲区接收完全相同的字节数。第 i 个服务器从所有工作线程接收平衡融合缓冲区的第 i 个分区,将它们累加起来,然后将结果发送回所有工作线程。由于所有服务器都平等地参与对每个平衡融合缓冲区取平均值中,因此可以高效地利用服务器带宽。

通过对向后传递的高效 AllReduce 重叠实现最佳 GPU 使用率

SageMaker 的分布式数据并行库实现了 AllReduce 操作与向后传递的最佳重叠,可显著提高 GPU 利用率,并通过优化 CPU 与 GPU 之间的任务来实现近线性的缩放效率和更快的训练速度。该库在 GPU 计算梯度时并行执行 AllReduce,不会占用任何额外的 GPU 周期,这使得库的速度更快。

  • 利用 CPU:该库使用 CPU 来处理 AllReduce 梯度,从而将此任务从 GPU 中分载出去。

  • 提高 GPU 使用率:集群的 GPU 专门用于计算梯度,从而在整个训练过程中提高其利用率。

SageMaker 分布式数据并行架构

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

  1. 该库将秩分配给 GPU(工作线程)。

  2. 每次迭代时,库都会将每个全局批次除以工作线程总数(世界大小),然后将小批次(批次分片)分配给工作线程。

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

    • 批次分片(小批次)是每次迭代分配给每个 GPU(工作线程)的数据集的子集。

  3. 该库会为每个工作线程启动训练脚本。

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

  5. 库可同步各工作线程的模型权重和梯度,以聚合成为单个经过训练的模型。

以下架构图显示了该库如何为由 3 个节点组成的集群设置数据并行性的示例。


        SageMaker 分布式数据并行架构图

要开始使用 SageMaker 分布式数据并行库,请参阅步骤 2:使用 SageMaker Python SDK 启动 SageMaker 分布式训练作业,以通过 Amazon SageMaker Python SDK 设置 SageMaker 估算器,并参阅使用模型并行性运行 SageMaker 分布式训练作业以使用 SageMaker 分布式数据并行库调整训练脚本。