一、前言

分布式训练系统目标就是将单节点模型训练转换成等价的分布式并行模型训练。对于大语言模型来说，训练过程就是根据数据和损失函数，利用优化算法对神经网络模型参数进行更新的过程。

单节点模型训练系统结构如上图所示，主要由数据和模型两个部分组成。

训练过程会由多个数据小批次（Mini-batch）完成。图中数据表示一个数据小批次。训练系统会利用数据小批次根据损失函数和优化算法生成梯度，从而对模型参数进行修正。

二、并行训练

针对大语言模型多层神经网络的执行过程，可以由一个计算图（Computational Graph）表示。这个图有多个相互连接的算子（Operator），每个算子实现一个神经网络层（Neural Network Layer），而参数则代表了这个层在训练中所更新的权重。计算图的执行过程可以分为前向计算和反向计算两个阶段。

• 前向计算的过程是将数据读入第一个算子，计算出相应的输出结构，然后依此重复这个前向计算过程，直到最后一个算子结束。
• 反向计算过程是根据优化函数和损失，每个算子依次计算梯度，并利用梯度更新本地的参数。在反向计算结束后，该数据小批次的计算完成，系统就会读取下一个数据小批次，继续下一轮的模型参数更新。

根据单个设备模型训练系统的流程，可以看到，如果进行并行加速，可以从数据和模型两个维度进行考虑。

• 数据进行切分（Partition），并将同一个模型复制到多个设备上，并行执行不同的数据分片，这种方式通常被称为数据并行（Data Parallelism，DP）。
• 模型进行划分，将模型中的算子分发到多个设备分别完成，这种方式通常被称为模型并行（Model Parallelism，MP）。
• 当训练大语言模型时，往往需要同时对数据和模型进行切分，从而实现更高程度的并行，这种方式通常被称为混合并行（Hybrid Parallelism，HP）。

2.1 数据并行

在数据并行系统中，每个计算设备都有整个神经网络模型的完整副本（Model Replica），进行迭代时，每个计算设备只分配了一个批次数据样本的子集，并根据该批次样本子集的数据进行网络模型的前向计算。假设一个批次的训练样本数为N，使用M 个计算设备并行计算，每个计算设备会分配到N/M 个样本。前向计算完成后，每个计算设备都会根据本地样本计算损失误差得到梯度Gi（i 为加速卡编号），并将本地梯度Gi 进行广播。所有计算设备需要聚合其他加速度卡给出的梯度值，然后使用平均梯度

对模型进行更新，完成该批次训练。

数据并行训练系统可以通过增加计算设备，有效提升整体训练吞吐量，每秒全局批次数（Global Batch Size Per Second）。与单个计算设备训练相比，其最主要的区别在于反向计算中的梯度需要在所有计算设备中进行同步，以保证每个计算设备上最终得到的是所有进程上梯度的平均值。常见的神经网络框架中都有数据并行方式的具体实现，包括TensorFlow DistributedStrategy、PyTorch Distributed、Horovod DistributedOptimizer 等。由于基于Transformer 架构的大语言模型中每个算子都是依赖单个数据而非批次数据，因此数据并行并不会影响其计算逻辑，一般情况下，各训练设备中前向计算是独立的，不涉及同步问题。数据并行训练加速比最高，但要求每个设备上都备份一份模型，显存占用比较高。

2.2 模型并行

模型并行往往用于解决单节点内存不足的问题。以包含1750 亿参数的GPT-3 模型为例，如果模型中每一个参数都使用32 位浮点数表示，那么模型需要占用700GB（即175G×4 Bytes）内存。如果使用16 位浮点数表示，每个模型副本也需要占用350GB 内存。H100 加速卡也仅支持80GB 显存，无法将整个模型完整放入其中。模型并行可以从计算图角度，用以下两种形式进行切分：

• 按模型的层切分到不同设备，即层间并行或算子间并行（Inter-operator Parallelism），也称之为流水线并行（Pipeline Parallelism，PP）。
• 将计算图层内的参数切分到不同设备，即层内并行或算子内并行（Intra-operator Parallelism），也称之为张量并行（Tensor Parallelism，TP）。
• 混合并行将多种并行策略如数据并行、流水线并行和张量并行等混合使用