摘要
我们提出了一种新的网内聚合框架 PANAMA，用于在共享集群实现分布式机器学习(ML)训练，可服务于各种工作负载。PANAMA 包含两个关键组件:(1)定制的网内硬件加速器，可以在不影响精度的前提下支持线性浮点梯度聚合; 以及(2)轻量级负载均衡和拥塞控制协议，该协议利用 ML 数据并行作业的独特通信模式，使不同作业之间公平共享网络资源，同时保障长期作业的高吞吐率以及短期作业和其他时延敏感作业的低时延。我们在 10 Gbps 收发器和大规模模拟的基础上创建了基于 FPGA 的原型，从而评估 PANAMA 的可行性。模拟结果表明，PANAMA 能够将大型工作的平均训练时间减少 1.34 倍。更重要的是，通过大幅降低大型数据并行作业在网络上的负载，非聚合流量(特别是对延迟敏感的短期流量)也能获得明显的好处，其 99%的完成时间提升了 4.5 倍。
1. 简介
对精确机器学习(ML)模型不断增长的需求导致了深度神经网络(DNN)数据集和模型规模的稳步增长。因此，训练现代 DNN 模型远远超出了单个设备的能力，今天训练大型模型需要数千个加速器(Huang et al.，2018; Lepikhin et al.，2020; Moritz et al.，2018; Sun et al.，2019)。
一些学术机构和公司最近提倡使用网内聚合(Costa et al.，2012; Mai et al，2014)，以提高分布式数据并行 ML 训练工作负载的性能(Bloch, 2019; Klenk et al.，2020; Li et al.，2019; Sapio et al.，2021)。通过在网络交换机内部而不是在终端主机上聚合梯度，可以缓解数据并行训练的通信瓶颈，减少训练时间。
然而，现有建议侧重于相对简单的场景，将网内聚合限制在单个交换机(Lao et al.，2021; Sapio et al.，2021)、少量交换机(Li et al.，2019)或单个 ML 作业上。但在实践中，随着神经网络模型的普及和多样性的增长，为了减少成本和资源浪费，第一方和第三方训练工作负载都越来越多的向共享集群过渡(Abadi et al.，2016; Amazon, 2021; Azure, 2021; Google,2021; Jeon et al.，2019)，跨越数百个机架并执行大量不同的 ML 作业。这些工作包括传统 ML 工作，如 K-means 聚类，以及最近的强化学习和深度学习等。从网络角度来看，这意味着有大量异构流，从几个 KB 到几十 GB 大小不等(Abadi et al.，2016; Azure, 2021; Google, 2021; Jeon et al.，2019)，其中只有一小部分可能需要聚合。如果没有合适的机制来大规模有效共享网络资源，在更现实的情况下，网络内聚合的实际可行性就会变得令人怀疑，这一点在某些早期部署的情况下得到了证实: 例如，当在更大规模测试网络内聚合时，NVIDIA 观察到由于拥塞造成吞吐量低于预期(Klenk et al.，2020)。
我们在本文中提出 PANAMA(ProgrAmmable Network Architecture for ML Applications, ML 应用程序可编程网络体系架构)来解决这一问题，这是一个为具有异构工作负载的共享环境量身定制的网络内聚合框架。PANAMA 由两个主要部分组成。首先是一种全新设计的聚合硬件加速器，可以同时支持多个活跃的训练工作。现有的基于可配置匹配表(RMT, Reconfigurable Match Table)架构的可编程交换机(Bosshart et al.，2013)，例如 Tofino 交换机(Intel，2018)可用于网内聚合(Lao et al.，2021; Sapio et al，2021)，但缺乏浮点运算支持，并且流水线架构缺乏灵活性(Gebara et al.，2020)，从而造成准确性和性能下降。相比之下，我们选择了一种串连模式(bump-in-the-wire) 方法，在这种方法中，网内加速器与交换机解耦。该设计的灵感来自于最近在云上部署的可编程网卡(NIC)(Firestone et al.，2018)，在不牺牲训练精度并且不要求对交换机逻辑进行任何更改的情况下提供了最大的灵活性。我们的设计可以在线速(10-100 Gbps)下运行，并能够有效利用逻辑区域。
PANAMA 的第二个组成部分是与轻量级拥塞控制协议相结合的负载均衡机制，允许在不同作业之间高效、公平的共享网络带宽。PANAMA 将聚合流量分布在多个聚合树中，从而减少网络热点，提高网络性能。现有数据中心拥塞控制协议(Alizadeh et al.，2010; Kumar et al.，2020; Mittal et al.，2015; Zhu et al.，2015)不适合网内操作，因为需要假设服务器之间是点对点连接，而不是基于树的配置。相比之下，我们的新型拥塞控制协议利用了网内聚合的独特特性(和机会)来提高性能，同时确保加速器上有限的缓冲区空间被 ML 作业公平共享，而不会溢出。此外，我们的拥塞控制算法兼容现有数据中心使用的基于 ECN 的协议(Alizadeh et al.，2010; Zhu et al.，2015)，从而能够与其他非网内聚合流量公平共享网络资源。
我们的工作揭示了以前没有注意到的网内聚合的好处，研究界怀疑网内聚合的实际价值，值得注意的是，梯度聚合只构成整个训练任务的一小部分，因此对整体训练时间的影响有限。我们认为，也许与直觉相反，网内聚合的真正动机与其说是为了提高训练工作本身的性能，不如说是为了减少数据并行梯度交换产生的通信量。将网络中的流量聚合在一起可以大幅降低网络使用量(单个数据并行作业在执行期间可以生成超过 1 PB 的数据)，从而为其他流(包括非数据并行的 ML 作业和延迟敏感流)释放网络资源。
我们基于 FPGA 的原型演示了硬件加速器的可行性，并提供了广泛的模拟分析，以评估大规模的 PANAMA 负载均衡和拥塞控制的好处。结果表明，PANAMA 缩短了 99%时延敏感型任务的完成时间，提升幅度最高达 4.5 倍(平均为 2 倍)，同时减少 ML 平均训练时间最高达 1.34 倍。
2. 网内聚合示例
最常见的分布式训练方法之一是数据并行训练，其中神经网络被复制到 N 个工作节点(或副本)上，每个工作节点处理一小部分训练数据(小批量)来计算其局部模型梯度。在每次迭代中，工作节点必须通过交换和聚合梯度来同步各自的模型，以确保收敛(Narayanan et al.，2019)。这个步骤被称为 allreduce，可以使用参数服务器(Li et al.， 2014)或 Ring-AllReduce (Sergeev & Balso, 2018; Uber Eng., 2017)。
allreduce 步骤对网络架构造成了很大的压力，因为在整个训练过程中，需要多次交换整组模型梯度。例如，对于有 1000 个工作节点的训练工作，需要 1000 次迭代才能达到目标精度的 1 GB DNN 模型，产生的总流量大约为 2 PB，其中包括发送梯度，以及接收聚合值。最近的提议提倡使用网内聚合来提高 ML 训练工作负载的性能(Bloch, 2019; Klenk et al.，2020; Lao et al.，2021 年; Li et al.，2019; Sapio et al., 2021)。直觉上，在 allreduce 步骤中减少网络上传输的数据量可以缩短通信时间，从而让训练更快。
然而，分析结果显示，网内聚合导致的整体训练时间改进只有 1.01-1.8 倍(Sapio et al.， 2021 中的表 2)。这有两个原因。首先，计算时间占据了整体训练时间的很大一部分。因此，只有一小部分时间花在通信上，从而限制了可达到的最大收益。其次，即使通信时间包含了整个训练时间，与最先进的 Ring-AllReduce 策略相比，网内聚合理论上可以实现最大的改进是在大消息通信的情况下，通信时间减少 2 倍(Klenk et al.，2020)。
为了验证这一观察结果，我们使用第三代英伟达 GPU(Pascal P100、Volta V100 和最近推出的 Ampere A100)来训练五个流行的图像分类模型。选择这些 DNN 模型是因为它们涵盖了广泛的规模和计算要求。我们使用 TensorFlow 框架和 ImageNet 数据集来训练这些模型，并在 TensorFlow 基准套件(TensorFlow, 2021)中指定了批量大小。
为了避免分布式训练框架或低效的网络栈造成的任何测量偏差，我们在单 GPU 节点上训练模型来测量每次迭代的计算时间，并将单次迭代中交换的梯度大小除以链路带宽来估计通信时间。为了在不同 GPU 之间保持一致，假设 P100、V100 和 A100 的链路带宽分别为 10Gbps、40Gbps 和 100Gbps。
图 1. 单个网内聚合作业的预期加速受限于通信时间与总训练时间的比率。
图 1a 显示了不同 DNN、GPU 和网络带宽之间的通信时间与总训练时间(通信+计算)的比率。该图显示，用于通信的训练时间比例在 0.11-0.70 之间，随着网络速度的提高，这一比例也在降低。这表明网内聚合对数据并行作业的好处也在随着时间的推移而减少。图 1b 说明了这一点，它显示了网内聚合对具有 100Gbps 链接的 A100 GPU 的预期训练时间加速在 1.06-1.28 之间(而对具有 10Gbps 链接的 P100 来说是 1.15-1.53)。事实上，我们假设网内聚合的最佳情况是梯度计算和通信之间没有重叠。这是非常保守的，因为现代分布式 ML 框架利用重叠来最小化通信对训练时间的影响，因此，我们预计减少通信的好处将不那么明显。
相反，在本文中，我们认为网内聚合的真正机会在于，通过减少注入网络的整体数据并行流量来提高共享集群中非数据并行作业的性能，从而为其他流量释放出网络带宽。如图 2 所示，共享的 ML 集群由一组异质流量组成，其传输速率从几 KB 到几百 GB 不等。图的左边部分(蓝条)显示了包含模型梯度的并行数据流的传输大小，我们称之为聚合流，以表明它们是网内聚合的良好候选者。图的右边(绿条)显示了不适合网内聚合的流量，即非聚合流，包括不需要任何聚合的流量，例如，(a)数据集传输和(b)由并行流水线(Huang et al.，2018; Narayanan et al.，2019)或并行模型(Shoeybi et al., 2020)产生的流量，以及需要聚合但网内聚合没有明显好处的短流，例如(c)由强化学习训练(RL)产生的流(Li et al.，2019; Moritz et al.，2018)和(d)更传统的 ML 作业，如 k-means 聚类(Rossi & Durut，2011)。
图 2. 共享 ML 集群中的数据传输大小。
我们发现，虽然非聚合流不适合网内聚合，但可以间接的从网内聚合中受益。然而，现有的网内聚合解决方案并不是为在共享环境中运行而设计的。正如我们在第 7 节中所展示的，现在的数据中心解决方案只支持聚合和非聚合流的共享，但没有适当的负载均衡和拥塞控制机制来应对多租户环境。相比之下，PANAMA 结合了对多个 ML 作业的原生硬件支持以及为网内聚合量身定制的拥塞控制协议，即使在大量聚合作业的情况下，也能接近非聚合流的理想性能，同时，减少聚合作业本身的训练时间。
3. PANAMA 概述
本节将提供 PANAMA 的概要介绍，接下来在第 4 节和第 5 节中将详细介绍其关键组成部分。我们假设传统数据中心网络是类似于图 3 那样的多层 CLOS 拓扑结构(Al-Fares et al.，2008; Singh et al.，2015)。PANAMA 集群中的每个交换机(PSwitch)包含一个传统交换机，例如 Broadcom Tomahawk(Broadcom，2020)，与串连聚合加速器相连。
① 工作节点安置。 当提交新的 ML 训练作业时，数据中心调度器确定最佳的分布式训练策略(Jia et al.，2019; Narayanan et al.，2019; Sergeev & Balso，2018），并将作业实例化在一组直接运行在集群服务器或虚拟机上的工作节点上。这些工作节点可能在同一个机架上，也可能分布在多个机架上，PANAMA 对工作节点的安置不做任何假设。网内聚合的流量(聚合流量)的选择基于运营商的偏好，可以在控制器逻辑中配置。在我们的实验中，我们将所有大小超过 40MB 的 reduce 数据并行流标记为聚合流量。这个阈值是我们根据经验选择的，因为在较短的流量中看到的好处几乎可以忽略不计。
图 3. PANAMA 高层工作流。
② PSwitch 的选择和初始化。 当 PANAMA 控制器为一项工作选择网内聚合时，需要初始化连接工作节点所有生成树的 PSwitch。如果所有工作节点都在同一个机架上，这些树只是包含单一架顶(ToR)交换机，如果工作节点分布在多个机架上，则可能包括多层交换(见图 3)。PANAMA 利用现代数据中心的多路径连接，将单一工作的聚合流量分布在多个树上，从而提高网络效率，减少拥堵。
PANAMA 控制器为每个聚合树生成唯一的 IP 组播地址，并沿着树的路径在 PSwitch 内配置转发条目。在上游方向，PSwitche 被配置为向聚合树根部转发聚合数据包，而在下游方向，聚合数据包使用本地 IP 组播支持(如果可用)或通过在转发表中添加单个条目被路由回工作节点。PANAMA 控制器还在路径上用作业状态和执行聚合所需的信息初始化网内聚合加速器(第 5 节)。
③ 工作节点设置。 最后，PANAMA 控制器将选定的工作节点配置为使用网内聚合，并通知他们选定聚合树的 IP 组播地址。PANAMA 不需要服务器上任何特定的硬件支持，PANAMA 通信库可以取代主流 ML 框架所使用的通信库，例如 NCCL 和 MPI(NVIDIA, 2021; The Open MPI Project, 2021）。PANAMA 库将梯度封装成加速器支持的数据包格式(见图 4)。
图 4. PANAMA 聚合包格式。
4. 网络设计
本节我们将介绍 PANAMA 如何解决在一个由聚合流和非聚合流组成的共享环境中运行网内聚合的挑战(见图 2)。这需要解决两个问题。首先是如何在多层数据中心中平衡聚合流量，以有效利用网络资源，并尽量减少非聚合流量的拥堵。PANAMA 通过利用多个聚合树来解决这个问题。第二个问题是如何在聚合流和非聚合流之间公平分享网络带宽，而不会造成硬件加速器缓冲区溢出。为此，我们提出了一种基于 ECN 的拥塞控制算法，对聚合包和整个结构的 PAUSE 帧设置了拥塞窗口上限，以确保即使在拥堵的情况下，也能保证确定性的无损运行。
4.1. 路由和负载均衡
现在数据中心网络通常依靠等价多路径(ECMP, Equal Cost Multi Path)协议实现网络负载均衡，并确保属于某个 TCP 流的所有数据包通过相同的路径进行路由，从而避免接收方失序(Hopps，2000)。由于大多数数据中心的流量都很短(Alizadeh et al.，2010; 2013; Greenberg et al.，2009)，ECMP 通常足以确保流量在网络中合理分布。然而，聚合流量通常非常大，按照 ECMP 的要求，将这些流量限制在单一路径(聚合树)上，会造成严重的网络不平衡。这对争夺带宽的延迟敏感的短流量特别有害，它们将面临更多的排队延迟。为了避免这种情况，PANAMA 在每个训练任务中使用多个聚合树，将流量分散到多个路径上，避免出现拥堵热点。
如前所述，PANAMA 控制器向工作节点提供一组 IP 组播地址，代表为一项工作选择的聚合树，工作节点以轮询的方式将梯度包分配到不同的树上。例如，在图 3 的拓扑结构中，有四个聚合树AggTreei,i=1,...,4和 8 个 ID 为pj,j=1,...,8的聚合包，假设工作节点开始向每个树发送单个数据包，我们的协议通过如下方式实现负载均衡: p1,p5→AggTree1;p2,p6→AggTree2;p3,p7→AggTree3;p4,p8→AggTree4，其中→表示每组数据包的目的地聚合树。例如，p1 和 p5 被发送到AggTree1。注意，数据包的编号也反映了发送顺序，由于轮流发送，p5 只有在前面的数据包被发送到其他树之后才被发送。这种机制平衡了各树的负载，并且由于顺序是确定的，因此不需要工作节点之间的协调。非聚合流量不受此影响，使用运营商定义的负载均衡协议进行转发。
4.2. 拥塞控制
本节我们将介绍 PANAMA 拥塞控制协议，首先概述聚合流量的拥塞控制协议需求，然后解释我们提出的协议如何满足这些需求。
4.2.1. 需求
R1 对多点通信的支持。 传统拥塞控制协议假定在 (源、目的) 对之间进行单播、点对点通信。与此相反，网络内聚合涉及不同实体之间的多对多通信。因此，典型的基于丢包或拥塞通知的限速机制并不直接适用。一个天真的解决方案是在 PSwitch 中实施拥塞控制，用逐跳流量序列取代树状的多点流量。然而，这将浪费宝贵的芯片封装面积。
R2 小型缓存。 由于硬件加速器需要在数百个端口上以线速运行，我们不能依靠 DRAM 之类的外部存储器，而只能使用片上缓存。片上存储器消耗了大量芯片面积，因此限制片上存储器的大小并通过在多个工作中公平分享有限的存储器并有效使用就显得至关重要。与传统网络的一个关键区别是，不仅因为拥堵需要缓存，而且还需要存储聚合包，直到来自所有工作的聚合包到达交换机。这就引入了流量之间的依赖性。由于只有在收到所有数据包后才能计算出结果，工作节点的发送速度必须匹配，否则，较快发送端的数据包必须被缓存，直到收到最慢发送端的数据包，从而浪费了可用于其他 ML 工作的资源。这一属性将网内聚合与表面上类似的 co-flow 抽象区分开来（Chowdhury & Stoica，2012），在后者中，单个流可以使用不同的速率。
R3 与传统协议的兼容性。 协议的一个关键需求是能够与主流拥塞控制协议共存。特别是，该协议应该是 TCP 友好的，因为 TCP 是数据中心事实上的标准拥塞控制协议。在交换机上使用加权公平队列(WFQ, Weighted Fair Queues)来分离聚合流和非聚合流可能看起来是个简单的解决方案，但事实并非如此，它不能在属于不同工作的聚合流之间提供公平的共享，而且还涉及到动态选择分配给两个流量类别中每一个的最佳权重的复杂任务。
R4 无损操作。 与传统 TCP 点对点流量不同，如果 PANAMA 中的某个聚合数据包丢失，就需要重新传输几个数据包，这可能会大大降低整个吞吐量。因此，即使在高网络负载下，确保缓冲区永不溢出也是至关重要的。
图 5. PANAMA 拥塞控制算法。
4.2.2. 设计
接下来我们讨论 PANAMA 拥塞控制协议的关键特性，并解释如何满足上述需求。我们在图 5 中提供了协议的伪代码，该协议是终端主机上的 PANAMA 通信库的一部分，使用这一协议不需要对训练框架进行任何更改。
隐式确认。 现有数据中心拥塞控制机制利用网络交换机(如丢包或 ECN 标记)和终端主机(如 RTT)的信号来检测拥堵的开始并调整发送方的数据包发送速率(Alizadeh et al.，2010; Dong et al.，2015; Ha et al.，2008; Handley et al.，2017; Mittal et al.，2015; Zhu et al.，2015)。然而，网内聚合不能复用这种机制，因为新的数据包是在网络内的每个交换机上通过将几个传入的数据包聚合成一个来构建的。这破坏了数据包和对应的确认消息之间的一对一映射。相反，我们的设计利用了网内聚合操作的独特属性。由于聚合结果的数量等于本地计算梯度的数量，每个工作节点都期望为每个发送的聚合数据包提供一个结果包。因此，工作节点可以将聚合结果包视为隐含的确认信号，以增加窗口大小，如图 5 所示。这就克服了在 PSwitch(R1)上维持每个流量拥堵状态的需要。此外，拥塞控制在每个聚合树上单独运行，避免了在多条路径上重新排序数据包的需要，当与我们的负载均衡协议(4.1)相结合时，可以提供多路径拥塞控制协议的好处(Peng et al.，2013)，而没有额外的复杂性。
ECN 标记。 我们的拥塞控制协议受到 DCTCP(Alizadeh et al.，2010)的启发，依靠 IP 头中的 ECN 标记对观察到的网络拥堵做出反应。在 PANAMA 中，我们扩展了这一机制，使聚合工作速率在工作节点之间同步。网内聚合的一个明显特征是，当数据包在聚合树中向上移动时，PSwitch 必须在多个数据包中聚合梯度并产生聚合结果数据包，而这可能会导致丢失网络拥堵状态信息。然而，在 PANAMA 中，PSwitch 内的聚合加速器保留了聚合数据包的 ECN 字段信息，每个硬件加速器对收到的数据包的 ECN 字段进行比特或(OR)操作，将 ECN 位镜像到生成的聚合数据包的 IP 头中(见图 4)。因此，聚合数据包将携带 ECN 位返回给所有的工作节点。与传统的基于 ECN 的拥塞控制方案不同，因为结果包被用作隐式确认，因此不需要将 ECN 回传给发送端。工作节点检查结果包，如果 ECN 位被设置，就会调整发送速率，详见图 5。这种机制确保每个聚合树的拥塞窗口在聚合树(R2)中的工作节点之间以同步的方式增长和缩小。此外，由于拥塞控制机制与 DCTCP 相匹配，可以保证与现有的传统协议兼容，正如我们在第 7 节的评估中显示的那样(R3)。
拥塞窗口的上限。 为了避免由于加速器缓冲区溢出而导致数据包丢失，PANAMA 的拥塞协议对训练工作节点的拥塞窗口大小设置了上限，以匹配每个聚合树加速器中的最小可用缓冲空间，从而确保总是可以容纳传入的数据包，而不会因为缺乏可用的缓冲空间而被丢弃(R4)。为了维护最新的可用缓冲空间，硬件加速器使用一个称为cwnd_cap的字段更新每个聚合数据包。数据包中为cwnd_cap保留了 16 位，以便在聚合数据包上升到聚合树根部时，捕捉到加速器上存储数据包的最小可用内存。每个加速器根据活跃训练作业的数量(packet_memory/num_jobs)计算其可用内存，如果可用缓冲空间小于收到的聚合数据包的cwnd_cap，则覆盖cwnd_cap，否则保留最小cwnd_cap值。然后，最终值与梯度聚合结果一起被发送给所有工作节点。如上所述，收到聚合数据包被视为确认信号，使工作节点能够发送下一组数据包。类似于 TCP 协商窗口大小(数据包在工作节点创建时被假定为具有固定大小)，cwnd_cap的值被用作每个工作节点当前最大可处理数据包数量的上限。我们依靠标准以太网流量控制机制，通过 PAUSE 帧确保网络内的加速器不会造成交换机缓冲区溢出，从而形成端到端无损架构。在数据包损坏或失败的情况下，仍然可能有丢包(Zhuo et al.，2017)，但可以使用简单的超时机制来处理这个问题。由于我们协议的无损属性，超时值的设置不需要太激进，从而防止虚假重传。
