鉴于基础模型的出现和成功，许多企业对使用云原生方法进行大型模型训练越来越感兴趣。一些 AI 从业者可能认为，分布式训练任务要实现高 GPU 利用率，唯一的途径是在高性能计算（HPC）系统上运行，例如那些通过 Infiniband 互连的系统，并且可能不考虑使用以太网连接的系统。我们演示了 PyTorch 最新的分布式训练技术，Fully Sharded Data Parallel (FSDP)，如何在 IBM Cloud 中使用商用以太网网络成功地扩展到参数量超过 100 亿的模型。
PyTorch FSDP 扩展
随着模型变大，标准的数据并行训练技术只有在 GPU 能容纳模型的完整副本及其训练状态（优化器、激活等）时才能工作。然而，GPU 内存的增长未能跟上模型规模的增长，因此出现了训练此类模型的新技术（例如 Fully Sharded Data Parallel、DeepSpeed），这些技术使我们能够在训练期间有效地将模型和数据分布到多个 GPU 上。在这篇博文中，我们展示了如何在使用 PyTorch 原生 FSDP API 将模型大小增加到 110 亿时，将模型训练显著扩展到 64 个节点（512 个 GPU）。
什么是 Fully Sharded Data Parallel？
FSDP 通过使用 wrapping policy 将模型参数、梯度和优化器状态分片到 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 重试），以充分利用计算和通信的重叠，如下文所述。扩展（Scaling）定义为 N 个节点按批量大小归一化的每迭代秒数与单节点该值的比率，表示随着添加更多节点，我们对额外 GPU 的利用效率。
实验结果
使用 T5-3B 配置（BF16 混合精度、激活检查点和 Transformer wrapping policy）的第一组实验表明，随着我们将 GPU 数量从 8 个增加到 512 个（分别对应 1 到 64 个节点），扩展效率达到 95%。我们在未修改现有 FSDP API 的情况下取得了这些结果。我们观察到，对于这种规模，在基于以太网的网络上，有足够的带宽来实现通信和计算的持续重叠。
然而，当我们将 T5 模型大小增加到 110 亿时，扩展效率大幅下降到 20%。PyTorch profiler 显示通信和计算的重叠非常有限。对网络带宽使用的进一步调查显示，重叠不足是由单个数据包通信中的延迟引起的，而不是所需的带宽问题（实际上，我们的峰值带宽利用率只有可用带宽的四分之一）。这使得我们假设，如果通过增加批量大小来增加计算时间，我们可以更好地重叠通信和计算。然而，考虑到我们已经达到 GPU 内存的最大分配量，我们必须找到重新平衡内存分配的机会，以便增加批量大小。我们发现模型状态分配了比实际需要多得多的内存。这些预留内存的主要功能是准备好在通信期间积极地发送/接收张量，而过少的缓冲区可能导致等待时间增加，过多的缓冲区则导致批量大小减小。
为了获得更好的效率，PyTorch 分布式团队引入了一个新的控制旋钮 rate_limiter，它控制为发送/接收张量分配多少内存，从而减轻内存压力并为更大的批量大小提供空间。在我们的案例中，rate_limiter 可以将批量大小从 20 增加到 50，从而将计算时间增加了 2.5 倍，并使得通信和计算的重叠程度大大提高。通过这个修复，我们将扩展效率提高到了 >75%（在 32 个节点上）！
对限制扩展效率的因素的持续调查发现，rate limiter 正在创建一个循环出现的 GPU 空闲时间流水线气泡。这是因为 rate limiter 对每组内存缓冲区的分配和释放采用了块和刷新（block and flush）方法。通过等待整个块完成才启动新的 all_gather，GPU 在每个块的开始时处于空闲状态，等待新的一组 all_gather 参数到来。通过转向滑动窗口（sliding window）方法，消除了这个气泡。在完成单个 all_gather 步骤及其计算后（而不是完成一整块），内存被释放，下一个 all_gather 会以更均匀的方式立即发出。这一改进消除了流水线气泡，并将扩展效率提高到了 >90%（在 32 个节点上）。
图 1：T5-XL (30 亿) 和 T5-XXL (110 亿) 模型从 1 个节点扩展到 64 个节点
图 2：随着节点数量增加，T5-XL (30 亿) 和 T5-XXL (110 亿) 的 TFLOPs/秒 使用情况
IBM Cloud AI 系统与中间件
本文所使用的 AI 基础设施是 IBM Cloud 上的大规模 AI 系统，包含近 200 个节点，每个节点配备 8 块 NVIDIA A100 80GB 显卡、96 个虚拟 CPU 和 1.2TB CPU 内存。节点内的 GPU 显卡通过 NVLink 连接，卡间带宽为 600GBps。节点之间通过 2 条 100Gbps 以太网链路连接，使用基于 SRIOV 的 TCP/IP 协议栈，提供 120Gbps 的可用带宽。
IBM Cloud AI 系统自 2022 年 5 月起已投入生产，并配置了 OpenShift 容器平台以运行 AI 工作负载。我们还为生产 AI 工作负载构建了一个软件栈，提供端到端的训练工作负载工具。中间件利用 Ray 进行预处理和后处理工作负载，并利用 PyTorch 进行模型训练。我们还集成了 Kubernetes 原生调度器 MCAD，它通过作业队列、gang scheduling、优先级和配额管理来管理多个作业。多网卡 CNI 发现所有可用的网络接口，并将它们作为一个单一的网卡池处理，从而在 Kubernetes 中实现网络接口的优化使用。最后，CodeFlare CLI 支持使用桌面 CLI 对整个栈进行统一的可观测性监控（例如，GPU 利用率、损失、梯度范数等应用指标）。
图 3：基础模型中间件栈
结论与未来工作
总之，我们展示了如何在非 InfiniBand 网络上实现 FSDP API 的出色扩展。我们发现了限制 110 亿参数模型训练扩展效率低于 20% 的瓶颈。在确定问题后，我们通过一个新的 rate limiter 控制解决了这个问题，确保了预留内存和通信重叠相对于计算时间的更优平衡。通过这一改进，我们在 256 个 GPU 上将 110 亿参数模型的训练扩展效率提高到 90%（提升了 4.5 倍），在 512 个 GPU 上达到 80%。此外，即使我们将 GPU 数量增加到 512 个，30 亿参数模型也表现出极好的扩展性，效率达到 95%。
使用 Kubernetes、普通以太网和 PyTorch 原生 FSDP API，实现高达 110 亿参数模型如此高的扩展效率，这在业界尚属首次。这一改进使用户能够在混合云平台上以经济高效且可持续的方式训练大型模型。
我们计划继续研究仅解码器模型的扩展，并将这些模型的大小增加到 1000 亿以上的参数。从系统设计的角度来看，我们正在探索 RoCE 和 GDR 等能力，这些能力可以改善以太网网络上的通信延迟。
致谢
本博文的撰写得益于 PyTorch Distributed 团队和 IBM 研究团队的贡献。
来自 PyTorch Distributed 团队，我们要感谢 Less Wright, Hamid Shojanazeri, Geeta Chauhan, Shen Li, Rohan Varma, Yanli Zhao, Andrew Gu, Anjali Sridhar, Chien-Chin Huang 和 Bernard Nguyen。
来自 IBM 研究团队，我们要感谢 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 (110 亿) 架构有两种类型的 T5 块，一种是编码器（encoder），另一种是解码器（decoder）。遵循 Megatron-LM 的方法，每次矩阵乘法需要 2m×k×n FLOPs，其中第一个矩阵大小为 m×k，第二个为 k×n。编码器块包含自注意力（self-attention）层和前馈层（feed forward layers），而解码器块包含自注意力层、交叉注意力（cross-attention）层和前馈层。
注意力（包括自注意力和交叉注意力）块包含 QKV 投影，需要 6Bsh2 次操作；注意力矩阵计算需要 2Bs2h 次操作；对值进行注意力计算需要 2Bs2h 次计算；以及注意力后的线性投影需要 2Bsh2 次操作。最后，前馈层需要 15Bsh2 次操作。
一个编码器块的总操作次数是 23Bsh2+4Bs2h，而一个解码器块的总操作次数是 31Bsh2+8Bs2h。总共有 24 个编码器块和 24 个解码器块，以及 2 次前向传播（因为我们丢弃了激活）和 1 次反向传播（相当于 2 次前向传播），最终的 FLOPs 计算结果为 96×(54Bsh2+ 12Bs2h) + 6BshV。其中，B 是每个 GPU 的批量大小，s 是序列长度，h 是隐藏状态大小，V 是词汇表大小。我们对略有不同的 T5-XL (30 亿) 架构重复了类似的计算。
