集合操作
动机：AI 训练
正如你可能在新闻中读到的那样，人工智能（AI）是一个非常活跃的研究领域。现代 AI 系统需要在海量数据上训练模型。
在这些笔记中，我们将完全忽略这些模型的工作细节。你只需要知道，我们从一些未训练的模型开始：可以把它想象成一个充满随机数的大矩阵。然后，我们使用大量的训练数据来训练这个模型：可以把这想象成用训练数据和模型运行许多矩阵乘法运算（即乘法和加法）。最终，输出的是一个训练好的模型：可以把它想象成之前的那个大矩阵，但现在里面充满了有用的数字。
实际上，AI 训练比这复杂得多。例如，训练过程是迭代的：你会用一些训练数据运行模型，看看效果如何。然后，你会根据所犯的错误计算一个误差项，并利用它来更新模型。我们不会关心这些细节。我们只会把训练看作一个黑箱，它在非常大的数据集上运行大量的矩阵乘法。
分布式训练
AI 训练任务太大，无法在单台计算机上串行运行。如果通过一次一个地相乘数字来进行矩阵乘法，训练任务永远也完成不了。相反，我们需要并行化这些任务，以便同时运行许多操作（例如乘法）。分布式计算有许多方法，每种方法都沿着不同的维度对任务进行并行化：
我们可以分割训练数据，让每个节点在不同的子集上进行训练。
我们可以分割模型本身，让每个节点训练模型的不同子集。
我们可以对操作进行流水线处理，让每个节点运行不同的操作子集。例如，如果期望的操作是 “加 5” 然后 “平方”，我们可以进行分割，让你的节点执行加法，然后把结果传递给我，让我的节点执行平方。这样，每一份数据都先经过你的节点，再经过我的节点，以完成整个操作。
同样，我们将完全忽略工作是如何分配的细节。我们有一些大型任务，并且它已被分解为更小的子任务。
我们真正关心的一件重要事情是这些节点如何相互同步。这些节点通常需要相互通信，以确保它们的状态一致。此外，在运行一些操作后，可能每个节点都有输出的一部分，所有人都需要协调将这些部分组合成完整的输出。
结合我们对训练模型和分布式计算的理解，我们可以对分布式训练有一个非常高层次的概述：
- 将任务分割为子任务。每个节点运行一个子任务。
- 每个节点完成子任务后，所有节点交换大量状态。
- 进行下一个任务，并对下一个任务重复步骤 1-2。
我们的重点是第二步中的数据交换，以及如何使这种数据交换高效。
同样，我们不关心正在交换的确切数据是什么。根据我们分配工作的方式以及我们正在构建的特定 AI 模型，我们交换的数据的性质可能会略有不同。我们的重点是数据的交换方式。
分布式训练基础设施
当我们将训练任务分配给许多节点时，每个节点究竟是什么？
每个节点可以是运行标准 CPU 的计算机，但实际上，节点是专门的 GPU（图形处理器）。这些是专门设计用于高效运行 AI 操作（例如矩阵乘法）的处理芯片。除了 GPU，节点也可以是 TPU（张量处理单元），这是谷歌开发的 AI 优化芯片。
单个训练任务可以在几百个甚至几万个节点上运行，这取决于任务的大小、场景以及每个节点的性能。
GPU 互连在类数据中心网络中，这为我们提供了之前见过的数据中心优势：节点在物理上彼此靠近（例如，在同一建筑物中）。节点按结构化拓扑组织（例如 Clos 网络）。节点是同质的（都以相同方式构建），并且链路具有很高的带宽。
如果你走进一个 AI 训练数据中心，你会看到服务器像其他任何数据中心一样按机架组织。然而，与我们迄今为止看到的其他数据中心不同，每台服务器包含一个或多个用于 AI 计算的 GPU。服务器也可以有一个常规的通用 CPU 用于杂项操作，尽管 CPU 通常不是很强，也不做大量的计算工作。服务器上的所有 GPU 使用相同的 NIC 与其他服务器交换数据。
由于每台服务器有多个 GPU，我们必须稍微修改我们的网络拓扑抽象。和以前一样，服务器通过交换机和高带宽链路连接。然而，我们现在还必须考虑同一服务器上的两个节点相互通信的可能性。同一服务器内的通信与跨服务器通信相比效率极高，因此我们可以将服务器内部通信建模为具有无限带宽和零延迟的链路。
每个 GPU 可以有自己的专用内存，我们可以使用像 RDMA（远程直接内存访问）这样的技术来加速一个 GPU 内存和另一个 GPU 内存之间的数据传输。
机架之间的互连有许多不同的拓扑，但就我们而言，我们将使用胖树 Clos 拓扑来连接机架。无论使用哪种拓扑，某些 GPU 对会更近（例如，同一服务器中的 GPU 可以不使用网络进行通信），其他 GPU 对会更远（例如，不同服务器但同一机柜 / 机架中的 GPU，通过单个交换机连接），还有一些 GPU 对会最远（例如，不同机架中的 GPU，通过多跳连接）。较近的 GPU 对比较远的 GPU 对可以以更高的带宽和更低的延迟进行通信。总之，如果你选择任何一对节点，有些对的连接会比其他对更好。
也存在其他拓扑。TPU 上直接内置了路由器，因此可以将 TPU 直接连接到网络中，完全不需要交换机。TPU 的一种常见拓扑是将它们互连在 3D 环面中，看起来像一个边缘环绕的立方体。例如，如果你到达立方体的顶部并通过向上的链路传输，你会到达立方体的底部。或者，如果你到达立方体的前面并通过面向前方的链路传输，你会到达立方体的后面。就像在 Clos 拓扑中一样，有些节点对更近（例如直接相邻），而其他节点对更远（例如多跳之外）。
集合通信：定义
现在我们知道了任务（分布式计算）和我们运行任务的基础设施（GPU/TPU 的类数据中心网络），我们可以正式定义我们要解决的问题。
集合通信的教科书定义是：一组节点以协同方式交换数据，作为群体计算的一部分。通俗地说，其思想是许多节点协同工作以实现一个共同目标，并且节点在这个过程中必须交换数据。
集合通信背后的思想和术语最初是几十年前在超级计算机的背景下发展起来的。由于最近 AI 的进步，这个话题再次成为活跃的研究领域。集合通信库的现代实现包括 NCCL（英伟达）、MSCCL（微软）、TCCL（迅雷研究组）等。如果你感兴趣，NCCL 的代码可以在网上找到。
集合通信与我们迄今为止看到的所有内容有什么不同？我们将关注 3 个主要差异。
高度结构化的通信：到目前为止，当我们考虑网络时，我们已经抽象掉了正在交换的数据。我们提前不知道谁想通信，我们构建网络以便任何一对主机可以在任何时候通信。
相比之下，在集合通信中，节点有一个非常具体的目标要实现，并且我们提前知道这个目标。这意味着与通用互联网不同，我们对通过网络交换的数据以及数据需要交换的时间有非常明确的结构。换句话说，我们有一套严格规定的数据交换和计算流程，所有节点将协同完成。
专用网络基础设施：到目前为止，我们构建的网络可以支持多个连接同时进行。即使在数据中心网络中，多个租户也可以同时通过数据中心网络发送数据。
相比之下，AI 训练任务非常大，通常在专用基础设施上运行。训练任务是网络上运行的唯一任务，没有其他数据通过网络发送。这意味着我们可以准确预测在任何给定时间使用的带宽量。
数据在交换过程中被转换：到目前为止，当我们考虑通过互联网发送数据时（例如 HTTP/TCP/IP 栈），我们的思维模型是服务器有一些数据（例如一个文件），他们想将该数据的副本发送给用户。
相比之下，在运行集合操作时，数据在通过网络发送时可以被转换。这与我们迄今为止看到的任何内容都不同。这些操作通常相当简单（例如计算总和），但这确实意味着发送方发送的数据不一定与接收方收到的数据相同。
我们可以为我们构建的每个 AI 模型从头设计协调通信方案，但这将很繁琐，并导致大量重复工作。相反，我们将定义一组基本的通信模式，称为集合操作。然后，我们可以使用这些集合操作作为基本构建块，为特定任务设计协调通信方案。你可以将基本的集合操作视为分布式通信的 API，例如用户可用的库函数。然后，用户可以以各种方式调用这些集合函数来实现他们的特定目标。
事实证明，我们只需要相对较少的原始集合操作，AI 训练中的大多数任务都可以分解为这些操作，并表示为这些操作的各种组合。
我们的重点将是这些集合操作是什么，以及它们在网络中是如何实现的。我们不会讨论为什么 AI 训练会产生这些特定的集合操作。我们选择这些特定操作作为基本构建块的原因更多地与 AI 计算的性质有关，这超出了我们的范围。
集合操作：设置
现在我们将定义 7 个基本的集合操作。我们将通过指定输入（操作前每个节点持有的数据）和相应的输出（操作后每个节点持有的数据）来定义这些操作应该做什么。我们不会指定操作在网络中是如何实现的（这将在后面介绍）。
输入：有 p 个节点。在我们的示例中，我们将 p 设置为 4，但其他值也可以。
每个节点有一个 p 元素的向量数据。在这些示例中，你可以将数据视为一个由 4 个整数组成的数组。实际上，这些数据也可以是更高维的，例如矩阵的 4 行，或者训练数据的 4 个等大的块。
输出：元素以某种指定的方式在节点之间移动。输出指定作为该操作的结果，哪些值放入哪些框中。
此外，有时元素可以被聚合（例如求和）。同样，输出指定这个特定操作执行的计算（如果有的话），以及将计算结果放入哪些框中。
在集合操作发生之前，需要进行一些额外的协调，以便每个节点知道自己的编号和节点总数（例如，“你是节点 1，总共有 4 个节点”）。这种额外的协调超出了我们的范围，但你可以想象，某个集中式调度器或控制器会将这些信息分发给节点并设置任务。
为了执行集合操作，每个节点并行地、同时运行完全相同的代码。每个节点独立调用相同的集合操作来启动操作，当操作完成时，输出应该与操作定义匹配。理想情况下，节点具有相同的硬件资源，以便它们都同时完成。如果有些节点比其他节点慢，操作是阻塞的，这意味着我们必须等待所有节点完成操作后才能进行下一个任务。
总之，集合操作由设置任务的控制器协调。操作是同步的（所有节点同时开始）、同质的（理想情况下所有节点同时完成）和阻塞的（必须等待所有节点完成才能继续）。
设置完成后，我们现在可以看看这 7 个集合操作是如何定义的。这些操作大致可以分为两类：4 个操作是关于重分布（移动数据而不进行转换），3 个操作是关于合并（将许多数据片段聚合为单个输出）。
操作：广播（Broadcast）
中文描述：获取指定根节点中的整个向量，并将该整个向量的副本发送到每个节点。
注意：此图显示的是以节点 1 为根节点的广播操作，但我们也可以使用不同的根节点执行该操作。运行广播操作的用户必须将根节点指定为操作的 “参数” 之一。
注意：非根节点中的输入向量不用于创建输出。你可以把它们看作是函数中实际未使用的参数。
注意：每个节点的输入和输出向量不一定必须存储在相同的位置。如果你使用相同的内存地址来保存输入和输出向量，那么一些操作（如广播）会用输出数据覆盖输入数据。或者，你可以使用不同的内存地址来保存输出向量。
操作：散射（Scatter）
中文描述：获取指定根节点中的整个向量。将该向量的第 i 个元素发送到第 i 个节点。
注意：与广播操作一样，你可以指定任何节点作为根节点。同样，与广播操作一样，非根节点的输入向量不用于创建输出（可以理解为：函数中未使用的参数）。
操作：收集（Gather）
中文描述：构建一个新向量，其中第 i 个元素定义为来自第 i 个节点的第 i 个元素。将此向量发送到指定的根节点。
注意：在这个操作中，非根节点的接收缓冲区中没有存储任何内容。
操作：全收集（AllGather）
中文描述：构建一个新向量，其中第 i 个元素定义为来自第 i 个节点的第 i 个元素。将这个新向量的副本发送到每个节点。
替代描述（与上述等效）：节点 i 广播其第 i 个元素，使其成为每个节点输出向量的第 i 个元素。
操作：归约（Reduce）
中文描述：计算所有向量的按元素求和，并将得到的和向量发送到指定的根节点。
在这些笔记中，我们将求和用作我们的归约操作，但也可以有其他归约操作。例如，我们可以在归约操作（或归约散射或全归约）中用乘法代替加法。归约操作通常是可结合和可交换的，这大致意味着你可以以任何顺序执行它们，仍然得到相同的结果（例如加法是可结合和可交换的）。
操作：全归约（AllReduce）
中文描述：计算所有向量的按元素求和，并将得到的和向量的副本发送到所有节点。
操作：归约散射（ReduceScatter）
中文描述：计算所有向量的按元素求和。将和向量的第 i 个元素发送到第 i 个节点。
替代描述（与上述等效）：每个节点的第 i 个元素被求和，得到的和（标量）被发送到节点 i。
对偶（Duals）
有些操作对是彼此的对偶。大致来说，这意味着一个操作是另一个操作的逆操作。例如，在数学中，你可以说平方和平方根是彼此的对偶。
在检查一对操作是否形成对偶对时，我们忽略任何归约计算。我们只关心输出中哪些框被写入，而不管写入这些框的值是什么。
广播和归约是彼此的对偶。广播从根节点的 4 个框中读取，写入所有节点的 16 个框中。归约则相反：它从所有节点的 16 个框中读取，写入根节点的 4 个框中。
散射和收集是彼此的对偶。散射从根节点的 4 个框中读取，写入第 i 个节点的第 i 个框（总共 4 个框）。收集则相反：它从第 i 个节点的第 i 个框中读取（总共 4 个框），写入根节点的 4 个框中。
全收集和归约散射是彼此的对偶。全收集从第 i 个节点的第 i 个框中读取（总共 4 个框），写入所有节点的 16 个框中。归约散射则相反：它从所有节点的 16 个框中读取，写入第 i 个节点的第 i 个框（总共 4 个框）。
全归约没有对偶。或者，你可以将全归约视为自身的对偶，因为它从所有 16 个框中读取并写入所有 16 个框中。
当我们开始考虑这些集合操作的实现时，对偶的概念很有用。对于特定的拓扑和路由方案，一个集合操作及其对偶将具有相同的性能（例如相同的总带宽使用量），因为在集合操作及其对偶中发送和接收的数据总量是相同的。
组合操作（Compositing Operations）
用户可以组合多个操作来获得他们期望的操作。
例如，全归约可以等效地表示为先执行归约散射，再执行全收集。
（注：文档部分内容可能由 AI 生成）