引言 — 许多超大规模企业都在竞相构建大型 GPU 集群，通常配备 64K 或更多 GPU，以适应所有类型的生成式 AI (genAI) 训练工作负载。虽然这些大型 Transformer 模型和数据集的大小可能需要数千个 GPU 进行训练，但在 GPU 之间提供任意到任意的非阻塞网络连接可能是一种过度设计。
了解 GenAI 训练工作负载的模型分区和流量模式可以帮助优化网络拓扑，并实现 GPU 和以太网交换机的有效利用。
本文的主要参考文献见文末，前文介绍了GPU 集群规模、模型分区。
本文为第二部分，继续介绍了
GPU集群的分区并行
张量并行
每个流水线阶段中的矩阵 (GEMM) 操作可以分布在多个 GPU 上。张量并行利用这一事实来创建 2D 模型并行（流水线 + 张量），以减少流水线深度和训练时间。
矩阵乘法的并行化非常简单。每当输入矩阵 ( X ) 与权重矩阵 ( Y ) 相乘时，都可以将其分为Nt 个离散部分矩阵乘法，如上图 5 所示，其中Nt为 2。
如果将Nt个部分矩阵乘法分配给Nt 个离散 GPU，则必须将输入X广播到所有这些 GPU。我们将这些 GPU 称为张量并行 GPU。
每个 GPU 都会发生部分乘法以获得结果Zt。每个张量并行 GPU 都需要将其部分结果传达给所有其他张量并行 GPU。每个 GPU 通过连接部分结果（列并行）或添加结果（行并行）来计算最终结果（Z）。此结果将用于该流水线中的后续计算。
每个微批次的Nt GPU之间的all-to-all通信需要高带宽。通信的大小取决于微批次大小和隐藏层（矩阵乘法中使用的权重）的大小。由于高带宽要求，每个流水线中参与张量并行的 GPU 数量通常仅限于 GPU 服务器或节点内的 GPU 数量。这些服务器内 GPU 通过高速 NVlink 和 NVSwitches 连接。回想一下，当两个 GPU 在服务器内时， H100服务器中的 GPU 到 GPU 带宽是它们在两个不同的服务器上时的 9 倍。
如图 6 所示，当一个流水线阶段中的张量并行 GPU 与下一个流水线阶段中的 GPU 交换中间结果时，两个连续流水线阶段的张量并行组之间也需要进行all-to-all的通信。
在上面的矩阵示例中，如果下一个流水线阶段的 GPU 位于不同的服务器中，Nvidia 不会将最终结果Z广播到下一个流水线阶段的所有张量并行组，而是提供像分散-聚集（scatter-gather）这样的集合，如 Megatron-LM 论文中所述。结果可以在发送端分成大小相等的块，每个 GPU 通过叶交换机将一个块发送到下一个流水线阶段中相同张量等级（轨道）的 GPU（图 5.a）。因此，如果每个流水线阶段有八个张量并行 GPU，数据通信量可以减少八分之一。使用此方案，在接收端，每个张量并行 GPU 都可以通过 NVlinks 执行所有聚集以获取所有块并计算最终结果Z，然后再将其用于进一步的矩阵乘法。
梯度聚合流量
梯度聚合涉及聚合所有模型副本中每个参数的梯度。模型副本中的每个 GPU 都参与梯度聚合，所有模型副本中同一rank/流水线中的 GPU 参与其中。因此，每个模型副本中有Nm 个GPU，每次迭代都有Nm 个并行梯度聚合线程（每个线程中有 Nd 个GPU 参与）。
传统上，人们使用 Ring-All-Reduce 之类的方案将梯度以环形模式从一个 GPU 发送到另一个 GPU，其中每个 GPU 将从前一个 GPU 收到的梯度与其本地计算的参数梯度聚合在一起，然后再将其发送到下一个 GPU。这个过程很慢，因为梯度聚合是按顺序完成的，最终结果也会按顺序在环形拓扑中传播回所有 GPU。
Nvidia 集体提供了一种双二叉树机制，它为梯度聚合提供了全带宽和对数延迟。有关更多详细信息，请参阅论文：https://developer.nvidia.com/blog/massively-scale-deep-learning-training-nccl-2-4/。
当使用二叉树进行梯度聚合时，所有模型副本中同一流水线阶段的 GPU 可以排列为树结构的一部分。每个叶节点将其存储的所有参数的梯度发送到其父节点，然后将其与来自其兄弟叶节点的相应梯度相加。此过程以递归方式继续，直到梯度在树的根节点处聚合。
在根节点获得所有梯度的总和后，必须将聚合梯度发送回该树中的所有节点，以更新模型参数的本地副本。根节点首先将聚合梯度发送给其子节点，然后子节点再将其传递给其子节点，直到所有节点都收到它为止。
在双二叉树方法中，使用跨数据并行组的同等级 GPU 构建两个二叉树。第一棵树的叶节点是另一棵树的中间节点。每棵树聚合一半的梯度。如论文所述，在双二叉树中，每个 GPU 最多可以有两个父 GPU 和两个子 GPU，并且性能（训练时间）远优于大型集群中的环形拓扑。对于大型集群，如果仅使用叶交换机即可访问子 GPU 和父 GPU，则部分梯度聚合可以使用叶交换机进行。但梯度聚合（或数据并行流量）还需要使用主干/聚合交换机来聚合所有无法通过叶交换机访问的数据并行 GPU 等级。
虽然树形结构具有较低的延迟，但它们也会在网络中引入许多 2 对 1 和 1 对 2 的流量模式，从而可能造成短暂的拥塞情况。Ring-all-reduce 具有 1 对 1 的流量模式，因此仍是一些超大规模网络运营商的首选，用于减少主干-叶子流量。
GPU 内存优化
GPU 内存存储其所属流水线/张量分区的参数、梯度、优化器状态、中间激活和输入数据。它还需要临时/暂存空间来进行计算。
对于混合精度训练，参数、梯度和优化器状态所需的存储约为 ( 4*P + K*P )。使用 Adam 的优化器时， K为 12。一万亿参数模型需要约 24TB 的存储空间。
此外，中间激活（来自前向传递，用于反向传递）占用的额外空间与模型的批大小和隐藏层大小成正比。激活重新计算，即在反向传递中再次计算中间激活，可以减少内存需求，但需要额外的计算。不过，我们可能需要 1-2TB 的内存来存储输入激活。由于内存碎片等原因，存在额外的暂存空间和效率低下的问题。
总体而言，1.5 万亿GPT-4模型的总内存为 32TB，效率为 75%。如果每个 GPU 有 80GB，则需要 400 个 GPU 才能容纳一个模型副本。
对于Nd模型副本，可以通过在每个模型副本中仅存储一小部分参数和/或梯度和优化器状态来优化内存。在此方案中，GPU 在每次计算之前动态获取存储在其他 GPU 中的参数/状态。这称为“分片”。它增加了额外的通信开销，但减少了内存占用（以及训练所需的 GPU 总数）。微软的论文讨论了约 100B 个参数模型的分片。对于像GPT-4这样的万亿参数模型，分片如何改善 GPU 规模尚不清楚。
GPU-GPU 流量要点
流水线分区的张量并行 GPU 之间的通信需要非常高的带宽。模型分区框架应尽可能将它们保持在一个服务器节点内。
分散-聚集方法可以减少张量并行 GPU 从一个流水线阶段到下一个流水线阶段（在不同服务器中）之间的全部通信。这使得 GPU 服务器能够以轨道优化拓扑连接，以进行流水线并行流量（不同流水线阶段之间的流量），其中一组服务器中的第 N 个GPU可以通过第 N个叶交换机（也称为轨道交换机）以无阻塞带宽相互通信。
数据并行流量主要为梯度聚合。它发生在所有数据并行组中同一rank/流水线的 GPU 之间。分层树聚合在这些 GPU 之间创建许多 2 对 1 或 1 对 2 的流量模式。传输大小与每个 GPU 等级中存储的参数成正比。
一旦集群中的 GPU 被划分为数据/张量和模型并行，它们之间的通信模式就会在每次训练迭代中重复。次优分区的副作用（拥塞、长尾延迟等）会在所有迭代中累积，并对作业完成时间产生不利影响。
通过在所有数据并行 GPU 上分片参数，可以进一步减少集群中的 GPU 数量 - 这会增加更多的数据并行通信，但可以减少集群大小。
状态空间/划分方法
决定张量、管道和数据并行 GPU 的最佳组合的状态空间很大，并且取决于许多因素。
数据并行 GPU 组过多意味着梯度聚合的通信量增加。由于梯度聚合是下一次迭代的障碍，流水线中的停顿会导致计算停顿，从而降低 GPU 利用率。对于给定的批次大小，使用许多数据并行组可以减少小批次和微批次的大小。这可能无法有效利用 GPU 计算资源（因为每个管道和张量组中的计算与微批次大小成正比）。
更深的流水线会导致在下一次迭代之前执行流水线刷新操作时产生更多气泡。它们每次迭代的延迟也更长。
对于给定的小批量，微批量 ( Bu )的数量越大，流水线刷新停滞的影响就越小。但是，增加给定的小批量的微批量也意味着每个微批量的大小会减小，这可能会导致 GPU 计算利用率不足。
如果张量并行组中有超过 8 个 GPU，则意味着这些组中的 GPU 需要使用低带宽连接来连接叶交换机以传输高带宽流量，从而形成瓶颈。几乎所有模型分区方法都试图将张量并行 GPU 的数量保持在低于每台服务器的可用数量。
Nvidia 的Super POD提供了多达 256 个 GPU，这些 GPU 通过 NV 交换机 ( GH200 ) 的层次结构进行连接。使用Super POD时，拥有八个以上的张量并行 GPU 是可以的。
模型状态的分片使数据并行 GPU 之间的通信量增加了约 1.5 倍，但可以减少训练所需的 GPU 总数。这将整体上改善训练时间/成本。
很难手动将模型划分到多个 GPU 上，使得每个分区都符合 GPU 的内存限制，同时充分利用 GPU 的计算能力。Nvidia提供了开源框架 (Alpa/Ray)来自动执行此状态空间搜索，并考虑 GPU 集群拓扑。
此外，基于特定集体操作所需的拓扑和通信模式，NVIDIA Collective Communications Library(NCCL) 构建了横跨参与计算的 GPU 和节点的环或树。这些结构旨在最大限度地减少争用并最大限度地提高吞吐量。
服务器间流量
在训练期间，服务器间流量使用 GPU Direct RDMA 将数据（中间结果、梯度等）从一个 GPU 内存传输到另一个（在不同的服务器上）。GPU Direct RDMA 是 RDMA 的专门版本。它扩展了标准 RDMA 的功能，允许在 GPU 内存和远程设备（如另一个 GPU 或 NIC）之间直接传输数据，而无需主机 CPU 参与。
由于以太网无处不在，且可用于构建结构的交换机/路由器生态系统丰富，因此所有超大规模企业和公共数据中心都在大力投资构建以太网结构。在以太网中，RoCEv2（基于融合以太网/IP 的 RDMA）用于承载这种服务器间流量。这种流量的交换/路由方式与任何其他 IP 流量一样。
RDMA 写入涉及以下步骤
1 - 在发送和接收 GPU/流之间建立队列对 (QP)。使用带外通信共享 QP 信息。在整个训练期间，QP 可以建立一次。
2 - 将 QP 转换为准备状态以发送/接收交易
3 - 准备 RDMA 进行写入（发送方/接收方内存地址、传输大小）
4 – 该操作排队进入已建立的 QP 的发送队列。请注意，数据仍在 GPU 内存中。
5 - 发送服务器上的 RDMA 网络接口卡 (NIC) 接管，从指定的 GPU 内存读取数据，并通过网络将其发送到目标服务器。它使用 GPU 结构的 MTU 大小将写入分解为网络上的多个事务。
6 - 在接收端，RDMA NIC 将传入的数据直接写入指定的 GPU 内存位置
7 - 发送方为 QP 中的每个 RDMA 操作（写入/读取/发送/接收）分配一个序列号。接收方使用序列号来检测丢失的操作。在传统的 RDMA NIC 中，网络不会对 QP 内的数据包进行重新排序，如果检测到丢失的序列号，则接收方会从该点开始停止接收数据包，并请求发送方从丢失的序列号重新传输所有数据包。这称为回退 N 次重传，效率极低，因为它浪费了网络带宽并增加了较长的延迟。
一些 NIC 支持选择性 NACK，它们请求仅重新传输丢失的数据包。一些 NIC（如 Nvidia 的 ConnectX NIC）允许网络对数据包进行重新排序（有限重新排序）。在此模式下，NIC 将操作无序（OOO）直接写入 GPU 内存，而不会触发向发送方的重新传输。NIC 内部的硬件可以使用位图跟踪最多 N 个操作（N 对应于带宽延迟乘积或 RTT），并按顺序将元数据传送给 GPU。此机制巧妙地使用 GPU 内存来存储 OOO 数据包，并且可以在不占用 NIC 内存空间的情况下实现。
参考文献：
Sharada Yeluri： GPU Fabrics for Gen AI Workloads
https://community.juniper.net/blogs/sharada-yeluri/2024/01/02/gpu-fabrics-for-genai-workloads
高阅读量文章
