SageMaker 分布式数据并行库简介 - Amazon SageMaker

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

SageMaker 分布式数据并行库简介

SageMaker 分布式数据并行度 (SMDDP) 库是一个集体通信库,可提高分布式数据并行训练的计算性能。SMDDP 库通过提供以下内容来解决关键集体通信操作的通信开销。

  1. 该库AllReduce针对以下内容进行了优化 AWS。 AllReduce是一项关键操作,用于在分布式数据训练期间每次训练迭代结束时同步各个 GPU 的梯度。

  2. 该库AllGather针对以下内容进行了优化 AWS。 AllGather是分片数据并行训练中使用的另一项关键操作,这是一种节省内存的数据并行技术,由 SageMaker 模型并行度 (SMP) 库、 DeepSpeed 零冗余优化器 (ZerO) 和完全分片数据并行度 (FSDP) 等常用库提供。 PyTorch

  3. 该库通过充分利用 AWS 网络基础设施和 Amazon EC2 实例拓扑来优化 node-to-node 通信。

SMDDP 库可以在您扩展训练集群时提供性能改进,从而提高训练速度,并具有近乎线性的扩展效率。

注意

SageMaker 分布式训练库可通过训练平台内的 Hugging Face 的 AWS PyTorch 深度学习容器 SageMaker 获得。要使用这些库,你必须使用 SageMaker Python SDK 或通过 Python SDK (Boto3) 或 SageMaker API。 AWS Command Line Interface在整篇文档中,说明和示例都侧重于如何将分布式训练库与 SageMaker Python SDK 配合使用。

SMDDP 集体通信操作针对 AWS 计算资源和网络基础设施进行了优化

SMDDP 库提供了针对 AWS 计算资源AllReduce和网络基础架构进行了优化的AllGather集合操作的实现。

SMDDP 集体行动 AllReduce

SMDDP 库实现了AllReduce操作与向后通道的最佳重叠,从而显著提高了 GPU 利用率。它通过优化 CPU 和 GPU 之间的内核操作,实现了近线性的缩放效率和更快的训练速度。该库AllReduce在 GPU 计算梯度时并行执行,而不会占用额外的 GPU 周期,这使得该库可以实现更快的训练。

  • 利用 CPU:该库使用 CPU 进行AllReduce渐变,从而将此任务从 GPU 中卸载。

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

以下是 SMDDP AllReduce 操作的高级工作流程。

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

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

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

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

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

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

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

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

SMDDP AllReduce 和数据并行架构图

SMDDP 集体行动 AllGather

AllGather是一种集体操作,其中每个 worker 都从一个输入缓冲区开始,然后将所有其他工作程序的输入缓冲区连接或收集到输出缓冲区中。

注意

SMDDP AllGather 集体操作已在 PyTorch v2.0.1 smdistributed-dataparallel>=2.0.1 及更高版本的 Deep Lear AWS ning Containers (DLC) 中提供。

AllGather大量用于分布式训练技术,例如分片数据并行性,其中每个 worker 都拥有模型的一小部分或一个分片层。工作人员在向前和向后传递AllGather之前调用以重建分片层。收集完所有参数后,继续向前和向后传递。在向后传递过程中,每个 worker 还会调用ReduceScatter收集(减少)渐变并将它们分解(分散)成渐变碎片以更新相应的分片层。有关这些集体操作在分片数据并行性中的作用的更多详细信息,请参阅 SMP 库关于分片数据并行性的实现、 DeepSpeed文档中的 ZerO 以及有关完全分片数据并行性的博客。PyTorch

由于在每次迭代中 AllGather 都会调用诸如此类的集体操作,因此它们是 GPU 通信开销的主要来源。更快地计算这些集体运算可以直接转化为更短的训练时间,而不会对收敛产生任何副作用。为了实现这一目标,SMDDP 库提供了AllGather针对 P4d 实例进行了优化的功能。

SMDDP AllGather 使用以下技术来提高 P4d 实例的计算性能。

  1. 它通过具有网状拓扑的弹性结构适配器 (EFA) 网络在实例之间(节点间)传输数据。EFA 是一种 AWS 低延迟、高吞吐量的网络解决方案。节点间网络通信的网状拓扑更适合 EFA 和 AWS 网络基础设施的特性。与涉及多个数据包跳跃的 NCCL 环形拓扑或树状拓扑相比,SMDDP 可以避免从多个跳点累积延迟,因为它只需要一个跳。SMDDP 实现了一种网络速率控制算法,该算法可以平衡网状拓扑中每个通信对等体的工作负载,从而实现更高的全球网络吞吐量。

  2. 它采用基于 NVIDIA GPUDirect RDMA 技术 (GDRCopy) 的低延迟 GPU 内存复制库来协调本地 NVL ink 和 EFA 网络流量。GDRCopy 是 NVIDIA 提供的低延迟 GPU 内存复制库,它提供 CPU 进程和 GPU CUDA 内核之间的低延迟通信。借助这项技术,SMDDP 库能够管理节点内和节点间的数据移动。

  3. 它减少了 GPU 流式多处理器的使用,从而提高了运行模型内核的计算能力。P4d 和 p4de 实例配备了 NVIDIA A100 GPU,每个 GPU 都有 108 个流媒体多处理器。虽然 NCCL 最多需要 24 个流媒体多处理器来运行集体操作,但 SMDDP 使用的流式多处理器少于 9 个。模型计算内核会获取保存的流式多处理器以加快计算速度。