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

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

大模型这话题，最近确实被炒得很热。但不管外界怎么争论，有一点几乎没争议：大模型能力，已经是人工智能继续往前走的基石了。

跟传统小模型一比，大模型对大规模分布式并行训练的需求，那真不是一个量级。原因其实很简单，就两个。

第一，大模型本身太大，受限于当前GPU显存，只能把模型拆成好几块，分散存到多个GPU上。举个例子，文心大模型有2600亿个参数。一个80G显存的A800 GPU，算上训练过程中的计算状态，大概只能塞下10亿到20亿参数。光存这个模型，就得一两百块GPU——这规模已经相当可观了。

第二，参数越多，需要的计算量越大，就得拿更多的GPU来加速。所以GPU数量得成倍往上翻。

通常我们按任务的GPU卡数来划分训练规模：百卡以下算小规模，百卡到千卡是中规模，千卡以上算大规模，超过万卡就是超大规模。按这个标准，千卡以上的大规模并行训练，才是大模型成功的关键。

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

先说说数据并行。这是用得最广的一种。每个GPU都复制一份相同的模型，数据集切成多份，每个GPU各处理一部分。一轮迭代完后，所有GPU同步各自的梯度，再算出下一轮的参数。这种并行中，网络通信主要是对梯度的Allreduce操作。通信量跟模型参数规模成正比，参数一上千亿，数据量就非常大了。

接着是流水线并行。神经网络模型一般有很多层，大模型的Transformer也一样。这种策略是按神经元的层次来拆分模型，分给不同的GPU。不同GPU之间需要传层间的点对点数据：正向计算时的激活值，反向计算时的梯度值。虽然每次迭代会传很多次，但单次数据量不大，对网络性能的要求相对低一些。

最后是张量并行。多个GPU同时算一个张量计算，比如矩阵乘法。算完后，局部结果要做全局的Allreduce同步。结果大小不仅跟模型有关，还跟训练用的batch size有关，通常都很大。而且一次迭代里可能要多次Allreduce，所以张量并行对网络带宽的需求最高。

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

实际训练大模型时，三种策略往往要一块用，才能把性能拉满。

在单台机器的多张GPU之间，用张量并行，充分发挥NVLink的高带宽优势。

但一个模型太大，单台机器的8张GPU往往装不下，所以多台机器之间用流水线并行，搭成多路流水线，形成一个完整的训练单元。

最后，为了进一步提速，还要上多路数据并行。数据并行里，每个并行单元叫DP组，每组内部同时包含张量并行和流水线并行。

数据并行中的Allreduce操作，是在每个DP组里编号相同的GPU之间进行的。比如配置8路张量并行、4路流水线并行、3路数据并行，那总共有32个Allreduce组，每组3张GPU做梯度同步。

注意，数据并行里每张GPU卡要对大约10GB级别的数据做Allreduce。这一步才是大模型训练对网络带宽的主要需求。

大模型训练对网络的要求

为了满足训练需求，高性能网络得盯住三个核心目标：超大规模、超高带宽、超长稳定。

先说规模。规模直接决定训练速度。举个例子，一个1750亿参数的模型，用2000张GPU要训练超过100天，但换成8000张GPU，就能缩短到30天左右。对大模型这种需要快速迭代的来说，规模优势太重要了。

再看带宽。Allreduce带宽直接影响整体效率。数据显示，单GPU的Allreduce带宽到5GB/s时，加速比只有70%；要干到90%的加速比，单GPU的AllReduce带宽得提到20GB/s，相当于把400G网卡的能力用满。

最后是稳定性。长时间训练，稳定性是性命攸关的事。假设单GPU的月可用性99.9%，千卡规模下，一个月内因故障中断的概率是60%。换成8000张GPU，这个概率飙到99%。就算GPU可用性提到99.99%，8000卡下的中断概率仍然接近50%。所以网络必须比GPU还稳，才能减少中断次数，降低检查点（checkpoint）的开销。

为支撑大规模AI训练，需要设计AIPod网络，采用3层无收敛的CLOS组网结构。服务器直接连Leaf交换机（LF），再上联Spine交换机（SP），最后通过SuperSpine交换机（SSP）实现跨通道通信。

大模型训练中，同号GPU之间的通信占主导，所以AIPod采用8通道架构：每台服务器上的8个GPU分别接到8个不同的Leaf交换机，形成一个汇聚组。每个汇聚组最多支持512张GPU，整个集群可支持超过16K GPU（还能继续扩）。

虽然主要通信在通道内部，但跨通道通信也免不了。所以用SuperSpine交换机把不同通道的Spine连起来，实现全通道互通。AIPod的组网方式确保无收敛——交换机上联带宽等于下联带宽，保证集群内互通的带宽充足。

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