请关注公众号并设为星标，这样才能收到文章更新提示。
引言 — 本文讨论了深度学习框架中GPU组之间的集体通信操作。这些集体通信操作包括数据分发、收集和聚合等任务，对于模型训练和参数更新至关重要。文章详细介绍了不同的集体操作类型，如Reduce、AllReduce、ReduceScatter、Broadcast、AllGather和AlltoAll，以及它们如何通过组合实现更复杂的操作。特别强调了Nvidia的NCCL库在实现高效集体操作中的作用。此外，探讨了网络内卸载集体操作的概念，以及Nvidia的SHARP协议和以太网结构中网络内集体操作的发展前景，以及硬件的实现，并分别介绍了Juniper和Broadcom的芯片实现。
本文的主要参考文献见文末，相关文档已传至知识星球。
本文分为三个部分，分别介绍了：
操作的介绍
在网的卸载
硬件的实现
I. 操作的介绍
什么是集体操作（Collective Operations）？
在并行/分布式系统中进行计算时，“集体操作”指的是一组处理节点（如GPU）之间的通信操作，用于执行彼此需要协调的任务。这些任务可能包括：
将数据从一个节点分发到所有节点。
从所有节点收集数据到一个节点。
跨所有节点聚合数据
等等。
深度学习框架支持能够在GPU组之间进行集体通信的库。例如，Nvidia的集体通信库（NCCL）高效地实现了针对其GPU架构设计的集体操作。
当一个模型被部分地分到一组GPU上时，NCCL管理它们之间的所有通信。
举个栗子：
以多个模型副本的梯度聚合为例。在一个包含Nd个模型副本和每个副本包含Nm个GPU的集群中，在每次训练迭代结束时，将并行运行Nm个梯度聚合线程。
每个梯度聚合线程会聚合Nm个GPU上的所有参数的梯度。
这是一个归约操作，其中所有模型副本的每个参数的梯度都被平均得到最终梯度。用于梯度聚合的集体操作是“AllReduce”。
同样，当我们在所有模型副本中，跨同一rank的所有GPU进行参数分片时，在每次迭代之前，每个模型副本必须从同一rank的所有其他GPU中获取参数。这通过“AllGather”集体操作完成。
在训练迭代结束时，梯度被聚合并分布（或散布）到持有相应参数的GPU上。这一集体操作称为“ReduceScatter”。
ReduceScatter和AllReduce是什么区别呢？
AllReduce：归约所有数据并将结果发送回所有节点。适用于需要每个节点都持有完整归约结果的情况。
ReduceScatter：归约所有数据并将结果按块分散到各个节点。适用于参数分片等情况，每个节点只需要持有归约结果的一部分。
常用集体操作简述：
Reduce：从所有成员聚合数据（求和或平均）并将结果发送给一个成员。
AllReduce：从所有成员聚合数据并将结果发送回所有成员。
ReduceScatter：从所有成员聚合数据并将结果分散（每个成员得到结果的唯一子集）到所有成员。
Broadcast：将数据从一个成员发送到组中的所有其他成员。
AllGather：收集不同部分的数据并将其分发给所有成员。
Scatter：将一个成员的不同值分发给所有成员。
AlltoAll：将所有成员的数据分发给所有成员。
组合实现：
有些集体操作可以通过组合其他集体操作来实现。例如，AllReduce 可以通过先进行 ReduceScatter 操作，然后进行 AllGather 操作来实现：
ReduceScatter 操作首先聚合数据，然后将结果分散，这样每个成员仅持有结果的一部分。
AllGather 操作在此基础上收集并分发数据，使所有成员最终持有完整的归约结果。
这些集体操作在并行计算和分布式系统中非常重要，因为它们能够高效地进行数据传输和计算，优化系统性能。
AI/ML框架中的集体操作
NCCL（Nvidia Collective Communication Library）为集体操作实现了几种算法，如环形、树形和双二叉树，根据数据大小、GPU数量和网络拓扑选择最有效的算法。这些算法旨在优化集体操作的数据传输。
例如，在使用AllReduce集体操作进行梯度聚合时，梯度可以按环形模式从一个GPU发送到另一个GPU，每个GPU将从前一个GPU收到的梯度与本地计算的参数梯度相加，然后发送到下一个GPU。这种过程比较慢，因为梯度聚合是顺序进行的，最终结果也在环形拓扑中顺序传播回所有GPU。
在AllReduce集体操作中，GPU还可以排列成二叉树结构。每个叶节点将其存储的所有参数的梯度发送给其父节点，父节点将其与对应的兄弟叶节点的梯度相加。这个过程递归进行，直到在树的根节点处聚合所有梯度。根节点得到所有梯度的和后，需要将聚合的梯度发送回树中的所有节点，以更新它们本地的模型参数副本。根节点开始将聚合的梯度发送给其子节点，子节点再传给它们的子节点，直到所有节点都接收到梯度。
下图显示了七个GPU排列成二叉树进行梯度聚合的情况。假设这些GPU是连接到不同叶子和主干交换机的大型网络拓扑的一部分。图中还显示了GPU之间梯度聚合的流量模式。网络交换机是被动设备，在GPU之间交换结果。
环形算法
环形算法是一种常见的用于AllReduce的集体操作算法。它将数据按顺序在所有GPU之间传递，每个GPU在接收数据后与本地数据进行聚合，然后传递给下一个GPU。这种方法适用于数据量较小的情况，因为它需要较长的传输时间。
树形算法
树形算法通过将GPU排列成树形结构，使得梯度聚合过程更加高效。每个节点只需与其子节点和父节点进行通信，减少了总的通信量和时间。这种方法在处理大数据量时表现更好。
双二叉树算法
双二叉树算法结合了环形和树形算法的优点，通过在两个二叉树结构中进行数据传递和聚合，提高了数据传输效率和可靠性。
大规模深度学习模型的训练中，NCCL可以根据当前的网络拓扑和GPU数量，选择最优的集体操作算法来进行梯度聚合和参数更新，从而提高训练效率和减少通信开销。
这些集体操作和优化算法在实际应用中极大地提升了并行计算和分布式系统的性能，使得大规模模型训练和推理成为可能。
II. 在网的卸载
如果网络交换机可以帮助卸载部分或全部集体操作，会怎样呢？在这个例子中，每个交换机可以部分地对来自属于AllReduce集体操作的GPU的梯度进行求和，并将结果传递给下一个层级的交换机。最终的结果可以分层次地广播给该集体的所有GPU。
这种网络内聚合加速了梯度聚合：更少的跳数（因此延迟更低）和通过网络的流量更少。
这样做减少了集体操作的延迟和网络拥塞，同时将GPU的计算资源卸载，用于训练期间更重要的任务。同样，AllGather和Broadcast等其他集体操作也可以从网络交换机中的卸载中受益。
Nvidia在其InfiniBand和NVLink交换机中通过SHARP（可扩展层次聚合和归约协议）支持这一功能。SHARP（最新版本为v3）的主要目标是直接在网络中卸载并加速复杂的集体操作，减少需要通过网络传输的数据量，从而减少总体通信时间。SHARP是一个专有协议（不开源），仅适用于Nvidia的InfiniBand/NVLink交换机。它在Nvidia的AI堆栈中无缝工作。深度学习框架，如TensorFlow和PyTorch，通常依赖于MPI（消息传递接口）或NCCL进行集体操作。Nvidia确保这些框架可以利用SHARP以获得网络内计算的性能优势。
Nvidia没有发布比较使用和不使用SHARPv3集体操作的性能改进结果。然而，一些关于旧版本SHARP的文章显示，训练性能提高了17-20%。
以太网和在网集体操作（In-Network Collectives）
在使用ROCEv2进行GPU到GPU通信的以太网结构中，目前还没有针对网络内集体操作的开放标准。Ultra Ethernet Consortium正在开发一种新的传输协议（UET）及其相应的网络内集体操作技术。一旦标准成熟，这项技术可能会被交换机和AI框架采用。然而，从任何标准发布到在网络设备的硬件中本地部署，通常会有3-4年的滞后时间，并且需要几代硬件来优化这些功能。
III. 硬件的实现
在数据平面交换机中进行INC需要硬件支持各种操作：
解析新协议/包格式以识别数据包是否是集体操作的一部分
能够深入查看数据包（packet）并提取所有payload以进行处理
能够对payload进行操作并在多个数据包中累积结果，这需要支持各种浮点/整数格式的算术操作，而以太网交换机通常不具备这些实现
能够从内存中读取累积结果，创建新数据包并发送出去
能够处理拥塞并保持集体操作流量的服务质量（QoS）——特别是在芯片内同时进行多个并行集体操作时
能够从错误中恢复，并在集体操作失败时通知终端主机，其原因可能是由于网络或链路错误或GPU问题等。
实现集体操作的几种方法
方案一：集体操作的原生实现通过交换芯片硬件实现。但交换机添加此功能之前，业界标准必须成熟。但这些交换机如果用于更广泛的应用，这种额外需求带来的芯片面积/成本可能会在AI/ML领域外变得不具吸引力。
方案二：卸载到协处理器。将集体操作卸载到连接到交换机的协处理器。协处理器可以包含CPU核心，集体操作由在其上运行的软件实现。或者它可以是具有集体处理功能的FPGA/ASIC。这些协处理器通常处理一小部分广域网带宽以降低成本。这种方法使以太网交换机保持精简、低功耗和成本效益，并使数据中心能够选择性地在某些交换机上卸载集体操作。
方案三：具有可编程流水线处理或者包处理引擎的交换机。这种方法可以灵活地解析新包头和负载。此外，可编程交换机允许更快地实现新标准（因为需要的硬件更改很少或没有），并且可以支持多种集体操作协议。然而，这些交换机通常具有更低的带宽（5到10倍），使用这些交换机构建大型结构可能不具成本效益。但这些交换机可以作为主要高带宽交换机的协处理器来卸载集体操作和其他处理。
在网计算与Juniper的Trio架构
Trio6拥有1.6Tbps的带宽。其查找子系统包含许多数据包处理引擎（PPE）和用于专门功能（如哈希/记录、锁和过滤器）的加速器。通过高带宽交叉开关(cross bar)，PPE与加速器、数据路径和内存子系统通信。
当数据包到达时，前~200B被发送到一个空闲的PPE线程进行处理。如果需要更深入地查看数据包，PPE线程还可以读取数据包尾部的额外字节。这一特性对集体操作至关重要，因为操作数（梯度、参数）超过前200字节并占据整个数据包负载。
哈希引擎（加速器）包含多个哈希表，PPE可以在其中插入、查找或删除条目。哈希记录存储在数据内存中，并可以由PPE修改。它还支持通过定期扫描条目来进行有效期管理。
内存子系统包含用于数据包处理结构的数据内存（DataMemory）和用于带宽延迟buffer的数据包缓冲区。一个大型片上内存前端连接数据内存。片上内存分为地址空间的数据内存扩展和数据内存访问的大型缓存。它们之间的分区是可调整的。
数据包处理涉及许多读-改-写操作。有时，多个PPE可能访问同一位置（例如，计数器）以更新其内容。一个简单的方法是让每个线程完全拥有内存位置，直到其读-改-写完成。然而，这效率低下，可能大大降低性能。在Trio中，PPE将读-改-写操作卸载到内存子系统。内存子系统包含多个读-改-写引擎（每个处理特定地址子集），可以每周期处理这些请求（8字节）。当一个特定引擎收到对同一内存位置的多个请求时，它按顺序处理这些请求，保证更新的一致性。混合读取、写入和读-改-写操作时，无需显式一致性命令。“修改”操作支持各种位操作和两个32位整数加法。
该架构还允许通过读取数据内存的内容、使用PPE处理附加适当的头信息、将新数据包写入数据包缓冲区，并将其排队发送来创建新数据包。因此，Trio架构具有实现网络内集体操作所需的所有配置。
该架构这些处理在交换机中没有硬编码（hardcoded）。它具有支持和解析任何新协议（UEC或自定义协议）的灵活性。SIGCOMM论文使用网络内聚合（AllReduce）来展示Trio PFE在集体操作方面的灵活性。
网络内聚合流程
首先，需要定义集体数据包的头部（通常在UDP层下）和数据包格式。在SIGCOMM论文中，作者创建了一个自定义头部，称为Trio-ML头部。该头部包含足够的信息来识别聚合线程、源/目标GPU、梯度块和块中梯度的数量。ML框架通常允许插件支持自定义通信协议。
在训练开始前，在作业配置时，在哈希表中创建一条作业记录，并在作业完成之前保持不变。这些记录包含所有参与作业的源（GPU）、要聚合的块（梯度子集）的数量以及如何创建响应数据包的信息（包括数据包的目标等）。
当作业/块的第一个数据包到达时，PPE线程在哈希表中创建一个块记录。它跟踪集体操作的状态（在本例中为梯度聚合），包括尚未传送梯度的源以及数据内存中聚合缓冲区的指针。块聚合完成后，块记录被移除。
当数据包到达PPE线程时，它解析头部并在哈希表中查找作业/块ID。
如果不存在块记录，它会创建记录并为该块在数据内存中分配一个聚合缓冲区。
如果块记录已经存在，则进行读-改-写操作，将当前数据包的梯度聚合到聚合缓冲区中。
如果这是需要聚合梯度的最后一个源（块记录指示所有源的状态），则生成响应数据包，将其写入数据包缓冲区内存，并在排队子系统中排队。之后清理块记录。
在训练之后，系统（在控制平面中）可以清理作业记录以释放分配的空间。
在一个小型设置（使用六个A100 GPU服务器和用于图像识别的小型DNN模型）中，结果显示训练性能提高了1.8倍。这是一个概念验证（POC）设计。虽然这篇论文专门讨论了梯度聚合，但同样的工作流程可以支持其他数据缩减或数据收集的集体操作。
尽管使用数据包处理引擎（或网络处理器）的网络交换机在实现交换机内部的集体操作方面提供了最大的灵活性，但这些交换机不像使用固定流水线数据包处理的浅/深缓冲交换机那样高的端口密度。
此外，由于数据包处理预算有限，网络内集体操作花费的时间越多，数据包处理性能越低（这会增加满足线路速率所需的最小数据包大小）。
克服这一点的一种方法是将这些可编程交换机添加为连接到常规高带宽交换机的协处理器，并使用这些交换机的一部分广域网带宽与协处理器进行集体操作通信。
在使用可编程交换机构建支持网络内集体操作的结构时，没有一种方案适合所有情况。这取决于集群大小、正在训练的模型类型以及将要卸载的集体操作。
网络内集体操作的其他考虑因素
网络拥塞和流量控制：现有的拥塞通知机制仍然可以用来标记结果数据包的ECN标记。当接收主机接收到网络内部生成的ECN标记数据包时，它知道所有参与的源主机，并可以将拥塞信息传回发送主机。
内存需求：交换机的片外内存有限。这导致了同时运行多少集体线程的额外限制。此外，收集数据的集体操作（没有缩减）可能会进一步加剧内存容量压力。
在多个交换机上分散集体操作：每个叶交换机执行部分集体操作，并将部分结果发送到下游交换机（通常是主干交换机），可以减少每个交换机的内存需求、网络热点并均衡负载。通过优先处理关键的聚合流量并应用适当的QoS策略，网络可以确保即使在高需求期间关键数据流畅通。
错误处理/超时：如果网络交换机未接收到所有源的梯度，应具有超时机制来结束集体操作，并将部分结果或错误指示发送回所有源服务器。在Trio示例中，可以通过扫描哈希表并在预定义超时时间内未收到新流量时使条目老化来实现。
标准编程语言/互操作性：Juniper Networks的高级转发接口（AFI）通过允许第三方开发人员控制和管理一小部分转发路径图的虚拟容器（称为沙盒）来提供设备的部分可编程性。沙盒允许开发人员添加、移除和更改特定数据包操作的顺序。
更广泛地在网络设备中采用在网计算要求这些设备支持与P4编程语言兼容的可编程数据平面。Trio可以支持P4原语（正在开发中）。使用P4，开发人员可以在非常细粒度的级别上指定网络设备如何处理数据包，包括解析、修改和根据数据包内容做出转发决策。这种细粒度的控制对于在网络内部实现各种集体操作是必要的。P4还允许不同供应商之间的互操作性。然而，不同供应商具有不同的硬件架构和行为，实现真正的互操作性很难。P4在一个设备上编译的内容可能与另一个设备上具有不同硬件资源的设备不同。
在网计算与Broadcom的Jericho3-AI芯片
Tomahawk5在某些方面比Jericho3-AI更强大，但是它缓冲区要小得多，适合于hyper scale和CSP在CLOS网络中做spine，leaf的交换机。它有512个100Gbps的Serdes，以支持51.2Tbps的总带宽。
而Jericho3-AI只有304个同样的Serdes，其中144个下行链路，160个上行链路连接到更高的Ramon3，以形成spine，leaf架构，如下图：
这两个芯片的SerDes是一样的，都为Peregrine SerDes，可以直接驱动光模块的放大器，而不需要在前面加DSP，即支持LDO。此外，它可以驱动信号支持4米长的DAC线缆，而不需要retimer。Broadcom的SerDes真的强，不知道怎么做的。
Jericho3-AI支持AI网络中的集约操作带来的复杂流量处理，特别是all-to-all和all reduce操作，这两个操作存在于每一步分布式计算的最后一步。
这是LLM和推荐系统中重要的功能。
接下来看一个InfiniBand和Jericho3的以太网的对比。
我们知道NCCL是拓扑感知的，它可以知道高速高带宽的NVLINK和跨机器之间的IB/Ethernet之间的区别。
下图是InfiniBand和Jericho3-AI之间throughput的对比。
IV. 写在最后
通过Juniper的Trio架构和Broadcom的Jericho3-AI芯片的例子，文章展示了在网计算的灵活性和高效性，以及它们如何支持复杂的集体操作。这些技术的发展预示着未来在网络设备中实现更智能、更高效的数据处理的可能性。
参考文献：
1. Sharada Yeluri：
In Network Acceleration for AI/ML Workloads
https://www.linkedin.com/pulse/network-acceleration-genai-workloads-sharada-yeluri-g8xbc/?trackingId=SP2TC1lSS0eMtaCS8A2Rlw%3D%3D
2. NCCL 2.21 Collectives：
https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html
高阅读量文章
