鉴于基础模型的出现和成功，使用云原生方法进行大型模型训练正引起许多企业的日益关注。一些 AI 从业者可能认为，他们实现分布式训练作业高 GPU 利用率的唯一方法是在 HPC 系统上运行它们，例如那些通过 Infiniband 互连的系统，并且可能不考虑以太网连接的系统。我们演示了最新的分布式训练技术，PyTorch 的完全分片数据并行 (FSDP)，如何在 IBM Cloud 中使用商用以太网网络成功扩展到 10B+ 参数的模型。
PyTorch FSDP 扩展
随着模型变得越来越大，数据并行训练的标准技术仅在 GPU 可以容纳模型的完整副本及其训练状态（优化器、激活等）时才有效。但是，GPU 内存的增加速度并未跟上模型尺寸的增加速度，并且出现了训练此类模型的新技术（例如，完全分片数据并行、DeepSpeed），这使我们能够在训练期间有效地将模型和数据分布在多个 GPU 上。在这篇博文中，我们将演示一种通过使用 PyTorch 原生 FSDP API 将模型训练显着扩展到 64 个节点（512 个 GPU），并将模型大小增加到 11B 的路径。
什么是完全分片数据并行？
FSDP 通过将模型参数、梯度和优化器状态分片为 K 个 FSDP 单元来扩展分布式数据并行训练 (DDP) 方法，这通过使用包装策略来确定。FSDP 通过显着减少每个 GPU 上的内存占用并重叠计算和通信，在资源和性能方面实现了大型模型训练效率。
资源效率通过减少内存占用来实现，方法是让所有 GPU 拥有每个 FSDP 单元的一部分。为了处理给定的 FSDP 单元，所有 GPU 通过 all_gather 通信调用共享其本地拥有的部分。
性能效率是通过将即将到来的 FSDP 单元的 all_gather 通信调用与当前 FSDP 单元的计算重叠来实现的。一旦当前 FSDP 单元被处理，非本地拥有的参数将被丢弃，从而为即将到来的 FSDP 单元释放内存。此过程通过计算和通信的重叠实现训练效率，同时还减少了每个 GPU 所需的峰值内存。
在接下来的内容中，我们将演示 FSDP 如何使我们能够在整个分布式训练作业中保持数百个 GPU 的高利用率，同时通过标准以太网网络运行（系统描述在本博客末尾）。我们为实验选择了 T5 架构，并利用了 FSDP workshop 中的代码。在我们的每个实验中，我们都从单节点实验开始创建基线，并报告以秒/迭代为单位的指标，并按批大小进行归一化，以及根据 Megatron-LM 论文计算 teraflops（有关 T5 的 teraflop 计算的详细信息，请参见附录）。我们的实验旨在最大化批大小（同时避免 cudaMalloc 重试），以便充分利用计算和通信中的重叠，如下所述。扩展定义为 N 个节点的秒/迭代与单节点的秒/迭代之比，以批大小归一化，表示当我们添加更多节点时，我们如何很好地利用额外的 GPU。
实验结果
我们使用 T5-3B 配置（BF16 混合精度、激活检查点和 transformer 包装策略）的第一组实验表明，当我们将 GPU 数量从 8 个增加到 512 个（分别为 1 个节点到 64 个节点）时，扩展效率为 95%。我们在不修改现有 FSDP API 的情况下实现了这些结果。我们观察到，对于此规模，基于以太网的网络具有足够的带宽来支持通信和计算的持续重叠。
然而，当我们将 T5 模型大小增加到 11B 时，扩展效率大幅下降至 20%。PyTorch profiler 显示，通信和计算的重叠非常有限。对网络带宽使用情况的进一步调查显示，重叠不佳是由单个数据包的通信延迟引起的，而不是所需的带宽（事实上，我们的峰值带宽利用率仅为可用带宽的 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，该调度器管理具有作业排队、组调度、优先级排序和配额管理的多个作业。多 NIC CNI 发现所有可用的网络接口，并将它们处理为单个 NIC 池，从而在 Kubernetes 中优化网络接口的使用。最后，CodeFlare CLI 支持使用桌面 CLI（例如，GPU 利用率、应用程序指标（如损失、梯度范数））对整个堆栈进行单窗格可观察性。
图 3：基础模型中间件堆栈
结论和未来工作
总之，我们演示了如何在非 InfiniBand 网络上实现 FSDP API 的显着扩展。我们确定了将 11B 参数模型训练的扩展效率限制在 20% 以下的瓶颈。在确定问题后，我们能够通过新的速率限制器控制来纠正此问题，以确保相对于计算时间，预留内存和通信重叠之间更加平衡。通过此改进，我们能够将 11B 参数模型的训练扩展效率提高到 90%（提高了 4.5 倍），在 256 个 GPU 时达到 90%，在 512 个 GPU 时达到 80%。此外，即使我们将 GPU 数量增加到 512 个，3B 参数模型的扩展性也非常好，效率达到 95%。
这是业界首次在使用 Kubernetes、普通以太网和 PyTorch 原生 FSDP API 的情况下，为高达 11B 参数的模型实现如此高的扩展效率。此改进使使用者能够在混合云平台上以经济高效且可持续的方式训练大型模型。
我们计划继续研究仅解码器模型的扩展，并将这些模型的尺寸增加到 100B+ 参数。从系统设计的角度来看，我们正在探索 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 FLOP，其中第一个矩阵的大小为 m×k，第二个矩阵的大小为 k×n。编码器块由自注意力层和前馈层组成，而解码器块由自注意力层、交叉注意力层和前馈层组成。
注意力（自注意力和交叉注意力）块由 QKV 投影组成，这需要 6Bsh2 次运算，注意力矩阵计算需要 2Bs2h 次运算，值注意力需要 2Bs2h 次计算，后注意力线性投影需要 2Bsh2 次运算。最后，前馈层需要 15Bsh2 次运算。
编码器块的总计为 23Bsh2+4Bs2h，而解码器块的总计为 31Bsh2+8Bs2h。总共有 24 个编码器块和 24 个解码器块，以及 2 个前向传递（因为我们丢弃了激活）和一个后向传递（相当于两个前向传递），最终的 FLOP 计算结果为 96×(54Bsh2+ 12Bs2h) + 6BshV。其中，B 是每个 GPU 的批大小，s 是序列长度，h 是隐藏状态大小，V 是词汇表大小。我们对 T5-XL (3B) 架构重复类似的计算，后者略有不同。
