AI大模型并行计算：所需的高性能网络方案

发布日期：2024-07-26 06:51:30 浏览次数： 3368

大模型AI训练对网络的要求

关于大模型的发展与应用，尽管存在诸多讨论，但无疑，大模型能力已成为人工智能未来发展的基石。

相较于传统的小模型，大模型对大规模分布式并行训练有着更高的需求。这主要源于两个方面的因素。首先，大模型本身规模庞大，受当前GPU显存限制，不得不将模型拆分成多个部分，分散存储在多个GPU上。例如，文心大模型拥有2600亿个参数，而一个80G显存的A800 GPU，加上训练过程中的计算状态，仅能容纳大约10亿至20亿参数。因此，仅仅存储2600亿的模型就需要一两百块GPU，这已经构成了相当大的规模。

其次，训练更多的参数需要巨大的计算量，这就需要引入更大规模的GPU进行加速。因此，所需的GPU数量也需要有数量级的提升。

我们通常根据任务的GPU卡数来定义训练的规模。百卡以下为小规模，百卡至千卡为中规模，千卡以上为大规模，超过万卡则称为超大规模。按照这种定义，千卡以上的大规模并行训练是大模型成功的关键所在。

分布式并行AI训练的策略主要有三种模式：

首先，数据并行是应用最广泛的策略。在这种策略中，每个GPU都持有模型的相同副本，而数据集则被分割成多个部分，每个GPU处理一部分数据。完成一轮迭代后，所有GPU会同步各自的梯度数据，并据此计算出下一轮迭代所需的参数。数据并行中的网络通信主要体现为对各个GPU上梯度的Allreduce操作，其通信量大小与模型参数规模成正比，对于参数规模达到千亿级别的大模型来说，数据通信量会非常大。

其次，流水线并行是另一种策略。神经网络模型通常由多层神经元构成，包括大型模型的底层Transformer模型也是如此。在这种策略中，模型被按照神经元的层次进行拆分，并分配到不同的GPU上。这种并行策略要求不同GPU之间进行层间点对点的数据传输，包括正向计算中的激活值和反向计算中的梯度值。虽然这种通信在一个迭代中会多次发生，但通信量通常较小，对网络性能的要求相对较低。

最后，张量并行是第三种策略。它涉及多个GPU同时执行一个张量计算任务，如矩阵乘法。在这种策略中，多个GPU需要对局部的张量计算结果进行全局的Allreduce同步。张量计算结果的大小不仅与模型本身有关，还与训练时使用的batch size相关，通常都非常大。因此，在一个迭代过程中，可能会多次进行Allreduce操作，这使得张量并行对网络带宽的需求最大。

AI大模型训练常用混合并行方案：

在训练大型模型时，我们通常会综合采用三种并行策略以优化性能。

首先，在单台机器内部的多GPU卡之间，我们采用张量并行，以充分利用NVLink提供的高带宽特性。

然而，由于模型规模庞大，单一机器上的8个GPU卡往往无法完全容纳整个模型。因此，我们会在多台机器之间采用流水线并行策略，通过构建多路流水线并行来形成一个完整的模型训练单元。

最后，为了进一步提高训练速度，我们还会使用多路数据并行。在数据并行中，我们将每个并行单元称为DP组，每个DP组内部同时包含张量并行和流水线并行。

在数据并行过程中，Allreduce操作实际上是在每个DP组的相同编号的GPU卡之间进行的。例如，如果配置为8路张量并行、4路流水线并行和3路数据并行，那么总共会有32个Allreduce组，每个组包含3张GPU卡进行梯度同步。

值得注意的是，数据并行中的每张GPU卡都需要对大约10GB级别的数据进行Allreduce操作。这一步骤是大模型训练中对网络带宽的主要需求。

大模型训练对网络的要求：

为满足大模型训练的需求，我们设定了AI高性能网络的三大核心目标：超大规模、超高带宽以及超长稳定。

首先，规模是决定模型训练速度的关键因素。如图所示，一个1750亿参数的模型，使用2000张GPU需要超过100天完成训练，而若增至8000张GPU，则能显著缩短至30天左右。对于追求快速迭代的大模型而言，这一优势至关重要。

其次，Allreduce带宽直接影响大规模分布式训练的整体效率。从图中可以看出，当单GPU的Allreduce带宽达到5GB/s时，整体加速比仅为70%。若想达到90%的加速比，单GPU的AllReduce带宽需提升至20GB/s，相当于单GPU完全利用400G网卡的能力。

最后，稳定性在长时间训练中尤为关键。以GPU可用性为例，假设单GPU的月可用性为99.9%，则在千卡规模下，模型在一个月内遭遇故障导致中断的概率为60%。若使用8000张GPU，这一概率将高达99%。即使GPU可用性提升至99.99%，8000卡下的中断概率仍接近50%。因此，网络必须保证更高的可用性，以减少模型训练的中断次数，降低检查点（checkpoint）的频率和开销。

为了满足大规模AI训练的需求，需要设计AIPod网络，选择了3层无收敛的CLOS组网结构来支撑AIPod网络的超大规模。该结构由服务器、Leaf层交换机（LF）、Spine交换机（SP）和SuperSpine交换机（SSP）组成。服务器直接连接到Leaf层交换机，随后通过Spine交换机实现互联，并最终通过SuperSpine交换机实现不同通道间的通信。

在大模型训练中，同号GPU卡之间的通信占据主导。因此，AIPod网络采用了8通道架构，确保每个服务器上的8个GPU分别连接至8个不同的Leaf交换机，形成一个汇聚组。这样的设计使得每个汇聚组下最大可支持512张GPU，整个集群则可支持超过16K GPU（可进一步扩展）。

尽管主要通信发生在同一通道内，但跨通道通信仍不可避免。为此，需要通过SuperSpine交换机连接不同通道的Spine交换机，实现全通道通信。AIPod的组网方式确保无收敛，即交换机的上联带宽与下联带宽相等，从而确保集群内互通带宽的充足性。

为了尽可能支撑更大的规模，在选择交换机的时候，会选用当前顶级容量的交换芯片，比如曾经的 12.8T 或者 25.6T 芯片，很快行业客户现在已经演进到了单芯片 51.2T 的交换机。