随着Llama2、Falcon和StarCoder等几个公开可用的模型的发布，大型语言模型（LLM）训练在过去一年中变得越来越流行。现在，用户正在训练前所未有规模的LLM，参数范围从10亿到超过1750亿个。训练这些LLM需要大量的计算资源和时间，因为必须使用数百到数千个图形处理单元（GPU）来处理当今庞大的训练数据集和模型大小。分布式训练中的一个瓶颈是由NVIDIA Collective Communication Library (NCCL)处理的GPU通信。在某些大规模分布式训练作业中，与实际的GPU计算相比，花费在GPU通信上的时间更长。为了减轻GPU通信瓶颈并实现更快的训练，Amazon SageMaker 兴奋地宣布优化的AllGather集体操作作为SageMaker分布式数据并行库（SMDDP）的一部分。AllGather 是流行的内存高效数据并行解决方案（如 DeepSpeed Zero Redundancy Optimizer (ZeRO) 和 Fully Sharded Data Parallelism (FSDP)）中使用最频繁的集体操作，也是GPU通信开销的主要贡献因素。在本文中，我们对SMDDP的工作原理进行了高级概述，展示了如何在Amazon SageMaker训练脚本中启用SMDDP，以及可以期待的性能改进。
解决方案概述
传统的数据并行训练涉及在多个GPU上复制整个模型，每个模型使用数据集中的不同片段进行训练。在反向传播过程中，梯度在GPU工作者之间平均，以使得每个模型副本在不同的数据片段上进行训练时更新相同的梯度值。这种技术通过并行化训练数据的消费，大大加快了对庞大数据集的训练速度。然而，如今的一些大型模型（如 Llama2 70B）远远超过了GPU内存的容量，这使得传统的数据并行无法使用。为了在克服有限的GPU内存的同时继续享受数据并行的好处，分片数据并行解决方案如 DeepSpeed ZeRO、PyTorch FSDP 和 Amazon SageMaker模型并行库变得越来越受欢迎。
在分片数据并行中，不再在GPU工作者上复制整个模型，而是将模型参数、梯度和优化器状态分解（即分片）并分布在训练作业中的各个GPU上。为了执行正向和反向传播计算，需要从其他GPU工作者的片段中收集参数，以形成一个或多个模型层。计算完成后，这些层将从内存中释放，以便收集下一组层。请注意，分片数据并行的变体中只对优化器状态和梯度进行分片，而不是模型参数。在这种类型的分片数据并行中仍然使用AllGather，但仅用于在正向传播计算之前收集来自其他GPU工作者的已由不同梯度或优化器状态片段更新的模型参数。有关更多详细信息，请参阅DeepSpeed ZeRO的不同阶段和SHARD_GRAD_OP FSDP分片策略。
在未分片参数时，每次执行AllGather集体操作都会执行—NCCL提供了该例程的标准开源实现。如下所示，每个参与AllGather的GPU工作者从一个输入缓冲区开始，最终得到来自其他工作者的所有输入缓冲区的串联。当AllGather用于分片数据并行时，输入缓冲区包含模型参数分片，而大型输出缓冲区包含从其他分片中实现的一个或多个模型层。
虽然NCCL通常用于分布式训练中的AllGather操作，但其底层实现并未针对亚马逊弹性计算云（Amazon EC2）实例的网络基础设施进行优化，因此其性能可能会减慢端到端训练的速度。SMDDP库是一种为NVIDIA GPU设计的集体通信库，可作为NCCL的替代品，为使用PyTorch进行分布式训练作业提供更好的性能。具体而言，SMDDP为p4d/p4de实例类型提供了AllGather的优化实现。
由于像AllGather这样的集体操作会阻塞正向传播和反向传播的计算，更快的操作执行直接转化为更短的端到端训练时间，并对收敛没有副作用。对于在分片数据并行训练中较少使用的其他集体操作，我们将回退到NCCL。
指南
AWS优化的AllGather
AWS优化的AllGather利用以下技术，在AWS基础架构上实现比NCCL更好的性能：
- 我们通过弹性Fabric适配器（EFA）网络在实例之间移动数据，采用全互联通信模式。EFA是AWS的低延迟高吞吐量网络解决方案，而全互联模式对于节点间网络通信更符合EFA和AWS网络基础设施的特性，相较于NCCL的环形或树状通信模式，所需的数据包跳数较少。
- 使用GDRCopy来协调本地NVLink和EFA网络流量。GDRCopy是一个提供CPU进程和GPU CUDA内核之间低延迟通信的库。借助这项技术，我们能够将节点内和节点间的数据移动进行流水线处理。
- 减少对GPU流处理器的使用，以释放更多计算能力给模型核心。AWS P4d/P4de实例配备了NVIDIA A100 GPU，每个GPU具有108个流处理器。虽然NCCL需要占用24个流处理器来执行集体操作，而SMDDP集体操作只使用最多九个流处理器。通过节省的流处理器，模型核心可更快地执行。
使用方法
SMDDP集体操作通过torch.distributed模块中的进程组抽象与PyTorch进行本地集成。进程组定义了常用集体操作（如AllGather、ReduceScatter、AllReduce等）的接口。用户可以编写通用的分布式代码，然后选择基于所使用计算设备的底层backend来提供这些操作的实现。CPU训练作业通常使用gloo或mpi后端，而NVIDIA GPU使用nccl后端。
SMDDP库通过将自身注册为进程组抽象中的自定义后端来实现。下面的代码片段展示了import语句。在选择基于GPU的分布式训练作业的后端时，只需将nccl替换为smddp。smddp后端与nccl后端遵循相同的语义，并支持相同的训练场景。
DeepSpeed
import smdistributed.dataparallel.torch.torch_smddpdeepspeed.init_distributed(dist_backend="smddp") # 替换为"nccl"FSDP
import smdistributed.dataparallel.torch.torch_smddpdist.init_process_group(backend="smddp") # 替换为"nccl"基准测试
我们对独立的AllGather性能进行了基准测试，其中集合操作在没有任何模型训练的情况下运行。下面是在32个p4d实例上比较NCCL和SMDDP AllGather的样本结果。X轴表示AllGather的输出大小，Y轴表示p4d的400 Gbps EFA网络的网络利用率。4个子图表示共有通信组模式，其中每个p4d实例参与AllGather操作的排名数分别为1、2、4和8。
这些微基准显示出SMDDP在两个关键特性上优于NCCL：
- 在所有配置中，SMDDP的峰值性能（约为90%的带宽利用率）高于NCCL（约为80%的带宽利用率）。
- 与NCCL相比，SMDDP在更小的缓冲区大小上达到了峰值性能。这特别改善了较小模型的训练速度或用户在DeepSpeed中设置较小的AllGather缓冲区大小的情况（其中AllGather大小不需要等于层大小）。
模型训练基准测试
在GPU通信是显著瓶颈的大规模训练作业中，SMDDP可以显著提高训练速度，该速度以模型每秒浮点运算数（TFLOPS/GPU）衡量。
| 配置 | 性能 | ||||
| 模型/训练 | 集群 | 分片数据并行解决方案 | 使用NCCL的模型TFLOPS/GPU | 使用SMDDP的模型TFLOPS/GPU | %加速 |
| 13B Llama2 序列长度：4096 全局批次大小：4M标记 | 64个p4d.24xlarge节点（512个NVIDIA A100 GPU） | PyTorch FSDP | 97.89 | 121.85 | 24.40% |
| 65B GPT-NeoX 序列长度：2048 全局批次大小：4M标记 | 64个p4d.24xlarge节点（512个NVIDIA A100 GPU） | DeepSpeed ZeRO 第3阶段* | 99.23 | 108.66 | 9.50% |
*使用了EleutherAI的Megatron-DeepSpeed存储库。还启用了张量并行度为8的张量并行。
注意：模型TFLOPS/GPU基于此处定义的模型FLOPS利用率计算，并且其他地方的基准测试数据可能引用硬件TFLOPS/GPU作为性能指标。硬件TFLOPS/GPU可以近似为4/3乘以模型TFLOPS/GPU。
结论
在本文中，我们向您展示了如何通过只更改两行代码，显着加快在Amazon SageMaker上的分片数据并行训练作业。大规模分布式训练在LLM的出现下越来越普遍，但随着规模的扩大，成本也越来越高。通过减少GPU之间的通信瓶颈，SMDDP帮助您在大规模上更快地进行训练，并节省计算资源。您可以在Amazon SageMaker Examples GitHub存储库中找到更多使用分片数据并行训练的SMDDP示例。