集合通信的实现

动机：实现 AllReduce

既然我们已经定义了 7 种集合通信操作，接下来可以思考如何在网络中实现它们。要实现一种集合通信操作，我们需要回答两个问题：使用什么拓扑结构连接节点？为了高效完成操作，节点之间需要交换哪些数据？

一旦确定了使用的拓扑结构和需要交换的数据，我们就可以分析设计的性能。我们使用的总网络带宽是多少？操作完成需要多长时间？还可以关注其他性能指标，但在本笔记中我们主要关注这两个。

为了衡量性能，我们定义一些变量。总共有 p 个节点。每个向量的总大小为 D 字节。这意味着每个向量元素（即图中的每个方框）的大小为 D/p 字节。

在本节中，我们将 p 设为 5，以便让一些演示更直观。注意，这也意味着每个向量现在有 5 个元素，而不是 4 个。（补充说明：请记住，向量代表任意数据，我们将每个向量分成 p 个大小相等的子向量，其中 p 是节点总数。p 从 4 增加到 5 并不一定意味着数据更多，可能只是将相同的数据分成 5 块而不是 4 块。）

在本节中，我们将重点介绍 AllReduce 集合通信的实现，不过这些思路也适用于其他集合通信操作。回想一下，AllReduce 会计算向量的按元素求和，然后将求和后的向量发送给所有节点。

方法 1：全连接 Mesh 网络

我们考虑的第一种拓扑结构是全连接 mesh 网络，其中每个节点都与其他所有节点有直接链路。

在这种拓扑结构下，我们可以通过以下步骤实现 AllReduce：首先，每个节点将其整个向量直接发送给其他所有节点。

然后，每个节点对收到的所有向量进行求和。

这种方法使用多少带宽？每个节点需要将其整个向量（D 字节）发送给其他 p-1 个节点，因此每个节点发送 D (p-1) 字节。总共有 p 个节点，因此发送的总数据量为 Dp (p-1)=O (D・p²) 字节。

这种方法需要多长时间？这取决于节点和链路的确切资源限制，但假设没有资源限制，所有向量的发送都可以同时进行，在一个时间步内完成。也就是说，节点 1 使用其所有 3 条 outgoing 链路同时向其他所有节点发送数据。同时，节点 2 也可以使用其所有 3 条 outgoing 链路同时向其他所有节点发送数据。假设没有资源限制，这种方法需要一个时间步完成，其中每个节点在每个时间步需要发送和接收 2・D・(p-1) 字节（每个节点发送 D・(p-1) 字节，接收 D・(p-1) 字节，两者相加得到额外的系数 2）。

方法 2：在单个节点进行归约

在接下来的拓扑结构中，我们让单个节点承担所有计算工作：

要运行 AllReduce：首先，除节点 1 外的所有节点将其向量发送给节点 1。

然后，节点 1 计算求和结果，并将求和后的向量发送回所有其他节点。

这种方法使用多少带宽？每个节点（除节点 1 外）需要将其整个向量发送给节点 1，这意味着发送 D 字节。有 p-1 个节点需要发送数据，因此第一步发送的总数据量为 D (p-1) 字节。

然后，在第二步中，节点 1 必须将求和后的向量发送给所有其他节点。求和后的向量大小为 D 字节，需要发送给 p-1 个其他节点，因此第二步发送的总数据量也为 D (p-1) 字节。

两步总共发送的数据量为 2・D・(p-1)=O (D・p) 字节。注意，这比全连接 mesh 方法的 O (D・p²) 字节要好 p 个数量级。

这种方法需要多长时间？同样，这取决于确切的资源限制，但假设没有资源限制，所有节点都可以同时向节点 1 发送向量。然后，我们需要等待节点 1 计算求和结果。计算完成后，节点 1 可以同时将求和结果发送回所有其他节点。总的来说，这种方法需要 2 个时间步完成，其中节点 1 在每个时间步需要发送或接收 D・(p-1) 字节。

这里我们没有精确测量 “时间步” 的长度，但主要的对比点是，这种方法中，第一步的所有发送必须完成后，第二步的发送才能开始。相比之下，在第一种方法中，所有数据发送可以同时进行。

这种方法的一个缺点是节点 1 存在单点故障。这种方法在实践中并不常用。

方法 3：基于树的结构

在接下来的拓扑结构中，我们构建一棵二叉树。请记住，这里的 “二叉” 是指每个节点最多有 2 个子节点。

要运行 AllReduce：从最底层的叶节点开始，每个节点将其向量发送给其父节点。

当收到所有子节点的向量后，将这些向量与自己的向量求和。

然后，将得到的求和向量发送给父节点。

在树的各层重复此步骤后，根节点将得到总求和结果。

然后，在第二步中，根节点将总求和向量沿树向下发送给其子节点。当从父节点收到求和向量后，需要将该求和向量的副本发送给所有子节点。

这种方法使用多少带宽？在第一步中，每个节点最多从其子节点接收 2 个向量（回想一下：这是二叉树），每个节点向其父节点发送 1 个向量。这给我们每个节点的带宽上限为 3D 字节，因此第一步的总带宽为 3D・p 字节。

然后，在第二步中，每个节点从其父节点接收 1 个向量，并最多向其子节点发送 2 个向量。同样，每个节点的带宽上限为 3D 字节，因此第二步的总带宽为 3D・p 字节。

两步总共发送的数据量为 6・D・p=O (D・p) 字节。这比全连接 mesh 方法好 p 个数量级，与单节点归约方法相同。

这种方法需要多长时间？必须等待收到子节点的向量后，才能将（自己的向量与子节点向量的）求和结果发送给父节点。总的来说，这种方法沿树上行发送向量需要 O (log p) 个时间步，沿树下行发送总求和结果也需要 O (log p) 个时间步，总共需要 O (log p) 个时间步。每个节点在每个时间步需要发送或接收 3D 字节（注意，这比其他方法每个时间步的数据量要少）。精确的时间对比需要代入 D 的值和网络中的资源限制，但大致来说，这种方法需要更多的时间步，但每个时间步可能因为传输的数据量更少而完成得更快。

注意，我们在这个实现中利用了归约操作。每个节点将自己的向量与其子节点的向量求和，这样只需向父节点发送一个求和后的向量。在更朴素的方法中，每个节点会向父节点发送 3 个向量（自己的向量和两个子节点的向量），但我们利用归约节省了带宽。

更一般地说，归约类集合通信（Reduce、ReduceScatter、AllReduce）为我们提供了优化实现的机会。在 Reduce 和 ReduceScatter 中，接收的数据总量实际上小于发送的数据总量，我们可以在实现中利用这一点。例如，如果我们知道输出是所有向量的和，并且收到了两个向量，我们可以将它们求和，然后转发一个求和后的向量，而不是分别转发这两个向量。

方法 4：基于环的结构（朴素版）

在最后两种方法中，我们构建环形拓扑结构。注意，从节点 1 到节点 5 的环绕链路与其他链路没有特殊区别（即链路更长并不意味着什么）。

朴素版的 AllReduce 实现：节点 5 首先向左发送其向量。

当从右侧邻居收到一个向量时，将其与自己的向量求和。

然后，将得到的求和向量发送给左侧邻居。

最终，这个过程会在环中循环进行。

最后，节点 1 会计算出总求和结果。

然后，在第二步中，我们将总求和结果在环中发送，以便所有节点都能获得副本。节点 5 首先向左发送总求和结果。当从右侧邻居收到总求和向量时，将该求和向量的副本发送给左侧邻居。最终，这个过程在环中完成，所有节点都收到总求和结果的副本。

这种方法使用多少带宽？在第一步中，每个节点从右侧邻居接收一个向量，并向左侧邻居发送一个向量。这给我们每个节点的带宽上限为 2D 字节，因此第一步的总带宽为 2D・p 字节。

在第二步中，每个节点再次接收 1 个向量并发送 1 个向量。同样，每个节点的带宽上限为 2D 字节，因此第二步的总带宽为 2D・p 字节。

两步总共发送的数据量为 4・D・p=O (D・p) 字节。

这种方法需要多长时间？必须等待从左侧收到向量后，才能向右侧发送向量。总的来说，这种方法在第一步中环绕环需要 p 个时间步，在第二步中发送总求和结果环绕环也需要 p 个时间步，总共需要 2p=O (p) 个时间步。每个节点在每个时间步需要发送或接收最多 2D 字节。

与基于树的拓扑结构一样，精确的时间对比需要代入 D 的值和网络中的资源限制。大致来说，与前两种方法相比，这种方法需要更多的时间步，但每个时间步可能因为传输的数据量更少而完成得更快。

注意：我们选择节点 5 作为起点，但其他起点也可以。同样，我们也可以选择从左到右在环中传输，而不是从右到左。

方法 5：基于环的结构（优化版）

到目前为止，我们看到的方法都能得到正确的结果，但它们会产生突发性的工作负载。在朴素的基于环的方法中，每个节点大部分时间都处于空闲状态。在某个时刻，突然收到一个完整的向量，必须立即将其与自己的向量相加，并将结果发送到左侧。其他所有节点都必须等待完成这个操作。

为了减少工作负载的突发性，实现更均衡的负载，我们可以错开朴素环式 AllReduce 的步骤。一次性将整个向量发送到左侧会给左侧邻居带来突发的工作。相反，可以通过每个时间步发送一个元素，增量式地向左侧发送向量。

当收到一个元素（来自左侧）时，可以将该元素与其自身对应的元素相加。然后可以将得到的求和结果（仍然是单个元素）发送到左侧。

除了错开向量的发送外，注意起点也被错开了。不再是以节点 5 发送其所有元素为起点，现在改为第 i 个节点发送其第 i 个元素作为开始。

通过在这两个维度上错开操作（每个节点一次发送一个元素，每个节点从不同的元素开始），我们可以实现更均衡的工作负载。在每个时间步，每个节点从右侧接收正好一个元素，计算一个求和，然后向左侧发送正好一个元素。

如果我们重复 p 次，那么每个元素将完整地环绕环一周。

然而，并非所有节点都知道求和向量的所有元素，因此我们需要再环绕环一次。与朴素方法一样，在第二个循环中，当收到总求和的一个元素时，只需将副本发送到右侧。

观看这个动画演示时，请关注我们错开操作的两个维度。如果关注单个列，会注意到我们一次发送一个元素，一次接收一个元素。

另外，如果关注单个行，会注意到每个节点收到目前所有第 i 个元素的和，加上自己的第 i 个元素，然后将新的和发送到左侧。由于这个操作会循环经过所有节点，因此我们最终会将所有第 i 个元素相加。

总之，优化后的环式 AllReduce 与朴素的环式 AllReduce 执行的操作完全相同。唯一的区别是我们错开了向量的发送和接收，以减少每个节点工作负载的突发性。

优化后的环式 AllReduce 的带宽和时间分析与朴素的环式 AllReduce 相同。每个节点在第一步接收 / 发送 2D 字节，在第二步又接收 / 发送 2D 字节，总共 4・D・p=O (D・p) 字节。我们仍然需要 O (p) 个时间步来完成两次环内循环。

然而，优化后的方法每个时间步的带宽得到了改善。在朴素方法中，每个节点必须在单个时间步内接收和发送整个向量，单个时间步传输的总数据量为 2D 字节。在优化方法中，每个节点在每个时间步只需接收和发送单个元素，单个时间步传输的总数据量为 2D/p 字节。

覆盖网络与底层网络拓扑

回想一下，这些集合通信操作的定义是，用户（即 AI 训练程序）可以选择任意 p 台主机，并要求它们运行 AllReduce 操作。当用户选择 p 台主机时，它们不太可能已经以环形拓扑结构连接。即使主机本身没有物理上以环形拓扑连接，我们如何实现基于环的 AllReduce 呢？

答案是使用覆盖网络。我们可以绘制虚拟链路，将主机连接成环形拓扑：

从覆盖网络的角度看，当节点 D 向节点 B 发送向量时，节点 D 是沿着一条（虚拟）链路发送给其直接邻居。从底层网络的角度看，这个向量实际上需要经过多个跳才能到达目的地节点 B。

正如我们在讨论基于覆盖网络的多播时所看到的，覆盖网络的性能取决于覆盖网络拓扑与底层网络的匹配程度。在 AI 训练的背景下，性能尤为重要，因为我们要传输大量数据。

为了说明覆盖网络拓扑的重要性，假设 4 个节点要运行 AllReduce 操作。我们如何对节点进行编号以获得最佳性能？

首先，注意节点的任何编号都能产生正确的 AllReduce 结果。换句话说，任何节点都可以是节点 1，任何节点都可以是节点 2，依此类推。（并非所有集合通信操作都是如此，但 AllReduce 是这样。）

以下是节点的两种可能编号方式：

第一种方法的平均路径伸展度为 3.5。特别是，注意 C 到 D 和 B 到 A 的虚拟链路需要通过底层网络的许多链路。

相比之下，第二种方法的平均路径伸展度为 2.5。这组虚拟链路使环中的相邻链路在物理上更接近。

更一般地说，为了优化基于环的 AllReduce 的性能，我们希望相邻节点（例如节点 i 和节点 i+1）在网络中物理位置较近。

此图显示了任意的底层网络拓扑，但同样的思路也适用于我们用于 AI 训练的高度结构化的数据中心类拓扑。回想一下，在这些数据中心类拓扑中，有些节点具有高性能连接（例如同一台机器上的两个 GPU），而其他节点的连接性能较差（例如不同机架上的两个 GPU）。

AI 训练任务是可预测的，底层拓扑是固定且规则的。这意味着我们有很多机会来优化训练任务的性能。例如，我们可以将特定任务分配给特定节点，以便集合通信操作在附近的节点（例如同一机架中的所有节点）上执行。寻找优化 AI 训练任务的方法是一个活跃的研究领域。

抽象层次

总之，你可以从三个抽象层次来思考集合通信操作：

1. 定义。在最高的抽象层次上，我们通过指定输入和预期输出来定义操作。用户只需理解这些定义就能使用集合通信操作。用户不需要知道操作是如何实现的。

2. 覆盖网络。再往下一个抽象层次，我们可以思考在覆盖网络拓扑中交换哪些数据。在这个层次上，可以假设节点以有用的拓扑结构（例如树或环）组织，并且可以沿着该拓扑结构中的虚拟链路发送数据。

3. 底层网络。在最低的抽象层次上，我们思考虚拟链路（覆盖网络）如何对应到底层网络中的实际物理链路。当节点 5 向节点 4 发送向量时，该向量实际上需要通过多个物理路由器和链路进行转发。

（注：文档部分内容可能由 AI 生成）