【理论】千亿模型并行训练技术

千亿模型并行训练技术

摘要：
进行并行加速，可以从数据和模型两个维度进行考虑。

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

2.2. 模型并行（Model Parallelism，MP）
对模型进行划分， 将模型中的算子分发到多个设备分别完成

2.2.1. 流水线并行（Pipeline Parallelism，PP）
按模型的层切分到不同设备，即层间并 行或算子间并行（Inter-operator Parallelism）

流水线气泡（Pipeline Bubble）
计算图中的下游设备（Downstream Device）需要长时 间持续处于空闲状态，等待上游设备（Upstream Device）的计算完成，才能开始计算自身的任务。这种情况导致了设备的平均使用率大幅降低，形成了模型并行气泡（Model Parallelism Bubble），也称为流水线气泡（Pipeline Bubble）
2.2.2. 张量并行 （Tensor Parallelism，TP）
将计算图层内的参数切分到不同设备，即层内并行或算子内并行（Intra-operator Parallelism）

进行并行加速，可以从数据和模型两个维度进行考虑。首先可以对数据进行切分（Partition），并将同一个模型复制到多个设备上，并行执行不同的数据分片，这种方式通常被称为数据并行（Data Parallelism，DP）。还可以对模型进行划分， 将模型中的算子分发到多个设备分别完成，这种方式通常被称为模型并行（Model Parallelism，MP），包含流水线并行（Pipeline Parallelism，PP）和张量并行 （Tensor Parallelism，TP）。当训练超大规模语言模型时，往往需要同时对数据和模型进行切分，从而实现更高程度的并行， 这种方式通常被称为混合并行（Hybrid Parallelism，HP）

1. 数据并行

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

数据并行训练系统可以通过增加计算设备，有效提升整体训练吞吐量，每秒全局批次数（Global Batch Size Per Second) 。它和单计算设备训练相比，最主要的区别就在于反向计算中的梯度需要在 所有计算设备中进行同步，以保证每个计算设备上最终得到的是所有进程上梯度的平均值。

2. 模型并行

模型并行可以从计算图角度，以下两种形式进行切分：

（1）按模型的层切分到不同设备，即层间并 行或算子间并行（Inter-operator Parallelism），也称之为流水线并行（Pipeline Parallelism，PP）；（2） 将计算图层内的参数切分到不同设备，即层内并行或算子内并行（Intra-operator Parallelism），也称 之为张量并行（Tensor Parallelism，TP）。

1. 1. 1. 流水线并行

流水线并行（Pipeline Parallelism，PP）是一种并行计算策略，将模型的各个层分段处理，并将每个段分布在不同的计算设备上，使得前后阶段能够流水式、分批进行工作。

流水线并行通常应用于大规模模型的并行系统中，以有效解决单个计算设备内存不足的问题。图4.6给出了一个由四 个计算设备组成的流水线并行系统，包含了前向计算和后向计算。其中 F1、F2、F3、F4 分别代表 四个前向路径，位于不同的设备上；而 B4、B3、B2、B1 则代表逆序的后向路径，也分别位于四 个不同的设备上。然而，从图中可以看出，计算图中的下游设备（Downstream Device）需要长时 间持续处于空闲状态，等待上游设备（Upstream Device）的计算完成，才能开始计算自身的任务。

这种情况导致了设备的平均使用率大幅降低，形成了模型并行气泡（Model Parallelism Bubble），也称为流水线气泡（Pipeline Bubble）。

朴素流水线策略所产生的并行气泡，使得系统无法充分利用计算资源，降低了系统整体的计 算效率。为了能够减少并行气泡，文献 [131] 提出了 GPipe 方法，将小批次（Mini-batch）进一步划分成更小的微批次（Micro-batch），利用流水线并行方案，每次处理一个微批次的数据。在当前 阶段计算完成得到结果后，将该微批次的结果发送给下游设备，同时开始处理后一个微批次的数 据，这样可以在一定程度上减少并行气泡。如图所示，前向 F1 计算被拆解为了 F11，F12，F13，F14，在计算设备 1 中计算完成 F11 后，会在计算设备 2 中开始进 行 F21 计算，同时计算设备 1 中并行开始 F12 的计算。相比于最原始的流水线并行方法，GPipe 流 水线方法可以有效降低并行气泡。

b. 张量并行

张量并行（Tensor Parallelism，TP）需要根据模型的具体结构和算子类型，解决如何将参数切分到不同设备，以及如何保证切分后数学一致性两个问题。大语言模型都是以 Transformer 结构为基础，Transformer 结构主要由以下三种算子构成：嵌入式表示（Embedding）、矩阵乘（MatMul） 和交叉熵损失（Cross Entropy Loss）计算构成。这三种类型的算子有较大的差异，都需要设计对 应的张量并行策略[130]，才可以实现将参数切分到不同的设备。

对于嵌入表示层的参数，可以按照词维度切分， 每个计算设备只存储部分词向量，然后通过汇总各个设备上的部分词向量，从而得到完整的词向量。

矩阵乘（MatMul）的张量并行要充分利用矩阵了分块乘法原理。举例来说，要实现如下矩阵 乘法 Y = X × A，其中 X 是维度为 M × N 的输入矩阵，A 是维度为 N × K 的参数矩阵，Y 是 结果矩阵，维度为 M × K。如果参数矩阵 A 非常大，甚至超出单张卡的显存容量，那么可以把参 数矩阵 A 切分到多张卡上，并通过集合通信汇集结果，保证最终结果在数学计算上等价于单计算 设备计算结果。参数矩阵 A 存在两种切分方式：

(1) 参数矩阵 A 按列切块，将矩阵 A 按列切成：

(2) 参数矩阵 A 按行切块，将矩阵 A 按行切成：

按列切分：参数矩阵 A 分别将 A1，A2 放置在两个计算设备上。两个计算设备分别计算 Y1 = X × A1 和 Y2 = X × A2。计算完成后，多计算设备间进行通信，从而获取其它计算设备上的计算结果，并拼接在一起得到最终的结果矩阵 Y ，该结果在数学上与单 计算设备计算结果上完全等价。

按行切分：为了满足矩阵乘法规则，输入矩阵 X 需要按列切分 X = [X1|X2]。同时，将矩阵分块，分别放置在两个计算设备上，每个计算设备分别计算 Y1 = X1 × A1 和 Y2 = X2 × A2。计算完成后，多个计算设备间通信获取归约其他卡上的计算结果，可 以得到最终的结果矩阵 Y 。同样，这种切分方式，既可以保证数学上的计算等价性，并解决单计 算设备显存无法容纳，又可以保证单计算设备通过拆分方式可以装下参数 A 的问题。

Transformer 中的 FFN 结构均包含两层全连接（FC）层，即存在两个矩阵乘，这两个矩阵乘 分别采用上述两种切分方式，如图所示。对第一个 FC 层的参数矩阵按列切块，对第二个 FC 层参数矩阵按行切块。这样第一个 FC 层的输出恰好满足第二个 FC 层数据输入要求（按列切分）， 因此可以省去第一个 FC 层后的汇总通信操作。

多头自注意力机制的张量并行与 FFN 类似，因为 具有多个独立的头，因此相较于 FFN 更容易实现并行，其矩阵切分方式如图所示。

分类网络最后一层一般会选用 Softmax 和 Cross_entropy 算子来计算交叉熵损失（Cross Entropy Loss）。如果类别数量非常大，会导致单计算设备内存无法存储和计算 logit 矩阵。针对这一类算 子，可以按照类别维度切分，同时通过中间结果通信，得到最终的全局的交叉熵损失。首先计算 的是 softmax 值，公式如下：

其中，p 表示张量并行的设备号。得到 Softmax 计算结果之后，同时对标签 Target 按类别切分，每 个设备得到部分损失，最后再进行一次通信，得到所有类别的损失。整个过程，只需要进行三次小量的通信，就可以完成交叉熵损失的计算。