几十亿参数规模的大模型网络架构优化

作者系美团资深软件开发专家

原文链接 https://zhuanlan.zhihu.com/p/680573811

咨询网络相关问题可访问链接：

https://www.zhihu.com/consult/people/1403017576021184512

摘要

本文打破了在训练大型语言模型（LLM）构建any-to-any网络的传统的方案。我们了解，LLM 表现出一种独特的通信模式，其中只有一小部分 GPU 需要在其中进行高带宽any-to-any通信，用来实现接近最佳的训练性能。对于GPU集群而言，通信是微不足道的、稀疏的且相似的。我们提出了一种LLMs的通信要求非常相似的新网络架构。我们的架构将集群划分为非阻塞任意高带宽互连的 GPU 组,我们称之为 HB 域的互连。在 HB 域中，网络仅连接具有通信需求的 GPU。我们将该网络称为“rail-only”连接，并表明，与最先进的any-to-any Clos 网络相比，我们提出的架构可将网络成本降低高达 75%，且不会影响LLM训练性能。

1 、介绍

不断发展的大型语言模型 (LLM) 领域有望彻底改变我们对人类语言处理的理解，推动人工智能 (AI) 的技术进步。OpenAI 的 ChatGPT 在发布后三个月内为超过 1 亿活跃用户提供服务，使其成为有史以来增长最快的应用程序。除了聊天机器人之外 ，LLM正在逐渐渗透到我们的数字生活中。主要服务提供商将这些强大的模型集成到智能驾驶程序和搜索引擎中 ，改变人类与数字领域互动的方式。

LLMs是最大、计算量最大的深度神经网络 (DNNs) 之一。最新的 GPT4 模型估计有数万亿个参数，需要数月的时间来训练 。从历史上看，研究人员寻求通过优化并行化策略，复杂的调度，高级压缩来提高分布式DNN训练和推理的性能，甚至网络拓扑本身的重新配置。尽管做出了这些努力，LLMs仍然需要大量的原始计算能力。

2020 年的 GPT3 模型在 Nvidia 的 V100 GPU 上已经需要 355 个 GPU-年。随着摩尔定律放缓，LLM规模和计算需求的增长速度超过了加速器的进步，使得超大规模GPU集群不可避免。我们与业界领先的机器学习架构师的对话表明，下一代LLMs可能需要超过 30,000 个 GPU 的计算能力才能在合理的时间内完成训练。

以 GPU 为中心的集群通常采用两种类型的连接 。对于第一个，一些 GPU（例如，DGX H100 服务器有 8 个）通过 NVLink 等短程通信协议驻留在高带宽域（称为 HB 域）内。第二个连接形成一个网络，能够使用支持 RDMA 的 NIC 进行任意 GPU 通信，并以 Clos 网络的变体进行连接。该集群在该网络上使用 RDMA 协议，通过 GPU-Direct 完全绕过 CPU 和操作系统，从而受益匪浅 。然而，将 RDMA 网络扩展到数万个 GPU 具有挑战性。先前的工作表明，大规模 RDMA 网络容易出现死锁和 PFC 风暴 ，从而降低性能。此外，随着规模的扩大，Clos 架构的成本变得过高。数据中心提供商经常采用超额认购来降低集群成本，从而加剧死锁问题。

先前的工作提出了几种技术来支持大规模 RDMA 网络并降低其成本 。这些方法从根本上取决于网络能够进行any-to-any通信的假设。这一假设构成了数据中心几十年来概念化和开发的基石。

在本文中，我们挑战了这一假设，并表明 LLM 训练流量不需要网络中所有 GPU 之间的任意连接。我们认为，采用最佳并行化策略，LLM 训练工作负载仅在 GPU 的小子集中需要高带宽任意对任意连接，并且每个子集都适合 HB 域。在 HB 域中，通信仅发生在少数 GPU 对上，并且流量微不足道。因此，构建数据中心互连的传统任意对任意方法为分布式LLMs训练增加了不必要的复杂性和成本。

我们提出了一种能够准确反映 LLM 通信需求的网络架构。在此架构中，集群首先被划分为多个 HB 域，每个 HB 域均通过全平分带宽any-to-any互连。在 HB 域中，网络仅通过网络流量连接 GPU 组，而不是形成 Clos 来支持任意对任意通信。我们证明，我们以 LLM 为中心的网络架构具有与全平分带宽any-to-any的 Clos 集群相同的性能，同时将成本降低了 37% 至 75%。

2、动机

在本节中，我们首先介绍传统的以 GPU 为中心的集群的架构。然后，我们对 LLM 流量模式进行彻底分析，以激发网络架构与传统设计的对比。

2.1 最先进的 GPU 集群设计

传统的网络集群旨在使用多层 Clos 网络来服务 CPU 密集型工作负载，如图 1 所示。这种架构被称为 Fat-Tree 网络，在系统和网络社区中得到了深入研究。在典型的基于 Fat-Tree 的集群中，每台服务器都配备一个 NIC（40 Gbps 至 400 Gbps），并且 K 个服务器排列在连接到架顶式 (ToR) 交换机的机架中。然后，ToR 交换机连接到聚合交换机，以提供跨机架的连接，形成一个 Pod。最后，Pod 与主干交换机互连，允许 CPU 集群中的服务器之间进行any-to-any通信。传统的网络集群旨在使用多层 Clos 网络来服务 CPU 密集型工作负载，如图所示如图 1 [32, 33] 所示。这种架构被称为 Fat-Tree 网络，在系统和网络社区中得到了深入研究。在典型的基于 Fat-Tree 的集群中，每台服务器都配备一个 NIC（40 Gbps 至 400 Gbps），并且 K 个服务器排列在连接到架顶式 (ToR) 交换机的机架中。然后，ToR 交换机连接到聚合交换机，以提供跨机架的连接，形成一个 Pod。最后，Pod 与主干交换机互连，允许 CPU 集群中的服务器之间进行any-to-any通信。

相比之下，网络密集型机器学习工作负载的兴起导致了以 GPU 为中心的集群的主导地位，其中各个 GPU 都有专用的 NIC 。图 2 展示了典型 GPU 集群的网络架构。每个 GPU 都有两个不同的通信接口：(i) NVLink 接口，支持高带宽但短距离互连；(ii) 传统的支持 RDMA 的 NIC。NVLink 连接 K 个 GPU，为每个 GPU 提供太比特级的无阻塞any-to-any带宽输入/输出（例如，第四代 NVLink 为 7.68 Tbps）。这组具有快速互连的 GPU 形成高带宽域（HB 域）。传统上，HB 域仅限于单个服务器（例如，具有 K = 8 或 16 个 GPU 的 DGX 服务器）。然而最近，Nvidia 宣布推出 GH200 超级计算机，将 K = 256 个 Grace Hopper 超级芯片互连，形成跨多个机架的一个 HB 域 。

然而，即使使用 256 个 GPU，一些LLMs对于单个 HB 域的训练时间也太长。例如，假设 GPU 利用率完美，PaLM-540B 模型在 GH200 超级计算机上需要 117 天才能完成。这些模型需要跨多个 HB 域并行化。

为了能够跨多个 HB 域训练 LLM，GPU 集群操作员使用支持 RDMA 的 NIC 将多个 HB 域互连在一起。互连 HB 域的传统网络架构称为rail-optimized网络。在rail-optimized架构中，HB 域内的 GPU 标记为从 1 到 K。轨道是不同 HB 域上具有相同索引（或等级）的 GPU 的集合，通过轨道交换机互连。例如，图 2 分别以红色和黄色显示了轨道 1 和轨道 K。这些导轨交换机随后连接到骨干交换机，形成全平分any-to-any Clos 网络拓扑。该网络确保不同 HB 域中的任何一对 GPU 都可以以网络线路速率（GH200 为 400 Gbps Infiniband 网络）进行通信。例如，GPU 1、Domain 1 和 GPU 1、Domain 2 之间的流量仅通过 Rail Switch 1，而 GPU 1、Domain 1 和 GPU 2、Domain 2 之间的流量则通过各自的 Rail 和 Spine 交换机。

2.2 分析LLMs的网络流量

我们现在分析分布在由数百台 Nvidia GH200 超级计算机组成的集群中的 OpenAI 的 GPT3 、Meta 的 OPT3-175B 和 Google 的 PaLM-540B 的流量模式和迭代时间。这些LLMs代表了具有公开可用参数的尖端语言模型。每台 GH200 超级计算机均包含两层 NVSwitch 架构，可在 256 个 H100 GPU 上实现 2 Pbps 的全平分带宽（每个 GPU 7.68 Tbps）。此外，每个 GPU 都有一个 Connect-X7 HCA Infiniband 网络接口，为每个 GPU 的输入/输出提供 400 Gbps 的网络带宽。在此设置中，每台 GH200 超级计算机形成一个 HB 域。

我们的分析使用 Nvidia的基准并行化策略来确保最佳的 GPU 利用率。我们使用混合数据并行性（DP）和算子内模型并行性（MP），但我们的结论对于管道并行性仍然相似。

DP同步利用分层AllReduce算法来限制跨HB域的流量。我们为每个 LLM 模型采用最佳并行化策略，并计算最终的流量模式。请注意，对于 PaLM 模型，批量大小在整个训练过程中会发生变化。

图 3a 说明了 128 个 GH200 超级计算机集群的服务器对之间一次训练迭代的流量类型分布，图 3b 显示了每种流量类型的流量百分比。通信主要有两种类型：运营商内模型并行 (MP) 流量和数据并行生成的 AllReduce 流量。MP 流量发生在参与模型并行组的 GPU 内，该组始终适合 HB 域。对于数据并行性，分层 AllReduce 算法进一步将通信划分为两个阶段，其中第一阶段跨 HB 域同步参数，而第二阶段在 HB 域内同步参数。该算法确保第二级承载更多流量，以更好地利用可用带宽。虽然这些类型的流量在不同的 GPU 对之间不会重叠，但图 3a 表明超过 99% 的 GPU 对不承载任何流量，只有不到 0.25% 的 GPU 对在它们之间承载 MP 和第二阶段 DP 流量。同时，图 3b 显示这些流量类型占总传输数据的 90% 以上。回想一下，这两种类型的流量停留在 HB 域内，这表明 HB 域带宽的有效利用以及对互连 HB 域的网络结构的需求较低。这种模式在所有模型中都是一致的，这表明在 HB 域之上为 LLM 模型构建具有any-to-any连接的集群是过度的。

在 HB 域内，互连需要支持大量的any-to-any通信，以训练不同的LLMs。为了说明这一点，图 4 显示了 GPT3 和 OPT3-175B 训练迭代期间流量矩阵的热图。

在此图中，每组连续的 256 个 GPU 都驻留在同一个 HB 域内。图 4a 和 4b 展示了 HB 域内的通信模式。请注意，流量很大（服务器对之间的流量高达 1 GB），并且由于这些模型之间的并行化策略不同，因此模式会有所不同。对角线上的方块代表 MP 流量，而对角线外的其余部分代表第二阶段 DP 流量。HB 域内的任意互连可提供最大的灵活性以适应不同的LLMs。

然而，HB 域内所需的高带宽any-to-any连接在它们之间并不需要。图 4c 和 4d 缩小到前四个 HB 域，其中发生跨 HB 域的第一阶段 DP 流量。与 HB 域内的流量相比，跨 HB 流量要小得多（每个条目约 1 MB）并且更加稀疏。这是因为跨 HB 域通信仅发生在具有相同等级的 GPU 之间（即同一轨道上的 GPU）。此外，该模式是同质的，因为跨 HB 域通信模式在 LLM 之间保持相同（图 4c 和 d 中的非对角线）。

这些观察结果表明，可以删除不承载任何网络流量的链接，而不会损害 LLM 的训练性能。我们的分析表明，any-to-any 400 Gbps Clos 网络中 33% 的链路是可移除的。图 5 显示了该替代网络（标记为 Any-to-Any Trimmed 400 Gbps）与最先进网络相比的训练迭代时间。鉴于断开的链路不支持任何流量，这两种拓扑提供相同的性能。

为了正确看待这一点，我们还考虑了理想的性能，假设集群中的所有 GPU 都与 NVSwitch 互连，形成单片 7.68 Tbps 全平分带宽any-to-any网络，这意味着 19.2 倍的增长与最先进的技术相比，在全平分带宽中。这样的网络在实践中是不可能建立的。然而，性能提升与增加的带宽并不成正比，所有模型的性能提升范围为 1.06× 到 1.88×，如图 5 所示。这样的结果为 LLM 集群运营商重新评估传统的any-to-any 提供了令人信服的网络设计论据。

3 我们提出的LLMs集群

在本节中，我们提出了一种专门针对 LLM 集群的新网络设计。我们用数学模型参数化我们的设计，并提出一套全面的指南来确定此类集群的参数。

3.1 Rail-only 网络设计

我们提出了一种网络架构，可以跨所有 GPU 的any- to-any范式。图 6 展示了我们的网络架构，我们将其命名为“rail-only”。与图 2 所示的传统rail-optimized GPU 集群相比，我们的网络保留 HB 域并仅在同一轨道上提供连接。

实现我们提出的架构的一种直接方法是从图 2 中删除主干交换机，并将所有连接轨道交换机到主干的上行链路重新用作到 GPU 的下行链路。因此，每条rail都由专用但独立的 Clos 网络连接。在本文的其余部分中，我们的分析基于rail-only网络的实现，尽管其他技术也适用于其他rail互连。

我们的rail-only网络架构消除了不同轨道中具有不同等级的 GPU 之间的网络连接。然而，通过 HB 域转发数据，这种通信仍然是可能的。例如，从 GPU 1、域 1 到 GPU 2、域 2 的消息可以首先通过第一个 HB 域路由到 GPU 2、域 1，然后通过网络传输到最终目的地。尽管我们的分析表明 LLM 流量不需要此类转发，但控制消息、测量或训练该集群中的其他 DNN 模型可能需要这种连接。我们在第 5 节中提供了有关处理其他 DNN 模型的更多讨论。

3.2 Rail-only 网络分析

表 1 描述了我们分析中使用的参数。我们考虑一个具有 N 个 GPU 和大小为 K 的 HB 域的网络。每个 HB 域的带宽是 BF，网络带宽是 BS。我们重点分析使用混合两阶段数据和模型并行作为并行化策略，并推导出训练迭代时间和最佳模型并行组大小的数学模型。

单次迭代所需的总时间是计算时间和通信时间的总和，公式为：

我们将模型的计算持续时间建模为所需的浮点运算量与集群的总 FLOP 的比率。至于集体通信持续时间，我们仅计算带宽时间，因为延迟项对于我们选择的集体通信算法而言是恒定的，并且相对于带宽项较小。对于所有三种类型的通信，我们的计算基于带宽和延迟最优集体通信算法，该算法通过连接的 GPU 组内的所有对所有通信来实现。总迭代时间为：

利用这个公式，我们可以计算出最优的 NM ：

有趣的是，NM* 独立于 K、BF 和 BS，这表明不依赖于 HB 域的大小或任何互连的确切带宽。

3.3 HB 域的理想大小是多少？

设计过程中自然会遇到的问题是，HB 域的理想大小应该是多少？在图 7 中，我们改变 HB 域大小 (K)，并绘制大小为 16384、32768 和 65536 GPU 的 GPU 集群的 GPT3 和 OPT3-175B 的训练迭代时间。对于每个簇大小，我们使用根据 GH200 的带宽和计算能力参数计算出的相应最佳 NM*。我们还计算了“Any-to-Any 7.68 Tbps”的训练迭代时间作为理想情况的性能。回想一下，这个设计点代表了理想化的场景，其中每个 GPU 都与全平分 NVLink 结构连接，或者等效地 K = N 的情况，并且在实践中是无法实现的。理想情况下使用非分层AllReduce算法来利用统一的高带宽互连，其迭代时间为：

如图 7 所示，随着 HB 域大小的增大，性能增益会降低。HB 域大小从 32 到 64 的转变使 GPT3 和 OPT3-175B 的训练迭代时间分别减少了 12% 和 37%，而从 512 到 1024 的增量仅实现了 0.95% 和 7.5% 的增益。通信时间增益的减少可以归因于阿姆达尔定律，因为 DNN 的计算时间在所有实例中保持不变。我们认为，当前的 GH200 超级计算机的 HB 域大小为 256，只要选择适当的批量大小，就非常适合当今 LLM 训练的需求。我们推迟第 3.4 节中对批量大小的分析。与此同时，扩大这一规模的未来技术进步将进一步有利于训练性能，减少训练迭代时间，使其更接近理想情况，而不需要跨 HB 域的网络中的any-to-any连接。

3.4 批量大小对网络设计的影响

公式 2 表明，随着批量大小 b 的增加，最佳模型并行组大小会缩小。在这种情况下，HB 域内可以容纳更多数据并行通信，从而有利于整体性能。图 7a 和 7b 之间的比较分析揭示了这一优势，因为 GPT3 和 OPT3-175B 模型实际上具有相同的模型结构，并且训练迭代时间的差异仅来自批量大小的选择（GPT3 为 32M token，OPT3-175B 为 2M token。

为了进一步了解批量大小对训练时间的影响，我们分析了 PaLM-540B 模型在 32768 GPU 集群上的性能，同时将 HB 域大小从 K = 256 更改为 1024。在训练期间，PaLM 模型自动更改其批量大小从 512 到 2048 个序列（1M 到 4M token）。图 8a 绘制了迭代时间随批量大小变化的变化。对于所有 HB 域大小，迭代时间表现出相同的轨迹；然而，由于阿姆达尔定律，相对性能（HB 域大小的迭代时间与理想情况的迭代时间之比）随着批量大小的增加而提高。图 8b 代表了这一趋势。当 K = 256 时，随着批量大小从 512 个序列增加到 2048 个序列，相对性能从 53% 增加到 74%。

先前的研究表明，LLM 训练受益于较大的批量大小 ，使其非常适合我们的rail-only设计。此外，虽然批量大小参数通常是以 ML 为中心的优化指标，但我们的分析表明批量大小对综合训练性能的影响超出了收敛所需的迭代总数。

4、 网络成本分析

我们的rail-only网络架构通过消除未使用的网络连接，明智地减少了 LLM 训练的网络资源。本节将我们提出的方法与最先进的轨道优化 GPU 集群的网络成本进行比较。我们计算每个网络设计所需的交换机（#SW）和收发器（#TR）的数量，并根据先前工作1中报告的数字得出总网络设备成本。本节重点介绍仅使用电分组交换机来构建网络；然而，使用光纤直接连接技术可以进一步降低成本,我们将有关直连网络的讨论推迟到第 5 节。

我们在表 2 中列举了构建最先进的网络架构和我们提出的架构所需的交换机和收发器的数量，并考虑了可变的集群大小和网络交换机基数。请注意，对于最先进的架构，为了使用最少的网络资源，GPU 的每个轨道在某些情况下并不是物理上分离的。因此，数据中心操作员必须手动配置交换机以实现所需的跨轨隔离，从而实现轨优化设计。

表 2 的最后一列说明了我们的设计相对于最先进的设计所节省的交换机和收发器总成本。与最先进的设计相比，我们的rail-only设计显着降低了网络成本 37% 至 75%，同时实现了同等性能。这种减少源于消除核心层交换机并减少每个轨道内的交换机层数量。令人惊讶的是，基数为 64 的交换机在两种集群规模中都能实现最坏情况下的成本降低。在这种情况下，最先进的设计需要三层 Clos 网络，而rail-only设计则需要每个轨道两层。尽管如此，我们的设计仅需要交换机总数的四分之三，同时实现与最先进设计相同的性能。

5、讨论

LLM趋势。随着摩尔定律的放缓，目前 LLM 计算需求的增长速度超过了人工智能加速器和网络速度的进步，需要超大规模集群和更高效的互连。我们消除any-to-any网络连接的立场是满足LLMs训练网络要求和维持LLMs增长趋势的第一步。我们还承认在不影响性能的情况下减少语言模型的大小和资源需求的持续努力。这些工作补充了我们的设计，因为我们的设计减少了网络资源并即使对于较小的语言模型和集群也能保持性能。

LLM推理。本文探讨了LLM的培训工作量，但推理代表了LLM产品周期的另一个重要部分。推理比训练需要更少的计算资源，因为它涉及通过 LLM 移动更少的数据并且只计算前向传播 。因此，每个 HB 域自然成为一个推理服务域，并且rail-only连接有助于负载平衡多个推理域。我们将对 LLM 推理的详细研究留给未来的工作。

直连网络拓扑。正如第 4 节中提到的，数据中心运营商可以利用直连网络拓扑进行跨轨互连。为了最大限度地提高此类设计的有效性，我们建议通过 NIC 接口拆分来增加连接到每个 GPU 的网络接口数量 。此外，我们建议使用可重新配置的光开关为跨 HB 域的互连提供更大的灵活性。这样的设计还允许为即将到来的工作负载重新配置一些跨轨连接，这些工作负载的行为与LLM不同。我们相信，将我们提出的设计与光学可重构网络交换机相结合，开辟了 AI-ML 集群的新研究路线。

其他 ML 工作负载和限制。尽管我们提出的rail-only架构专注于专门针对LLM的网络设计，但与其他工作相结合时，我们的设计对于许多其他 DNN 工作负载来说是高效的。最近的工作试图使并行化策略和集体通信算法对任何 DNN 模型具有带宽感知能力，这些模型已经产生了类似于 LLM 的流量模式。对于需要 GPU 跨轨少量流量的并行化策略，集群可以使用第 3 节中描述的转发。我们设计的主要挑战是所有 GPU 之间的全面通信，这通常出现在具有大型嵌入表的推荐模型。转发方案会导致拥塞并降低点对点流量的性能。我们承认，any-to-any流量是机器学习工作负载中最具挑战性的流量模式之一。一些潜在的解决方案包括通过超额订阅网络重新引入小型any-to-any容量、利用快速可重新配置的网络结构，以及通过调整机器学习模型本身来减少最初生成的所有到所有流量。

容错性。乍一看，rail-only的容错能力可能比标准 Clos 网络要差。然而，假设任一网络中的铁路转辙器发生故障。连接到故障交换机的所有 GPU 将变得不可用，从而使两种拓扑在轨道交换机的容错方面相同。相反，我们的设计需要更少的交换机，这自然减少了故障点。数据中心运营商可以通过添加额外的轨道交换机来增加冗余容量，并且与最先进的任意网络设计相比，我们的设计仍然更具成本效益。直连网络还可以提高容错能力，因为即使控制平面发生故障，光交换机也可能保持功能。

6 结论

本文挑战了专门用于训练大型语言模型的 GPU 集群的传统any-to-any网络架构。我们提出了一种称为“rail-only”的新架构，它符合LLM的独特特征和需求，可降低高达 75% 的成本，同时保持与当前最先进的 Clos 网络相同的性能。

（正文完）