随着基础模型的出现和成功，使用云原生方法进行大型模型训练越来越受到许多企业的关注。一些 AI 从业者可能认为，他们实现高 GPU 利用率进行分布式训练作业的唯一方式是在 HPC 系统上运行，例如那些通过 Infiniband 互联的系统，而可能不会考虑以太网连接的系统。我们展示了最新的分布式训练技术，PyTorch 的完全分片数据并行（FSDP），如何成功扩展到 10B+参数规模的模型，并使用 IBM Cloud 中的商用以太网网络。
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 研讨会中的代码。在我们的每个实验中，我们从一个节点实验开始，以创建基线，并报告批大小归一化的秒/迭代指标，以及根据 Megatron-LM 论文计算 teraflops（见附录中 T5 的 teraflop 计算细节）。我们的实验旨在最大化批大小（同时避免 cudaMalloc 重试），以充分利用以下讨论中的计算和通信重叠。扩展定义为 N 个节点与单个节点批大小归一化的秒/迭代比，表示随着节点数量的增加，我们可以多好地利用额外的 GPU。
实验结果
我们的首次实验使用 T5-3B 配置（混合精度 BF16、激活检查点和 transformer 包装策略）展示了在将 GPU 数量从 8 个增加到 512 个（1 到 64 个节点）时，达到了 95%的扩展效率。我们实现了这些结果，而没有对现有的 FSDP API 进行任何修改。我们观察到，对于这个规模，基于以太网的网络上，有足够的带宽来支持通信和计算的连续重叠。
然而，当我们把 T5 模型的大小增加到 11B 时，缩放效率大幅下降到 20%。PyTorch 分析器显示，通信和计算的重叠非常有限。进一步调查网络带宽使用情况发现，重叠不佳是由单个数据包通信中的延迟造成的，而不是带宽需求（实际上，我们的峰值带宽利用率是可用带宽的四分之一）。这导致我们假设，如果我们通过增加批量大小来增加计算时间，我们可以更好地重叠通信和计算。然而，鉴于我们已经在最大 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 云 AI 系统和中间件
该工作的 AI 基础设施是 IBM 云上的一套大规模 AI 系统，由近 200 个节点组成，每个节点配备 8 张 NVIDIA A100 80GB 显卡、96 个 vCPU 和 1.2TB CPU RAM。节点内的 GPU 卡通过 NVLink 连接，卡间带宽为 600GBps。节点之间通过基于 SRIOV 的 TCP/IP 堆栈的 2 x 100Gbps 以太网链路连接，提供 120Gbps 的可使用带宽。
自 2022 年 5 月起，IBM 云 AI 系统已具备生产就绪性，并配置了 OpenShift 容器平台以运行 AI 工作负载。我们还构建了一个用于生产 AI 工作负载的软件堆栈，该堆栈提供端到端工具以训练工作负载。中间件利用 Ray 进行预和后处理工作负载，并使用 PyTorch 进行模型训练。我们还集成了 Kubernetes 原生调度器 MCAD，该调度器管理多个作业，具有作业排队、批量调度、优先级和配额管理功能。多 NIC CNI 发现所有可用的网络接口，并将它们作为单个 NIC 池处理，从而优化了 Kubernetes 中网络接口的使用。最后，CodeFlare CLI 支持使用桌面 CLI（例如，GPU 利用率、应用指标如损失、梯度范数）的单窗口全栈可观察性。
图 3：基础模型中间件堆栈
结论与未来工作
总结来说，我们展示了如何在非 InfiniBand 网络上实现 FSDP API 的显著扩展。我们确定了瓶颈，该瓶颈限制了扩展效率，使得 11B 参数模型训练的效率低于 20%。在确定问题后，我们通过新的速率限制器控制来纠正这个问题，以确保相对于计算时间，预留内存和通信重叠的更优平衡。通过这次改进，我们能够在 256 个 GPU 上实现 90%的扩展效率（提高了 4.5 倍），在 512 个 GPU 上实现 80%的扩展效率，用于训练 11B 参数模型。此外，当我们将 GPU 数量增加到 512 时，3B 参数模型的扩展效率仍然极高，达到 95%。
这是行业首次使用 Kubernetes、vanilla Ethernet 和 PyTorch 原生 FSDP API 实现高达 11B 参数模型的扩展效率。这次改进使用户能够在混合云平台上以成本效益和可持续的方式训练大型模型。
我们计划继续研究仅使用解码器模型的扩展，并将这些模型的大小增加到 100B+参数。从系统设计角度来看，我们正在探索 RoCE 和 GDR 等能力，这些能力可以提高以太网网络通信的延迟。
致谢
本博客得以实现，得益于 PyTorch 分布式和 IBM 研究团队的贡献。
我们谨代表 PyTorch 分布式团队感谢 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。
附录
特拉浮点计算
T5-XXL（11B）架构有两种类型的 T5 模块，一种是编码器，另一种是解码器。遵循 Megatron-LM 的方法，其中每个矩阵乘法需要 2m×k×n 次浮点运算（FLOPs），第一个矩阵的大小为 m×k，第二个矩阵的大小为 k×n。编码器模块由自注意力层和前馈层组成，而解码器模块由自注意力层、交叉注意力层和前馈层组成。
注意力（包括自注意力和交叉注意力）模块由 QKV 投影组成，需要 6Bsh 2 次操作，一个注意力矩阵计算需要 2Bs 2 h 次操作，一个对值的注意力需要 2Bs 2 h 次计算，而注意力之后的线性投影需要 2Bsh 2 次操作。最后，前馈层需要 15Bsh 2 次操作。
一个编码器块的总量为 23Bsh 2 +4Bs 2 h，而一个解码器块的总量为 31Bsh 2 +8Bs 2 h。总共有 24 个编码器和 24 个解码器块，以及 2 次正向传递（因为我们丢弃了激活值）和 1 次反向传递（相当于两次正向传递），最终的 FLOPs 计算结果为 96×(54Bsh 2 + 12Bs 2 h) + 6BshV。在这里，B 是每个 GPU 的批次大小，s 是序列长度，h 是隐藏状态大小，V 是词汇表大小。我们对 T5-XL（3B）架构进行类似的计算，略有不同。
