Rail-only: A Low-Cost High-Performance Network for Training LLMs with Trillion Parameters

作者/机构: Weiyang Wang (MIT CSAIL), Manya Ghobadi (MIT CSAIL), Kayvon Shakeri (Meta), Ying Zhang (Meta), Naader Hasani (Meta)

A1 主要贡献

本文针对超大规模大型语言模型（LLM）的训练，提出了一种低成本的网络架构。随着LLM模型规模和计算需求的增长速度超过硬件发展，超大规模GPU数据中心成为必然趋势。然而，将传统的Clos网络扩展到数万个GPU面临着成本高昂、功耗巨大以及死锁和PFC风暴等性能问题。

核心问题：当前用于GPU集群的“轨道优化（Rail-optimized）”Clos网络提供了任意GPU对之间的全对分带宽连接，但这对于LLM训练是否是必要的？扩展这种网络架构到30,000个GPU的规模，不仅成本高达2亿美元，峰值功耗也接近4.6兆瓦，导致数据中心供应商不得不采用超售（over-subscription）来控制成本，从而加剧了性能问题。

研究目标：本文旨在证明高效训练LLM并不需要在所有GPU之间建立任意到任意（any-to-any）的连接。基于此，论文提出了一种可立即部署的解决方案，旨在使用商用电交换机来降低LLM数据中心的成本和能耗。

创新贡献：
1. LLM训练流量模式分析：论文分析了LLM训练的流量模式，并指出在采用最优并行化策略时，LLM训练工作负载仅在一小部分GPU子集（如一个Nvidia DGX服务器）内需要高带宽的任意到任意连接。而跨平台的通信大多发生在集群中具有相同本地排名（rank）的少数GPU对之间。这一发现表明，传统的为实现任意到任意连接而构建的Clos网络中的Spine（核心）层交换机对于分布式LLM训练而言是不必要的，增加了额外的复杂性、成本和功耗。
2. 提出Rail-only网络架构：基于上述观察，论文提出了名为“Rail-only”的低成本网络架构。该架构移除了Clos网络中的Spine层交换机，仅为那些具有显著网络流量的GPU组（即“轨道”，Rail）提供连接。与最先进的Rail-optimized设计相比，Rail-only架构在不牺牲性能的前提下，移除了不承载重要流量的网络设备。
3. 支持MoE模型并评估性能：该架构通过转发机制支持需要全对全（all-to-all）通信的专家混合（MoE）模型，由此产生的完成时间开销仅为8.2%到11.2%。论文还分析了该设计的容错属性，并提供了对不同网络和训练参数性能影响的深入分析。评估结果显示，Rail-only网络相比传统方案，可将网络成本降低38%至77%，网络功耗降低37%至75%。

A3 背景知识与关键观察

II. 背景知识

A. 平台内连接：高带宽域

高带宽域的定义与特性。资源密集型的机器学习工作负载推动了以GPU为中心的平台的发展，这些平台为多GPU工作负载进行了优化。为了满足通信需求，这些平台在GPU的本地域内使用高带宽的本地互连技术。尽管不同制造商的平台在计算浮点性能（FLOPs）、GPU和CPU架构或物理互连技术上有所不同，但它们都具备一个共同属性：为GPU之间提供了数Tbps的内部带宽。例如，Nvidia的DGX H100服务器【【14, Dgx h100 computer, 2023, NVIDIA】】包含八个H100 GPU，通过NVSwitches互连，内部提供3.6 Tbps的无阻塞带宽。最近发布的GB200 NVL72计算机【【15, Gb200 nvl72 computer, 2024】】在一个机架内连接36个GB200 Superchips，每个GPU的第五代NVLink带宽达到7.2 Tbps。AMD MI300X平台则采用AMD的Infinity Fabric技术【【6, Amd instinct mi300x platform, 2023, AMD】】以全网状拓扑连接八个MI300X加速器，每个GPU带宽为3.6 Tbps。类似地，Nvidia的DGX GH200超级计算机利用多层NVSwitch拓扑将平台规模扩展到256个GPU，同时保持3.6 Tbps的全对分GPU间带宽【【16, Nvidia dgx gh200, 2023, NVIDIA】】。本文将具有Tbps级别内部连接的平台称为“高带宽（HB）域”，相应的互连称为HB互连（HBI）。

B. 平台间连接：NIC域

NIC域与Rail-optimized网络架构。虽然以GPU为中心的平台使用NVLink或Infinity Fabric等技术提供了高内部带宽，但它们只能扩展到有限数量的GPU。为了超越单个平台的规模，运营商依赖于以太网或Infiniband等传统网络技术来连接不同平台的网络接口卡（NIC）。本文将这种平台间网络称为“NIC域”。在NIC域中，最先进的互连技术是基于一种著名的网络架构，称为Rail-optimized网络【【17, Doubling all2all performance with nvidia collective communication library 2.12, 2022, NVIDIA Developer Blog】】。该架构被广泛用于高性能计算（HPC）工作负载。Rail-optimized网络比传统的以CPU为中心的数据中心网络更适合DNN，但由于其主要为HPC工作负载设计，未能充分利用LLM训练工作负载的独特流量模式。

Rail-optimized网络的设计与演进。传统的专为服务不可预测和突发性CPU密集型工作负载的数据中心设计，采用了Clos网络架构【【7, A scalable, commodity data center network architecture, 2008, ACM SIGCOMM】】【【18, Introducing data center fabric, the next-generation facebook data center network, 2014, Meta】】，以提供服务器对之间的任意到任意连接。用于GPU训练集群的Rail-optimized网络正是从数据中心Clos网络演变而来【【17, Doubling all2all performance with nvidia collective communication library 2.12, 2022, NVIDIA Developer Blog】】【【19, Nvidia dgx superpod: Next generation scalable infrastructure for ai leadership, reference architecture, 2023, NVIDIA】】，如图1所示。对于一个大小为K的HB域的GPU平台，总共有K个“轨道（rail）”，每个轨道由属于不同HB域但具有相同本地排名（local rank）的GPU组成【【20, Using multirail networks in high-performance clusters, 2001, IEEE International Conference on Cluster Computing】】。Rail-optimized网络将这些相同轨道的GPU连接到同一组交换机下，我们称之为轨道交换机（rail switches）。图1中分别用红色和黄色高亮显示了轨道1和轨道K。将相同排名的GPU连接到相同的轨道交换机确保了它们之间的延迟最低。这种连接方式是理想的，因为最优的DNN并行化策略会将其NIC域的流量集中在具有相同本地排名的GPU之间【【17, Doubling all2all performance with nvidia collective communication library 2.12, 2022, NVIDIA Developer Blog】】。

Rail-optimized网络的问题。Rail-optimized架构利用了DNN训练流量的强局部性，通过将相同排名的GPU连接到同一个ToR交换机，使得同一轨道内的GPU之间延迟很低。网络的其余部分则采用多层核心交换机（spine switches）来连接各个轨道交换机，从而形成一个全对分、任意到任意的Clos网络拓扑。这个网络确保了不同HB域中的任意一对GPU仍能以数百Gbps的线速进行通信。例如，在图1中，域1的GPU 1和域2的GPU 1之间的通信仅通过轨道交换机1，而域1的GPU 1和域2的GPU 2之间的通信则需要经过各自的轨道交换机和核心交换机。尽管Rail-optimized网络利用了流量局部性，但它忽略了一个根本问题：核心交换机是否是必需的？下一节将更详细地分析LLM训练流量，以探究无核心（spineless）网络架构设计的可能性。

III. LLM流量模式分析

A. MegatronLM的流量模式

LLM并行化策略产生的流量构成。我们通过模型超参数和并行化策略计算网络传输大小，分析了采用混合数据、张量和流水线并行化的LLM所产生的流量模式。我们首先研究了MegatronLM论文【【21, Efficient large-scale language model training on gpu clusters using megatron-lm, 2021】】表1中描述的一系列GPT模型，参数量从1456亿到1万亿不等，分布在多达3072个GPU上。这些模型分布在一个由最多384个DGX A100服务器组成的集群中，每个服务器的HB域大小为8。我们的分析采用了与MegatronLM相同的并行化策略以确保最优的GPU利用率，并使用基于环（ring-based）的集合通信，因为它是带宽最优且为NCCL中的默认算法。通信主要有三种类型：来自张量并行（TP）的AllGather和ReduceScatter流量、来自数据并行（DP）的AllReduce流量，以及来自流水线并行（PP）的点对点流量。图2a展示了单次训练迭代中每种通信类型的流量体积百分比，图2b则显示了所有GPU对之间的通信类型分布。

流量模式揭示了NIC域的低通信需求。张量并行（TP）的流量发生在参与同一个TP等级（rank）的GPU之间，而这些GPU通常被限制在一个HB域内。数据并行（DP）和流水线并行（PP）的流量则发生在NIC域，其流量体积远小于TP流量，如图2a所示。尽管不同并行化策略的流量不会在同一对GPU之间重叠，但图2b显示，超过99%的GPU对之间没有任何流量，而承载TP流量的GPU对不到0.04%。同时，图2a表明TP流量占总传输数据的75%以上。由于TP流量局限于HB域内，这表明HB域带宽被高效利用，而NIC域的需求很低。这一模式在我们研究的所有GPT模型中都是一致的，这说明在HB域之上为LLM模型构建一个具有任意到任意连接的GPU数据中心是过度的。

B. LLM在NIC域的流量

LLM的NIC域流量特征：稀疏且局限于轨道内。MegatronLM中采用的并行化策略在NIC域产生的流量与HB域相比微不足道。图3展示了训练GPT-1T模型的流量热图。在此图中，每连续8个GPU位于同一个HB域（橙色高亮），而编号相差8的GPU属于同一个轨道（粉色高亮）。图3a展示了一个流水线阶段内的流量模式，而图3b展示了前四个流水线阶段的流量。HB域内的流量体积非常大（GPU对之间约300 GB），而NIC域的通信量则降至约6 GB。更重要的是，NIC域的传输从不经过核心交换机——这些网络传输只发生在轨道内部。

最优并行化策略确保流量局限于轨道内。我们认为，对于所有没有稀疏MoE层的LLM，只要采用最优的并行化策略，其在NIC域内总是会产生稀疏、低容量且局限于轨道内的流量。根据设计，唯一离开HB域的流量是来自流水线并行的点对点流量，或来自张量并行（TP）和数据并行（DP）的集合通信流量（AllGather、ReduceScatter和AllReduce）。对于流水线并行（PP），由于LLM并行化的对称性，每个流水线阶段包含相同数量的GPU。因此，总可以将流水线阶段进行布局，使得跨阶段的流量只在NIC域中相同排名的GPU之间传输，从而保持在同一轨道内。

分层集合通信算法及其流量模式。张量并行（TP）和数据并行（DP）可能在HB域和NIC域中都引发集合通信流量。MegatronLM的例子总是将TP和DP分别限制在HB域和NIC域内。虽然这种划分对于LLM很常见，但并非普遍适用。例如，仅使用DP训练一个较小的模型会导致所有GPU参与同一个DP等级，从而在HB域和NIC域中都进行相同的AllReduce操作。在这些情况下，训练集群可以使用分层集合通信算法来实现近乎最优的性能。下面，我们介绍这些算法，分析其性能，并说明它们在NIC域的流量也保持在轨道内。

分层AllGather算法分析。分层集合通信算法是为多层网络拓扑设计的。我们在此介绍分层AllGather算法，并指出，对于LLM训练中出现的其他集合操作，ReduceScatter通过反转AllGather的调度可实现相同性能，而AllReduce则等同于一个ReduceScatter后跟一个AllGather。我们关注带宽性能而忽略延迟，因为在LLM训练期间数据传输量很大，通信运行时间主要由带宽决定。表I列出了本节使用的变量。我们假设执行AllGather操作的GPU被排列成一个 $x \times y$ 的网格，其中每 $x$ 个GPU属于同一个HB域，分布在 $y$ 个HB域中。基础的分层AllGather分两阶段执行：首先，该算法为每个轨道的GPU收集部分数据，而不在HB域内传输数据。如果执行AllGather的总数据大小为 $D$，则所有GPU在网络中交换的数据量为 $D(y - 1)/x$。此操作有效地为GPU创建了更大的数据分片，以便在每个HB域内重新运行AllGather。因此，每个HB域在第二阶段执行AllGather，总共产生 $D(x-1)$ 的传输量。假设HB域内的 $x$ 个GPU的带宽容量为 $CF$，NIC域中的 $y$ 个GPU的带宽容量为 $CS$，则总的AllGather运行时间为：

与PP通信类似，通过将逻辑上的 $x \times y$ 网格GPU适当地映射到GPU集群，该算法仅在同一轨道内的GPU之间产生流量。

分层通信在大型平台上的流量验证。图4a展示了使用跨越NIC和HB域的分层集合通信训练GPT-1T的流量模式，这与MegatronLM的情况不同。我们计算并分析了在一个由16台Nvidia DGX GH200超级计算机【【16, Nvidia dgx gh200, 2023, NVIDIA】】（共4096个GPU）组成的集群中，以4096的批大小训练GPT-1T模型的流量模式。每个DGX GH200超级计算机包含一个两层NVSwitch架构，为256个H100 GPU之间提供3.6 Tbps的GPU到GPU带宽。此外，每个GPU都有一个Connect-X7 HCA Infiniband网络接口【【16, Nvidia dgx gh200, 2023, NVIDIA】】，提供400 Gbps的进出GPU的网络带宽。在此设置中，每个DGX GH200超级计算机构成一个HB域。图4展示了此设置下的流量体积百分比和热图。并行化策略的总数据并行度为64，它跨越了每个HB域中的32个GPU以及网络中的两个HB域。图4b和4c展示了分层AllReduce算法的流量热图，该算法将AllReduce流量分配到每个DP组中。值得注意的是，网络流量保持在轨道内部（GPU编号相差256）。分层算法高效地利用HB域的带宽承载了98%的AllReduce流量，而网络域仅承载另外的2%。

C. 专家混合模型的全对全流量

MoE模型引入全对全通信。具有稀疏门控专家混合（MoE）层的LLM展现出与上述模型不同的流量模式。MoE层提供了一种在不显著增加计算需求的情况下增大LLM规模的替代方法。通过MoE，模型的一部分被一组“专家”神经网络替代，其中一个门控网络将每个令牌（token）路由到不同的专家，从而只激活模型的一部分。MoE的典型并行化策略是专家并行，即每个专家被分配到集群中的不同GPU上。

MoE的流量模式及其对网络的需求。与传统LLM不同，采用专家并行的MoE要求每个专家与模型的其余部分进行通信，从而产生密集的通信模式。确切的流量热图取决于门控网络和令牌分布。为简单起见，本节假设令牌分布是均匀的。图5显示了在16台DGX A100服务器上训练来自DeepSpeedMoE【【22, Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale, 2022】】的MoE-1.3B模型的流量矩阵。该模型包含128个专家。模型的静态部分使用数据并行（DP），而MoE部分使用专家并行。由于每个GPU包含一个不同的专家，均匀的令牌分布会在网络上产生一个均匀的全对全（all-to-all）流量模式。乍一看，这种流量模式使得Rail-optimized网络中的核心交换机至关重要，因为跨越不同轨道上的GPU之间的流量将通过核心交换机。然而，正如我们将在下一节中展示的，我们不必依赖核心交换机：使用HB域作为转发步骤可以实现接近最优的性能。

A2 方法细节

IV. Rail-only网络设计

基于以上观察，本节提出了一种名为Rail-only的新型网络架构，它背离了任意到任意GPU连接的范式。我们首先介绍Rail-only网络的架构设计，然后讨论其路由策略和容错特性。

A. 架构设计

Rail-only架构的核心思想。图6展示了我们的Rail-only网络架构。与图1所示的传统Rail-optimized GPU集群相比，我们的Rail-only网络保留了HB域，但取消了NIC域中所有GPU的全对分连接。取而代之的是，我们只确保每个轨道（rail）内的GPU之间通过全对分网络连接。

具体实现。一个直观的说明是我们提出的架构移除了核心交换机（图1中的spine switches），并将连接轨道交换机到核心交换机的上行链路（uplinks）重新用作连接到GPU的下行链路（downlinks）。因此，每个轨道都有一个专用的Clos网络来连接。与Rail-optimized架构相比，Rail-only设计通过为每个轨道构建多个较小的Clos网络来节省交换机和收发器的数量，这需要网络中使用更少层级的交换机。

B. Rail-only网络的路由

跨轨道通信通过HB域转发。我们的Rail-only网络架构移除了不同轨道中不同排名GPU之间的网络连接。但这种通信仍然可以通过HB域进行数据转发来实现。例如，在图6中，一条从域1的GPU 1到域2的GPU 2的消息，可以首先通过第一个HB域路由到域1的GPU 2，然后再通过网络传输到最终目的地。先前的工作已经表明，这种路由方案会引入“带宽税”（bandwidth-tax）【【13, TopoOpt: Co-optimizing network topology and parallelization strategy for distributed training jobs, 2023, USENIX NSDI】】【【23, Rotornet: A scalable, low-complexity, optical datacenter network, 2017, ACM SIGCOMM】】【【24, Expanding across time to deliver bandwidth efficiency and low latency, 2020, USENIX NSDI】】，即由于转发导致网络中的物理流量增加。然而，在本节中，我们表明由于HB域和NIC域之间的带宽不对称性，由带宽税引起的性能下降是微不足道的。

MoE模型全对全通信的性能分析。我们以第III-C节中描述的具有MoE层的LLM为例，因为它产生了具有挑战性的全对全通信模式。乍一看，这种流量模式对Rail-only网络是具有挑战性的，因为大部分全对全流量需要转发。然而，由于HB域比NIC域快得多，在HB域上转发网络流量只会产生很小的开销。下面，我们使用一个两步转发路由算法来推导Rail-only网络下均匀全对全流量的减速因子。值得注意的是，该策略已在NCCL中以“PCIe × NVLink”（PXN）的形式实现，用于在Rail-optimized网络的Spine层过载时避免拥塞【【17, Doubling all2all performance with nvidia collective communication library 2.12, 2022, NVIDIA Developer Blog】】。

性能开销的量化推导。我们使用表I中定义的相同变量进行以下推导。考虑一个 $x \times y$ 的GPU网格，其中 $x$ 个GPU位于一个HB域内，y个HB域通过Rail-only或Rail-optimized网络连接。回顾一下，HB域的带宽为 $C_F$，而NIC域的每个GPU对的带宽为 $C_S$。对于Rail-optimized网络，每个GPU通过HB域和全对分的NIC域直接向其目的地发送流量。假设每对GPU通信的数据量为 $D$，则总的全对全完成时间为：

对于Rail-only网络，两步算法首先在每个轨道内运行一个全对全通信，为每个GPU准备好在其轨道上发送的所有数据。然后，在每个轨道内，每个GPU在HB域中运行第二次全对全通信以完成传输，有效的数据分片大小为 $xD$。注意，第二步包含了带宽税。总传输时间为：

这两个表达式的差异在于 $y(x - 1)D/C_F$，这是在HB域内转发的成本。当 $y(x - 1) \approx x(y - 1)$ 时，减速因子大约为 $C_S/C_F$，对于DGX A100和DGX H100代的GPU平台，这个值分别为8.2%和11.2%。由于两个域之间的带宽不对称，这个因子很低。此外，这种减速仅适用于全对全通信，而根据图5，它只占总流量的27%。因此，根据阿姆达尔定律，这个小开销可以忽略不计。我们注意到，这种特性也使得我们的网络设计适用于其他类别的DNN模型。

C. Rail-only网络的容错属性

容错对大型GPU集群的重要性。容错对于需要进行长期LLM训练的大型GPU集群至关重要。本节研究了Rail-only网络与传统Rail-optimized网络在容错属性上的对比。

链路和交换机故障。假设一个轨道交换机或一条链路发生故障。在这两种网络架构中，连接到故障交换机或链路的GPU都将变得不可用，因此在交换机和链路的容错方面，两种设计是相同的。然而，我们的设计需要的交换机更少，这自然减少了故障点。数据中心运营商可以通过增加额外的轨道交换机来增加冗余容量，即便如此，我们的设计仍然比任意到任意网络设计更具成本效益。

GPU平台故障。对于由类似DGX的服务器组成的GPU集群，每个服务器就是一个HB域。当一个服务器发生故障时，网络运营商将任务迁移到另一个健康的服务器上。对于新服务器，Rail-only的连接性保持不变。对于类似GB200的集群，一个super-chip模块类似于一个服务器；因此，其故障模式与单个GPU故障相同，我们将在下面讨论。

HB域内有空闲GPU时的单GPU故障。对于单个GPU故障，我们分别讨论两种不同情况。第一种情况是，在发生故障的GPU所在的HB域中存在另一个空闲的GPU。在这种情况下，Rail-optimized网络可以直接用空闲GPU替换故障GPU，而不会破坏HB域的完整性。对于Rail-only网络，一个可能的解决方案是利用可重构光交换机。为了提高鲁棒性，我们在GPU和轨道交换机之间增加少量光交换机，以允许动态重构轨道。当一个GPU发生故障时，光交换机重新配置，将一个健康的空闲GPU替换故障的GPU。这种方法在概念上类似于主要使用可重构光交换机的网络设计的故障恢复机制【【13, TopoOpt: Co-optimizing network topology and parallelization strategy for distributed training jobs, 2023, USENIX NSDI】】【【25, Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings, 2023】】【【26, Jupiter evolving: Transforming google’s datacenter network via optical circuit switches and software-defined networking, 2022, ACM SIGCOMM】】。我们将对带有光交换机的rail-only网络的深入分析留作未来工作。

HB域全满时的单GPU故障。另一种故障模式是，在一个完全被占用的HB域中，一个GPU发生故障，需要从不同的HB域中找一个替代GPU。在这种情况下，Rail-only设计阻止了将故障GPU上的任务迁移到集群中另一个空闲GPU上，而这在Rail-optimized网络中是可能的。然而，即使在Rail-optimized网络中，这种解决方案也是不可取的。新的GPU不再属于与故障GPU相同的HB域，这会产生一个瓶颈，将HB域的性能降低到NIC域的水平。相反，我们提出了两种解决方案。对于HB域较小的平台，我们将包含故障GPU的整个HB域上的任务迁移到一个新的HB域。对于较大的HB域（例如，拥有K=256的DGX GH200超级计算机），这些HB域由一个带有光核心层的多层拓扑构成【【16, Nvidia dgx gh200, 2023, NVIDIA】】。一种潜在的方法是像前一种故障情况一样增加光交换机。当GPU发生故障时，光交换机重新配置，用一组健康的GPU替换一小组GPU（包括故障的那个），从而保持HB域的完整性。

A4 实验评估

实验环境

论文通过一个分析模型来评估Rail-only网络的性能，该模型模拟了使用张量并行（TP）、数据并行（DP）和流水线并行（PP）的LLM训练迭代时间，方法类似于Calculon【【27, Calculon: a methodology and tool for high-level co-design of systems and large language models, 2023, SC】】。

• 模型与数据集：评估中使用的模型包括不同规模的GPT模型（如GPT-22B, GPT-146B, GPT-1T）以及MoE-1.3B模型。这些模型的超参数和并行化策略参考了MegatronLM【【21, Efficient large-scale language model training on gpu clusters using megatron-lm, 2021】】和DeepSpeed-MoE【【22, Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale, 2022】】的公开数据。
• 硬件配置：
  • GPU：分析模型基于NVIDIA DGX A100和DGX H100/GH200平台的参数，模拟了高达65536个H100 GPU的集群规模。
  • 高带宽域（HB Domain）：HB域的大小（K）作为一个变量进行研究，从K=1到K=N（N为GPU总数）。其内部带宽也作为变量，从2.4 Tbps到9.6 Tbps。
  • 网络：NIC域的网络带宽（轨道内）从100 Gbps到400 Gbps不等。成本分析中考虑了基数（radix）为32、64和128的交换机。
• 软件与分析方法：
  • 使用基于先前工作【【28, Reducing activation recomputation in large transformer models, 2022】】提出的公式来计算每次迭代训练所需的总FLOPs。
  • 性能指标为硬件FLOPs利用率（HFU）和训练迭代时间。
  • 成本和功耗分析基于先前工作和供应商报告的交换机与收发器数据【【13, TopoOpt: Co-optimizing network topology and parallelization strategy for distributed training jobs, 2023, USENIX NSDI】】【【30, Nvidia docs hub: Mma4z00-ns400 400gb/s single-port osfp 400gb/s multimode sr4 50m connectivity scenarios, 2023, NVIDIA】】【【31, Nvidia docs hub: Qm9700/qm9790 1u ndr 400gb/s infiniband switch systems user manual specifications, 2023, NVIDIA】】。

实验结果

A. 迭代时间模型的准确性验证

论文首先通过将分析模型的硬件FLOPs利用率（HFU）估算值与已发表的文献【【28, Reducing activation recomputation in large transformer models, 2022】】中的真实结果进行比较，来验证其准确性。
- 实验内容：在一个由1到280台DGX A100服务器组成的集群上，比较不同GPT模型和集群规模下，模型计算的HFU与真实值的差异。
- 实验结果：如图7所示，对于GPT-1T模型，计算出的HFU与真实值仅相差1.8%。对于GPT-22B模型，估算误差为10.8%。误差主要来源于对GPU计算与通信的理想化建模、计算与通信重叠的假设以及对内存开销的低估。
- 分析结论：分析模型足够准确，可以用于评估Rail-only互连的训练迭代时间。

B. HB域的最佳大小

该实验旨在探究在Rail-only互连中，HB域的最佳大小是多少。
- 实验内容：在包含16384、32768和65536个H100 GPU的集群上，针对GPT-1T和GPT-146B模型，改变HB域大小（K），并绘制训练迭代时间。同时，与一个理想情况（所有GPU通过一个单一的NVLink fabric连接，即K=N）进行比较。
- 实验结果：如图8所示，随着HB域大小的增加，迭代时间减少。对于两个GPT模型，当HB域大小为256时，性能已接近最优。与理想情况相比，GPT-146B（K=256）慢4.1%，GPT-1T（K=256）慢0.9%。然而，随着HB域大小的增加，边际收益递减。
- 分析结论：HB域越大，网络开销越小。对于大型模型（GPT-1T），K=32之后收益递减。对于较小模型（GPT-146B），由于通信开销占比更大，增大HB域的收益更显著。因此，最佳HB域大小取决于GPT模型的规模。

C. HB域和网络带宽的影响

该实验分析HB域和NIC域的可用带宽对GPT-1T模型迭代时间的影响。
- 实验内容：分别针对HB域大小K=8和K=256，改变HB域带宽（2.4-9.6 Tbps）和轨道内网络带宽（100-400 Gbps），观察GPT-1T的迭代时间变化。
- 实验结果：如图9所示，增加任一带宽都会减少迭代时间。但K=8的情况对HB域带宽不太敏感（带宽翻4倍，性能提升8.0%），而对网络带宽更敏感（带宽翻4倍，性能提升35.9%）。K=256的情况则相反，对HB域带宽更敏感（性能提升13.3%），对网络带宽不敏感（性能提升8.0%）。
- 分析结论：随着未来HB域规模的增加，提高HB域的带宽比提高网络带宽更有益。

D. 批处理大小对网络设计的影响

该实验分析了全局批处理大小对训练时间和网络设计选择的影响。
- 实验内容：在一个32768个GPU的集群上训练GPT-1T模型，改变全局批处理大小（从256到4096），并观察不同HB域大小（K从8到32768）下的迭代时间和相对性能（与理想情况的比值）。
- 实验结果：如图10所示，随着批处理大小的增加，所有HB域大小的迭代时间都呈现相似的下降趋势。更重要的是，相对性能随着批处理大小的增加而提高。当K=256时，批处理大小从256增加到4096，相对性能从95%提高到99%。当K=8时，相对性能从65%提高到85%。
- 分析结论：对于HB域较小的集群，采用更大的批处理大小是可取的。由于LLM训练（特别是大型模型）本身就受益于大批量训练【【1, Language models are few-shot learners, 2020】】【【29, Scaling laws for neural language models, 2020】】，这使得Rail-only设计与LLM训练趋势完美契合。

E. 网络成本和功耗分析

该实验比较了Rail-only网络与最先进的Rail-optimized GPU集群在网络成本和功耗上的差异。
- 实验内容：计算构建两种网络架构所需的交换机和收发器数量，并根据公开数据推导网络设备成本和峰值功耗。考虑了不同集群规模和交换机基数。
- 实验结果：如表II所示，与最先进的设计相比，Rail-only设计显著降低了网络成本38%至77%（节省1.17亿至2.34亿美元），功耗降低37%至75%（节省1.7至6.9兆瓦），同时实现了同等的性能。这种节省源于消除了核心层交换机并减少了每个轨道内的交换机层级。
- 分析结论：Rail-only架构通过精简网络资源，能够为LLM训练带来巨大的成本和功耗优势。

A7 补充细节

VI. 相关工作

LLM趋势

LLM并行化与网络协同设计。随着摩尔定律放缓，LLM的计算和速度需求的增长速度超过了AI加速器和网络速度的进步，这使得超大规模集群和更高效的互连成为必需【【32, Sirius: A flat datacenter network with nanosecond optical switching, 2020, ACM SIGCOMM】】【【33, Openai: Ai and compute, 2023, OpenAI】】。MegatronLM系列工作开创了LLM并行化的先河【【21, Efficient large-scale language model training on gpu clusters using megatron-lm, 2021】】【【28, Reducing activation recomputation in large transformer models, 2022】】【【34, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020】】。本文提出的移除任意到任意网络连接的立场是对MegatronLM工作的补充。同时，我们也承认目前有工作致力于在不牺牲性能的情况下减小语言模型的大小和资源需求【【35, Hello dolly: Democratizing the magic of chatgpt with open models, 2023, Databricks】】。这些工作与我们的设计是互补的，因为我们的设计即使对于较小的语言模型和集群也能在减少网络资源的同时保持性能。同样，旨在通过量化和压缩直接减少通信量的研究方向，如DeepSpeed Zero++【【36, Deepspeed zero++: A leap in speed for llm and chat model training with 4x less communication, 2023, Microsoft】】，也对我们的方法起到了补充作用。

LLM推理

Rail-only在推理场景的潜力。本文探讨了LLM的训练工作负载，但推理是LLM产品周期的另一个重要部分。推理所需的资源较少，因为它涉及通过LLM移动的数据更少，并且只计算前向传播和多次传递以生成响应令牌【【37, AlpaServe: Statistical multiplexing with model parallelism for deep learning serving, 2023, USENIX OSDI】】。Pope等人为TPU-v4架构开发了特定的推理并行化策略【【38, Efficiently scaling transformer inference, 2022】】。对于我们的设计，每个HB域可以成为一个低延迟的推理服务域，而Rail-only连接有助于在多个推理域之间进行负载均衡。我们将对LLM推理的详细研究留作未来工作。

多作业训练

Rail-only对多作业环境的适用性。一个GPU集群同时训练多个较小的作业是很常见的。现有工作主要关注基于Clos的GPU集群，并提供了旨在实现更好公平性和更短作业完成时间的技术【【39, Gandiva: Introspective cluster scheduling for deep learning, 2018, USENIX OSDI】】-【【42, Cassini: Network-aware job scheduling in machine learning clusters, 2023】】。虽然本文专注于在大型集群上训练单个LLM，但Rail-only网络设计也适用于多作业环境。整个集群可以通过切片（tiling）的方式任意划分为更小的矩形分区，类似于TPU-v4的情况【【25, Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings, 2023】】。每个分区随后可以独立执行一个较小的训练作业。

ML基础设施和其他ML工作负载

软硬件协同设计与普适性。先前的工作已经阐明了为ML模型协同设计硬件和软件的好处。例如，谷歌的TPU集群针对在TPU上进行3D并行化的大模型训练进行了优化【【25, Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings, 2023】】，而Meta的Neo则专注于训练具有大型嵌入表的推荐模型【【43, Software-hardware co-design for fast and scalable training of deep learning recommendation models, 2023】】。我们的工作专注于设计一种能够高效训练LLM的成本效益型网络。尽管我们提出的Rail-only架构专注于LLM的网络设计，但由于转发开销较低（§IV-B），我们的设计在与其他工作结合时，对于许多其他DNN工作负载也是高效的。最近有工作试图使并行化策略和集合通信算法对任何DNN模型都具有带宽感知能力【【44, Alpa: Automating inter- and Intra-Operator parallelism for distributed deep learning, 2022, USENIX OSDI】】-【【46, Bandwidth optimal pipeline schedule for collective communication, 2023】】，这将产生非常适合Rail-only网络的流量模式。

A5 结论

本文深入研究并分析了采用混合并行化策略的LLM训练的流量模式。基于分析，我们提出了一种新颖的Rail-only网络架构，该架构与LLM的独特性质和需求高度契合。我们的架构在保持与最先进GPU网络相同性能的同时，实现了38%至77%的成本降低和37%至75%的功耗节省。