多 GPU 集体通信

当你将训练或推理扩展到多 GPU 时，需要找到一种在这些 GPU 之间通信的方式。像 NCCL 这样的库专门提供优化的原语来实现 GPU 间通信，这些原语被称为集体通信。

在训练和推理中，GPU 之间共享的消息类型和数量差异很大，但原语保持不变。本文我们将介绍所有重要的集体通信原语，助你开始扩展到多 GPU！

通信模型

在 GPU 之间发送消息的总时间可以建模为：

总时间 = α + nβ, β = 1/B

其中：

• α（alpha） 是每次通信的固定开销，例如建立连接、握手等的时间。它与发送数据量无关。

• β（beta） 是每字节的时间，是带宽的倒数（β = 1/B）。

• n 是消息大小（字节）。

• B 是链路带宽（例如 NVLink 为 900 GB/s）。

记住这句话：α 是每条消息支付一次的固定成本，nβ 随消息大小增长。

下面是一个例子，能帮助你更清晰地理解。在搭载 NVLink 的 NVIDIA H100 上：

• α ≈ 10 μs（建立成本）

• B ≈ 900 GB/s，因此 β ≈ 1.1 皮秒/字节

现在比较两种场景：

• 小消息（1 KB）： α = 10 μs，nβ = 1024 × 1.1 ps ≈ 1.1 ns。α 项大 10,000 倍。固定开销远超实际数据传输。

• 大消息（1 GB）： α = 10 μs，nβ = 1e9 × 1.1 ps ≈ 1.1 ms。nβ 项大 100 倍。带宽占主导，建立成本可忽略不计。

这种权衡决定了你选择哪种算法。当消息较小时，希望最小化通信轮次（每轮花费一个 α）。当消息较大时，希望通过对数据进行流水线操作来保持总线充分利用——每轮的固定开销几乎无关紧要。

通信算法

实现通信算法主要有两种：最小生成树（MST）和环形算法。

最小生成树（MST）

MST 算法针对较小消息非常高效，即适用于优先考虑低延迟的模型。它使用生成树结构来强制最小化数据传输轮次，但未能充分利用可用带宽。它假设一个网络，其中每个节点一次只能与另一个节点通信。

如果这还不好理解，别担心，我们稍后会通过比较 MST 和环形的例子来理解其权衡。

环形算法

在环形算法中，每个 GPU 按环形组织，一次仅与两个邻居通信。

• 消息分割 - 将大消息分成相等的 N 个块。环中的 P 个 GPU 每个获得 N/P 数据。

• 数据传输 - 每个 GPU 将其数据块发送到下一个相邻 GPU。每个 GPU 从左邻居接收一个块，并向右邻居发送一个块。

• 效率 - 所有 GPU 同时执行计算，带宽始终被占用。

现在，来看一个例子：我们希望将一条消息从一个 GPU 发送到节点中的所有其他 GPU。

MST 中的广播

这里有 6 个 GPU，尝试将 GPU 0 的数据发送给组内所有其他 GPU。

• 第 1 轮 - GPU 0 发送给 GPU 1 和 2。
• 第 2 轮 - GPU 1 发送给 GPU 3。
• 第 3 轮 - GPU 2 发送给 GPU 4 和 5。
• 总通信轮次 - log2(6) ≈ 2.585 轮

环形中的广播

• 第 1 轮 - GPU 0 发送给 GPU 1。
• 第 2 轮 - GPU 1 发送给 GPU 2。
• 第 3 轮 - GPU 2 发送给 GPU 3。
• 总通信轮次 = N - 1 = 3 轮

你可以看到 MST 立即优化了延迟，通信轮次比环形算法少。而环形算法通过分块数据并始终让每个节点执行操作，专注于提高带宽利用率。

在机器学习系统中，无论是训练还是推理设置，大多数消息都受带宽限制，因为我们需要移动大规模权重、激活值、张量分片等。

好，现在我们已经有了直观理解，是时候了解集体通信操作了。

广播（Broadcast）

广播将数据从一个 GPU（根节点）发送到组内的每个其他 GPU。每个 GPU 最终拥有完全相同的数据副本。我们在上面已经看到 MST 和环形是如何实现广播的。

在训练中，广播用于在开始时分发模型参数，以便所有副本从相同的初始状态开始。

归约（Reduce）

归约从每个 GPU 获取一个张量，并应用逐元素操作（求和、均值、最小值、最大值），在一个根 GPU 上生成单个结果张量。

这用于跨 GPU 聚合损失值，或在应用优化器之前对梯度求和。结果位于单个 GPU 上。

散射（Scatter）

散射将由一个 GPU 持有的张量分成相等的块，并将每个块发送给不同的 GPU。每个 GPU 最终拥有原始数据的一个唯一片段——没有两个 GPU 持有相同的块。

在数据并行中，散射是将不同的输入数据微批次分布到各 GPU 的自然方式。

收集（Gather）

收集是散射的逆操作。每个 GPU 将其张量发送到单个根 GPU，根 GPU 将它们拼接成一个更大的张量。

通常使用收集来将模型输出或激活值收集到 rank 0，用于检查点、日志记录或评估。

全收集（Allgather）

全收集类似于收集后跟广播。每个 GPU 将其数据发送给所有其他 GPU，每个 GPU 最终拥有来自所有 GPU 的完整拼接数据集。

这在张量并行中至关重要，其中每个 GPU 持有一个张量的分片，但某些操作（如注意力）需要完整张量。全收集允许每个 GPU 重构完整张量，而无需中央瓶颈。

归约散射（Reduce Scatter）

归约散射将归约和散射组合成单一操作。来自所有 GPU 的张量逐元素归约，结果被分成块，每个 GPU 接收一个块。

全归约（Allreduce）

全归约是分布式训练中最重要的集体通信之一。它对来自所有 GPU 的张量进行归约（求和或均值），并广播结果，使每个 GPU 最终拥有完全相同的归约后张量。

在 DDP 中，全归约是梯度同步的骨干。每个 GPU 计算本地梯度，全归约确保每个副本应用相同的权重更新。

总结

通过以上内容，你应该能够理解不同的集体通信原语，以及使用每种操作的权衡，从而将实验扩展到单 GPU 之外。

关键要点是，大多数 ML 工作负载受带宽限制，这就是基于环形的算法主导生产系统的原因。当延迟是瓶颈且消息较小时，基于 MST 的算法很有用。

在接下来的几篇博客中，我将更专注于分布式技术和更生产就绪的库，如 torchtitan。