随着基础模型的兴起和成功，许多企业对使用云原生方法进行大模型训练产生了浓厚兴趣。一些 AI 从业者可能认为，实现分布式训练作业的高 GPU 利用率的唯一途径是在高性能计算（HPC）系统（例如互连 Infiniband 的系统）上运行，而不考虑以太网连接的系统。我们展示了最新的分布式训练技术——PyTorch 的全分片数据并行（Fully Sharded Data Parallel，FSDP）——如何利用 IBM Cloud 中的通用以太网联网，成功扩展到参数规模超过 10B（百亿）的模型。
PyTorch FSDP 扩展
随着模型变得越来越大，标准的数据并行训练技术仅在 GPU 能够容纳模型的完整副本及其训练状态（优化器、激活值等）时才有效。然而，GPU 内存的增长速度未能跟上模型规模的扩大，因此出现了训练此类模型的新技术（例如全分片数据并行、DeepSpeed），这使我们能够在训练期间高效地将模型和数据分布在多个 GPU 上。在本篇博文中，我们将演示在将模型规模增加到 11B 时，如何利用 PyTorch 原生 FSDP API 实现模型训练到 64 个节点（512 个 GPU）的显著扩展。
什么是全分片数据并行（FSDP）？
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 论文计算 Teraflops（详情见附录）。我们的实验旨在最大化批量大小（同时避免 cudaMalloc 重试），以充分利用计算和通信的重叠（详见下文）。扩展性定义为 N 个节点与单节点相比，归一化批量大小后的“秒/迭代”之比，这代表了随着节点增加，我们利用额外 GPU 的有效程度。
实验结果
我们使用 T5-3B 配置（混合精度 BF16、激活检查点和 Transformer 封装策略）进行的第一组实验显示，随着 GPU 数量从 8 个增加到 512 个（分别为 1 到 64 个节点），扩展效率达到了 95%。我们在不对现有 FSDP API 进行任何修改的情况下实现了这些结果。我们观察到，在这种规模下，基于以太网的网络带宽足以支持通信与计算的持续重叠。
然而，当我们把 T5 模型规模增加到 11B 时，扩展效率大幅下降至 20%。PyTorch 分析器显示，通信与计算的重叠非常有限。对网络带宽使用情况的进一步调查发现，重叠效果差是由单个数据包通信的延迟引起的，而不是带宽不足（事实上，我们的峰值带宽利用率仅为可用带宽的 1/4）。这使我们推测，如果能通过增加批量大小来延长计算时间，就能更好地重叠通信与计算。但考虑到我们已经处于最大 GPU 内存分配状态，我们必须确定重新平衡内存分配的机会，以允许增加批量大小。我们发现模型状态占用的内存远超所需。这些预留的主要功能是在通信期间准备好预留内存以激进地发送/接收张量，过少的缓冲区会导致等待时间增加，而过多的缓冲区会导致批量大小减小。
为了实现更好的效率，PyTorch 分布式团队引入了一个新的控制旋钮 rate_limiter，它控制为张量发送/接收分配的内存量，从而减轻内存压力并为更大的批量大小留出空间。在我们的案例中，rate_limiter 能将批量大小从 20 增加到 50，从而使计算时间增加了 2.5 倍，并允许通信与计算之间实现更大的重叠。通过此修复，我们将（在 32 个节点时的）扩展效率提高到了 >75%！
对限制扩展效率因素的持续调查发现，rate_limiter 产生了一个重复出现的 GPU 空闲时间流水线气泡。这是因为 rate_limiter 在分配和释放每组内存缓冲区时使用了“阻塞并刷新”（block and flush）方法。由于在启动新的 all_gather 之前必须等待整个块完成，导致 GPU 在每个块的开始时处于空闲状态，等待新的一组 all_gather 参数到达。通过改用滑动窗口方法，消除了这个气泡。在完成单个 all_gather 步骤及其计算（而不是一整块）后，内存被释放，下一个 all_gather 以更均匀的方式立即发出。这一改进消除了流水线气泡，并将（在 32 个节点时的）扩展效率提升至 >90%。
图 1：T5-XL (3B) 和 T5-XXL (11B) 从 1 个节点到 64 个节点的扩展情况
图 2：随节点数量增加，T5-XL(3B) 和 T5-XXL (11B) 的 TFLOPs/sec 使用率
IBM Cloud AI 系统与中间件
用于此项工作的 AI 基础设施是 IBM Cloud 上的大规模 AI 系统，包含近 200 个节点，每个节点配备 8 个 NVIDIA A100 80GB 卡、96 个 vCPU 和 1.2TB CPU 内存。节点内的 GPU 卡通过 NVLink 连接，卡间带宽为 600GBps。节点通过 2 x 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 的显著扩展。我们确定了之前限制 11B 参数模型训练扩展效率不足 20% 的瓶颈。在发现问题后，我们通过新的 rate_limiter 控制进行了修正，确保预留内存与计算时间相关的通信重叠达到最佳平衡。通过这一改进，我们在训练 11B 参数模型时，在 256 个 GPU 上实现了 90% 的扩展效率（提升了 4.5 倍），在 512 个 GPU 上实现了 80% 的效率。此外，3B 参数模型扩展性极佳，即便在 GPU 增加到 512 个时，效率仍高达 95%。
这是业界首次在使用 Kubernetes、通用以太网和 PyTorch 原生 FSDP API 的情况下，为高达 11B 参数的模型实现此类扩展效率。这一改进使用户能够以经济高效且可持续的方式在混合云平台上训练超大规模模型。
我们计划继续研究仅解码器（decoder-only）模型的扩展，并将这些模型的规模增加到 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 FLOPs，其中第一个矩阵大小为 m×k，第二个为 k×n。编码器块由自注意力层和前馈层组成，而解码器块由自注意力层、交叉注意力层和前馈层组成。
注意力（自注意力和交叉注意力）块由 QKV 投影（需要 6Bsh2 次运算）、注意力矩阵计算（需要 2Bs2h 次运算）、针对值的注意力（需要 2Bs2h 次运算）以及后注意力线性投影（需要 2Bsh2 次运算）组成。最后，前馈层需要 15Bsh2 次运算。
一个编码器块的总运算量为 23Bsh2+4Bs2h，而解码器块为 31Bsh2+8Bs2h。总计 24 个编码器和 24 个解码器块，加上 2 次前向传递（因为我们丢弃了激活值）和 1 次反向传递（相当于两次前向传递），最终的 FLOPs 计算得出为 96×(54Bsh2+ 12Bs2h) + 6BshV。此处，B 是每个 GPU 的批量大小，s 是序列长度，h 是隐藏层状态大小，V 是词汇表大小。我们对 T5-XL (3B) 架构也进行了类似的计算，该架构略有不同。