可编程交换机如何卸载 AI 训练中的集体操作？.docxVIP

在本文中，Juniper技术专家兼高级工程总监SharadaYeluri深入研究了AI/ML训练/推理中使用的集体操作，并讨论如何将其中一些功能卸载到网络交换机上，以减少拥塞并提高结构的性能。最后，SharadaYeluri以Juniper的?Trio?架构为例，展示了可编程交换机如何无缝卸载集体操作。

什么是“集体操作”？

在由众多通过互联结构相连的处理节点（如GPU）构成的并行/分布式计算系统中，“集体操作”是指涉及一组处理节点间通信的一系列操作，用于执行协调性的任务。这些任务可能包括将数据从一个节点分发给所有节点、将所有节点的数据收集到一个节点、聚合所有节点之间的数据等。?

深度学习框架支持库，可实现GPU组之间的集体通信。Nvidia的集体通信库（NCCL）针对其GPU架构高效地实现了集体操作。当一个模型在一组GPU之间进行分区时，NCCL负责管理它们之间的所有通信。

??

下面是常用的集体操作：

Reduce：从所有节点聚合数据（求和或平均）并将结果发送给其中一个节点。?

AllReduce：聚合所有节点的数据，并将结果发送回所有节点。

AllReduce

ReduceScatter：聚合所有节点的数据，并将结果（每个节点获得结果的唯一子集）分发到所有节点。

ReduceScatter

广播：将数据从一个节点发送到组中的所有其他节点。

广播

AllGather：收集数据的不同部分，并将其分发给所有节点。?

Scatter：将一个节点的不同值分发到所有节点。?

AlltoAll：将所有节点的数据分发到所有节点。

有些集体操作可以通过使用集体操作的组合来实现。例如，AllReduce可以实现为先执行ReduceScatter操作，然后再执行AllGather操作。??

??

AI/ML框架中的集体操作

NCCL实现了多种集体操作算法，包括环形（Ring）、树形（Tree）以及双二叉树（DoubleBinaryTrees）等，根据数据大小、涉及的GPU数量和网络拓扑自动选择最为高效的算法。这些算法旨在优化集体操作中的数据传输。?

例如，在使用AllReduce集体操作进行梯度聚合时，梯度可以按照环形模式从一个GPU发送到另一个GPU，其中每个GPU都会将其从上一个GPU接收到的梯度与其本地计算出的梯度进行聚合，然后再将结果发送给下一个GPU。这个过程很慢，因为梯度聚合是按顺序完成的，最终的结果也会按顺序在环形拓扑中传回到所有GPU。?

在AllReduce集体操作中，GPU也可以排列为二叉树结构。在这个结构中，每个叶节点将其存储的所有参数的梯度发送到其父节点，父节点则会将接收到的梯度与来自其兄弟叶节点的对应梯度进行求和。此过程以递归方式继续，直到所有梯度在树的根节点处聚合。在根节点拥有所有梯度的总和后，必须将聚合梯度发送回该树中的所有节点，以更新其模型参数的本地副本。根节点首先将聚合梯度发送给其子节点，子节点再依次传递给它们的子节点，如此递归进行，直到所有节点都收到了更新后的梯度。

下图展示了7个GPU以二叉树的形式排列以进行梯度聚合。假设这些GPU是连接到不同叶交换机和主干交换机的大型网络拓扑的一部分。该图还显示了GPU之间梯度聚合的流量模式。在网络中，这些交换机是被动设备，负责转发GPU之间的通信结果。

梯度聚合流量（仅显示4个GPU）

集体操作的卸载

如果网络交换机可以帮助卸载部分或全部的集体操作，结果会怎样？在这个例子中，每个交换机可以对从属于AllReduce集体操作的GPU接收到的梯度进行部分求和，并将结果传递给下一层级的交换机。最终的结果将广播给该集体内的所有GPU。??

上图展示了网内聚合如何加速梯度聚合过程：通过减少数据跳转次数（从而降低延迟）和减轻网络流量。

这样做可以减少此集体操作的延迟和网络拥塞。它还可以在训练期间卸载GPU计算资源以专注于更重要的任务。同样，AllGather和Broadcast这两种集体操作也能从在网络交换机中卸载任务中获益。

Nvidia在其?InfiniBand?和NVLink交换机中通过SHARP（可扩展层次聚合和归约协议）支持此功能。SHARP的主要目标是在网络内部直接卸载并加速复杂的集体操作，减少需要在网络上传输的数据量，从而降低整体的通信时间。SHARP是一个专有的协议，仅与Nvidia的InfiniBand/NVLink交换机兼容。它在Nvidia的AI堆栈中无缝运作。

深度学习框架，如TensorFlow和PyTorch，通常依赖于MPI（消息传递接口）或NCCL来进行集体操作。Nvidia确保这些框架能够利用SHARP来获得网络内计算的性能优势。然而，Nvidia并未公开发布SHARP