鉴于基础模型的出现和成功，使用云原生方法进行大型模型训练正日益受到许多企业的关注。一些AI从业者可能认为，要为分布式训练作业实现高GPU利用率，唯一的方法是在HPC系统上运行它们，例如那些通过InfiniBand互连的系统，而可能不考虑以太网连接的系统。我们展示了PyTorch最新的分布式训练技术——完全分片数据并行（FSDP）如何成功地在IBM Cloud上使用商用以太网网络扩展到10亿以上参数的模型。
PyTorch FSDP扩展
随着模型变大，标准的数据并行训练技术只有在GPU能够容纳模型的完整副本及其训练状态（优化器、激活等）时才有效。然而，GPU内存的增长并未跟上模型大小的增长，因此出现了训练此类模型的新技术（例如完全分片数据并行、DeepSpeed），这些技术允许我们在训练期间有效地将模型和数据分布到多个GPU上。在这篇博客文章中，我们展示了如何在使用PyTorch原生FSDP API将模型大小增加到110亿参数时，将模型训练扩展到64个节点（512个GPU），从而实现显著的扩展。
什么是完全分片数据并行？
FSDP通过将模型参数、梯度和优化器状态分片为K个FSDP单元来扩展分布式数据并行训练（DDP）方法，这些单元由包装策略确定。FSDP通过显著减少每个GPU的内存占用以及重叠计算和通信来在资源和性能方面实现大型模型训练效率。
通过让所有GPU拥有每个FSDP单元的一部分来实现内存占用减少，从而实现资源效率。为了处理给定的FSDP单元，所有GPU通过all_gather通信调用共享其本地拥有的部分。
性能效率是通过将即将到来的FSDP单元的all_gather通信调用与当前FSDP单元的计算重叠来实现的。一旦当前FSDP单元处理完毕，非本地拥有的参数将被丢弃，从而为即将到来的FSDP单元释放内存。此过程通过计算和通信的重叠实现训练效率，同时还减少了每个GPU所需的峰值内存。
在下文中，我们将展示FSDP如何使我们能够在分布式训练作业中保持数百个GPU的高度利用率，同时在标准以太网网络上运行（系统描述在博客末尾）。我们选择T5架构进行实验，并利用FSDP研讨会的代码。在我们的每个实验中，我们都从单个节点实验开始创建基线，并报告以批处理大小归一化的每迭代秒数指标，并根据Megatron-LM论文计算每秒浮点运算数（有关T5的每秒浮点运算数计算的详细信息，请参阅附录）。我们的实验旨在最大化批处理大小（同时避免cudaMalloc重试），以充分利用计算和通信中的重叠，如下所述。扩展被定义为N个节点的每迭代秒数（按批处理大小归一化）与单个节点之比，表示随着更多节点的添加，我们能多好地利用额外的GPU。
实验结果
我们使用T5-3B配置（BF16混合精度、激活检查点和transformer封装策略）进行的第一组实验表明，随着我们将GPU数量从8增加到512（分别对应1到64个节点），扩展效率达到95%。我们无需修改现有FSDP API就实现了这些结果。我们观察到，在这种规模下，基于以太网的网络具有足够的带宽，可以实现通信和计算的持续重叠。
然而，当我们将T5模型大小增加到110亿时，扩展效率大幅下降到20%。PyTorch性能分析器显示，通信和计算的重叠非常有限。对网络带宽使用情况的进一步调查显示，重叠不佳是由单个数据包通信中的延迟引起的，而不是所需的带宽（事实上，我们的峰值带宽利用率仅为可用带宽的1/4）。这使我们假设，如果我们可以通过增加批处理大小来增加计算时间，我们就可以更好地重叠通信和计算。然而，鉴于我们已经达到最大GPU内存分配，我们必须找出重新平衡内存分配的机会，以允许增加批处理大小。我们发现模型状态分配的内存比实际需要的多得多。这些预留的主要功能是预留内存以在通信期间积极发送/接收张量，而过少的缓冲区可能导致等待时间增加，而过多的缓冲区则导致批处理大小更小。
为了实现更好的效率，PyTorch分布式团队引入了一个新的控制旋钮，即rate_limiter，它控制为张量发送/接收分配多少内存，从而减轻内存压力并为更高的批处理大小提供空间。在我们的例子中，rate_limiter可以将批处理大小从20增加到50，从而将计算时间增加2.5倍，并允许更大程度的通信和计算重叠。通过此修复，我们将扩展效率提高到>75%（在32个节点上）！
对限制扩展效率的因素的持续调查发现，速率限制器正在创建GPU空闲时间的循环管道气泡。这是由于速率限制器对每组内存缓冲区的分配和释放采用了阻塞和刷新方法。通过等待整个块完成，然后才启动新的all_gather，GPU在每个块的开始时都会空闲，同时等待新的all_gather参数集到达。通过改为滑动窗口方法缓解了此气泡。一旦单个all_gather步骤及其计算完成（而不是一整块），内存就会被释放，并且下一个all_gather会以更均匀的方式立即发出。这种改进消除了管道气泡，并将扩展效率提高到>90%（在32个节点上）。
图1：T5-XL（3B）和T5-XXL（11B）从1个节点到64个节点的扩展
图2：T5-XL（3B）和T5-XXL（11B）随节点数量增加的TFLOPs/秒使用情况
IBM Cloud AI系统和中间件
用于这项工作的AI基础设施是IBM Cloud上的一个大规模AI系统，由近200个节点组成，每个节点配备8块NVIDIA A100 80GB显卡、96个vCPU和1.2TB CPU RAM。节点内的GPU卡通过NVLink连接，卡到卡带宽为600GBps。节点通过2个100Gbps以太网链路连接，采用基于SRIOV的TCP/IP堆栈，提供120Gbps的可用带宽。
IBM Cloud AI系统自2022年5月起已投入生产，并配置了OpenShift容器平台以运行AI工作负载。我们还为生产AI工作负载构建了一个软件堆栈，提供端到端的训练工作负载工具。该中间件利用Ray进行预处理和后处理工作负载，并利用PyTorch进行模型训练。我们还集成了Kubernetes原生调度器MCAD，它通过作业排队、组调度、优先级和配额管理来管理多个作业。多网卡CNI发现所有可用的网络接口，并将它们作为单个网卡池处理，从而优化Kubernetes中网络接口的使用。最后，CodeFlare CLI支持通过桌面CLI（例如，GPU利用率、损失等应用程序指标、梯度范数）对整个堆栈进行单一观察。
图3：基础模型中间件堆栈
结论与未来工作
总而言之，我们展示了如何在非InfiniBand网络上实现FSDP API的显著扩展。我们确定了导致110亿参数模型训练扩展效率低于20%的瓶颈。在确定问题后，我们能够通过新的速率限制器控制来纠正此问题，以确保预留内存和通信重叠相对于计算时间达到更优化的平衡。通过此改进，我们能够将110亿参数模型的训练扩展效率提高到90%（提升4.5倍，在256个GPU上）和80%（在512个GPU上）。此外，30亿参数模型即使在我们将GPU数量增加到512个时，也能以95%的效率实现极好的扩展。
这是业界首次使用Kubernetes、通用以太网和PyTorch原生FSDP API，将参数高达110亿的模型训练扩展到如此高的效率。这项改进使用户能够以经济高效和可持续的方式在混合云平台上训练大型模型。
我们计划继续研究仅解码器模型的扩展，并将这些模型的大小增加到1000亿以上参数。从系统设计角度来看，我们正在探索RoCE和GDR等功能，这些功能可以改善以太网网络上的通信延迟。
致谢
这篇博客的完成得益于PyTorch Distributed和IBM Research团队的贡献。
我们感谢PyTorch Distributed团队的Less Wright、Hamid Shojanazeri、Geeta Chauhan、Shen Li、Rohan Varma、Yanli Zhao、Andrew Gu、Anjali Sridhar、Chien-Chin Huang和Bernard Nguyen。
我们感谢IBM Research团队的Linsong Chu、Sophia Wen、Lixiang (Eric) Luo、Marquita Ellis、Davis Wertheimer、Supriyo Chakraborty、Raghu Ganti、Mudhakar Srivatsa、Seetharami Seelam、Carlos Costa、Abhishek Malvankar、Diana Arroyo、Alaa Youssef和Nick Mitchell。
附录
每秒万亿次浮点运算（Teraflop）计算
T5-XXL（11B）架构有两种类型的T5块，一种是编码器，第二种是解码器。遵循Megatron-LM的方法，其中每个矩阵乘法需要2m×k×n次浮点运算，其中第一个矩阵的大小为m×k，第二个为k×n。编码器块由自注意力层和前馈层组成，而解码器块由自注意力层、交叉注意力层和前馈层组成。
注意力（自注意力和交叉注意力）块包括一个QKV投影，需要6Bsh2次操作，一个注意力矩阵计算需要2Bs2h次操作，一个对值的注意力需要2Bs2h次计算，以及注意力后线性投影需要2Bsh2次操作。最后，前馈层需要15Bsh2次操作。
一个编码器块的总浮点运算为23Bsh2+4Bs2h，而一个解码器块的总浮点运算为31Bsh2+8Bs2h。总共有24个编码器块和24个解码器块，以及2次前向传播（因为我们丢弃了激活）和1次后向传播（相当于2次前向传播），最终的浮点运算量为96×(54Bsh2+ 12Bs2h) + 6BshV。其中，B是每个GPU的批大小，s是序列长度，h是隐藏状态大小，V是词汇量大小。我们对T5-XL（3B）架构重复类似的计算，该架构略有不同。