分布式混合并行训练关键技术解读
为个人参与深度学习框架飞桨PaddlePaddle 开发时，梳理的个人笔记。
一、并行方式
1.数据并行（Batch维度）
数据并行分为了两种模式：Data Parallel（DP） 和 Distributed Data Parallel（DDP） 。
1.1 Data Parallel
DP是一种单进程多线程的并行策略，只能在单机上进行训练，从卡做Forward和Backward并行，主卡做梯度聚合和优化器更新，具体步骤如下：
- 单进程控制多GPU，即本质上是单进程多线程
- 首先将模型加载到主 GPU 上，再复制到各个指定从 GPU；
- 将输入数据按照 Batch 维度进行拆分，各个 GPU 独立进行 forward 计算；
- 将结果同步给主 GPU 完成梯度计算和参数更新，将更新后的权重参数复制到各个 GPU
存在的问题： 由于其是单进程控制多个GPU，故会存在GPU之间负载不均衡的问题，主GPU负载较大。
1.2 Distributed Data Parallel（DDP）
DDP 采用 AllReduce 架构，多进程的方式，突破锁的束缚。在单机和多机上都可以使用。负载分散在每个 GPU 节点上，通信成本（时间）是恒定的，与 GPU 数量无关，等于V/B（参数量/带宽）。DDP不需要通过主GPU分发全模型的参数到每个GPU上。使用ring-all-reduce的方式进行通讯，随着 GPU 数量 N 增加，总传输量恒定。也就是理论上，随着GPU数量的增加，ring all-reduce有线性加速能力。
- 在飞桨中，
paddle.DataParallel接口默认提供的是DDP功能 - 提供了
no_sync()接口，用于暂停梯度同步的上下文管理器。在 no_sync()中参数梯度只会在模型上累加；直到 with 之外的第一个 forward-backward，梯度才会被同步。
>>> import numpy
>>> import paddle
>>> import paddle.distributed as dist
>>> from paddle.autograd import PyLayer
>>> from paddle.distributed.fleet.utils.hybrid_parallel_util import fused_allreduce_gradients
>>> class cus_tanh(PyLayer):
... @staticmethod
... def forward(ctx, x):
... y = paddle.tanh(x)
... ctx.save_for_backward(y)
... return y
... @staticmethod
... def backward(ctx, dy):
... y, = ctx.saved_tensor()
... grad = dy * (1 - paddle.square(y))
... return grad
>>> class SimpleNet(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.linear = paddle.nn.Linear(2, 2)
... def forward(self, inputs):
... inputs = cus_tanh.apply(inputs)
... return self.linear(inputs)
>>> if __name__ == '__main__':
... dist.init_parallel_env()
... model = SimpleNet()
... model = paddle.DataParallel(model)
... opt = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
... for step in range(10):
... x_data = numpy.random.randn(2, 2).astype(numpy.float32)
... x = paddle.to_tensor(x_data)
... x.stop_gradient = False
... # step 1 : skip gradient synchronization by 'no_sync'
... with model.no_sync():
... y_pred = model(x)
... loss = y_pred.mean()
... loss.backward()
... # step 2 : fuse + allreduce manually before optimization
... fused_allreduce_gradients(list(model.parameters()), None)
... opt.step()
... opt.clear_grad()
1.3 数据并行使用技巧
1.3.1 学习率设置
数据并行模式下学习率的设置技巧，其基本原则是学习率正比于 global batch size。 与单卡训练相比，数据并行训练通常有两种配置：
- 一种是保持保持所有计算设备的 batch size 的总和（我们称为 global batch size）与单卡训练的 batch size 保持一致。这种情形下，由于数据并行训练和单卡训练的 global batch size 是一致的，通常保持数据并行模式下各个计算设备上的学习率与单卡训练一致。
- 另一种情形是，保持数据并行模式下每个计算设备的 batch size 和单卡训练的 batch size 一致。这种情形下，数据并行模式的 global batch size 是单卡训练的 N 倍。这里， N 指的是数据并行计算的设备数。因此，通常需要将数据并行模式下每个计算设备的学习率相应的设置为单卡训练的 N 倍。
- 这样，数据并行模式下的初始学习率通常较大，不利于模型的收敛。因此，通常需要使用 warm-up 机制。即，在初始训练时使用较小的学习率，并逐步缓慢增加学习率，经过一定迭代次数后，学习率增长到期望的学习率。
1.3.2 数据集切分
数据并行中，我们通常将数据集切分为 N 份，每个训练卡负责训练其中的一份数据。这里， N 是数据并行的并行度。如我们前面介绍的，每一个迭代中，各个训练卡均需要做一次梯度同步。因此，我们需要确保对于每个 epoch ，各个训练卡经历相同的迭代数，否则，运行迭代数多的训练卡会一直等待通信完成。实践中，我们通常通过数据补齐或者丢弃的方式保证各个训练卡经历相同的迭代数。
- 数据补齐的方式指的是，为某些迭代数少训练数据补充部分数据，从而保证切分后的各份数据集的迭代次数相同；
- 丢弃的方式则是丢弃部分迭代次数较多的数据，从而保证各份数据集的迭代次数相同。
通常，在每个 epoch 需要对数据做 shuffle 处理。因此，根据 shuffle 时机的不同，有两种数据切分的方法。
- 一种是在数据切分前做 shuffle；即首先对完整的数据做 shuffle 处理，做相应的数据补充或丢弃，然后做数据的切分。
- 另一种是在数据切分后做 shuffle；即首先做数据的补充或丢弃和数据切分，然后对切分后的每一份数据分别做 shuffle 处理。
2.张量并行
总体而言，是将张量操作划分到多个设备上，以加速计算或增加模型大小；对模型每一层的层内参数进行切分，即对参数矩阵切片，并将不同切片放到不同GPU上；比如将原本在单卡中的矩阵乘法，切分到不同卡中进行矩阵乘法。训练过程中，正向和反向传播计算出的数据通过使用 All gather 或者 All reduce 的方法完成整合。
在Tansformer中，该策略会把 Masked Multi Self Attention 和 Feed Forward 都进行切分以并行化。利用 Transformers 网络的结构，通过添加一些同步原语来创建一个简单的模型并行实现。张量并行适用于模型单层网络参数较大的情况。同时缺点也是十分明显：
- 当环境是多机多卡，张量并行所需的all-reduce通信需要跨服务器进行连接，这比单机多GPU服务器内的高带宽通信要慢（机间通信比卡间通信成本高）
- 高度的模型并行会产生很多小矩阵乘法，这可能会降低GPU的利用率。
张量模型并行需要解决两个问题： 参数如何切分到不同设备（切分方式）；以及切分后，如何保证数学一致性（数学等价）。本文以 NLP 中的 Transformer 结构为例，介绍张量模型并行的切分方式和随机性控制
2.1 Embedding 切分
如下图（a）所示。当采用模型并行时，Embedding 的参数被均匀切分到多个卡上。假设 Embedding 参数的维度为 N*D，并采用 K 张卡执行模型并行，那么模型并行模式下每张卡上的 Embedding 参数的维度为 N//K*D 。当参数的维度 N 不能被卡数 K 整除时，最后一张卡的参数维度值为 (N//K+N%K)*D 。以下图（b）为例，Embedding 参数的维度为 8*D ，采用 2 张卡执行模型并行，那么每张卡上 Embedding 参数的维度为 4*D 。
为了便于说明，以下我们均假设 Embedding 的参数维度值 D 可以被模型并行的卡数 D 整除。此时，每张卡上 Embeeding 参数的索引值为 [0, N/K) ，逻辑索引值为 [k*N/K, (k+1)*N/K) ，其中 k 表示卡序号，0<=k<K。对于输入索引 I，如果该索引在该卡表示的逻辑索引范围内，则返回该索引所表示的表项（索引值为 I-k*N/K ；否则，返回值为全 0 的虚拟表项。随后，通过 AllReduce 操作获取所有输出表项的和，即对应该 Embeding 操作的输出；整个查表过程如下图（b）所示。
2.2 Matmul 切分
2.2.1 列切分
对于矩阵乘操作，是按行或者列将矩阵切分 K 份。假设原始矩阵的维度为 M*N ，则按行切分后，各个卡上的矩阵维度为 M/K*N ；若按列切分，则各个卡上矩阵的维度值为 M*N/K 。图（a）给出单卡上的矩阵乘法。图（b）给出模型并行模式下的矩阵乘法，其中第二个矩阵按列切分到 2 张卡上；两张卡分别得到结果矩阵的一部分。最后，通过 AllGather 通信操作汇聚最终的结果。
2.2.2 行切分
下图给出按行切分矩阵乘法的示例图。其中，图（a）给出单卡上的矩阵乘法。图（b）给出模型并行模式下的矩阵乘法，其中第二个矩阵按行切分到 2 张卡上；第一个矩阵需要按列切分，以满足矩阵乘法的维度要求；两张卡分别得到结果矩阵的一部分。最后，通过 AllReduce 通信操作按元素累加结果矩阵得到最终的结果。
相对于列切分，每张卡上的通信量实际上是翻倍的？我们需要注意一下几点：
- 模型并行下，需要确保模型并行组中各个卡读取相同的数据；
- 模型并行下，除了被切分的算子对应的输出外，其它所有算子的输出在各个卡上是一致的。
3.流水线并行
通常来讲，训练更大规模的网络模型可以在多种任务上取得更好的效果，如提升图像分类任务的准确率。然而，随着参数规模的扩大，AI 加速卡存储（如 GPU 显存）容量问题和卡的协同计算问题成为了训练超大模型的瓶颈。流水线并行从模型切分和调度执行两个角度解决了这些问题，下面将以飞桨流水线并行为例，介绍下基本原理和使用方法。
流水线原理是将不同的 layer 分配给指定 GPU 进行计算，流水线并行只需其之间点对点地通讯传递部分 activations。具体步骤包括：
- 在流水线并行之中，一个模型的各层会在多个GPU上做切分。
- 一个批次（batch）被分割成较小的微批（Micro-Batches），并在这些微批上进行流水线式执行。
- 通过流水线并行，一个模型的层被分散到多个设备上。
- 当用于具有相同transformer块重复的模型时，每个设备可以被分配相同数量的transformer层。
- 在流水线模型并行中，训练会在一个设备上执行一组操作，然后将输出传递到流水线中下一个设备，下一个设备将执行另一组不同操作。
流水线并行的方法，解决了超大模型无法在单设备上装下的难题，也解决了机器之间的通信开销的问题，使得每台机器的数据传输量跟总的网络大小、机器总数、并行规模无关。如下图，在最简配置流水线并行模型下，任意时刻只有单个计算设备处于计算状态，其它计算设备则处于空闲状态，因此设备利用率和计算效率较差。
为了优化流水线并行中设备的计算效率，可以进一步将 mini-batch 切分成若干更小粒度的 micro-batch，以提升流水线并行的并发度，进而达到提升设备利用率和计算效率的目的。如下图所示，一个 mini-batch 被切分为 4 个 micro-batch；前向阶段，每个设备依次计算单个 micro-batch 的结果；从而增加了设备间的并发度，降低了流水线并行 bubble 空间比例，提高了计算效率。
如上图所示先进行前向计算，再进行反向计算，这种方式我们称之为 F-the-B 模式。不难看出这种 F-then-B 模式由于缓存了多个 micro-batch 的中间变量和梯度，显存的实际利用率并不高。接下来我们介绍一种前向计算和反向计算交叉进行的方式，即 1F1B 模型。在 1F1B 模式下，前向计算和反向计算交叉进行，可以及时释放不必要的中间变量。我们以下图 1F1B 中 stage4 的 F42（stage4 的第 2 个 micro-batch 的前向计算）为例，F42 在计算前，F41 的反向 B41（stage4 的第 1 个 micro-batch 的反向计算）已经计算结束，即可释放 F41 的中间变量，从而 F42 可以复用 F41 中间变量的显存。1F1B 方式相比 F-then-B 方式峰值显存可以节省 37.5%，对比朴素流水线并行峰值显存明显下降，设备资源利用率显著提升。
4.混合并行
5.MoE
通常来讲，模型规模的扩展会导致训练成本显著增加，计算资源的限制成为了大规模密集模型训练的瓶颈。为了解决这个问题， 《Outrageously large neural networks: The sparsely-gated mixture-of-experts layer》 提出了一种基于稀疏 MoE 层的深度学习模型架构，即将大模型拆分成多个小模型(专家, expert ), 每轮迭代根据样本决定激活一部分专家用于计算，达到了节省计算资源的效果； 并引入可训练并确保稀疏性的门( gate )机制，以保证计算能力的优化。
与密集模型不同，MoE 将模型的某一层扩展为多个具有相同结构的专家网络( expert )，并由门( gate )网络决定激活哪些 expert 用于计算，从而实现超大规模稀疏模型的训练。 以上图为例，示例模型包含 3 个模型层；如(a)到(b)，将中间层扩展为具有 n 个 expert 的 MoE 结构，并引入 Gating network 和 Top_k 机制，MoE 细节见图(c)，计算过程如下述公式。
上述第 1 个公式表示了包含 n 个专家的 MoE 层的计算过程。具体来讲，首先对样本 x 进行门控计算， W 表示权重矩阵；然后由 Softmax 处理后获得样本 x 被分配到各个 expert 的权重； 然后只取前 k (通常取 1 或者 2）个最大权重，最终整个 MoE Layer 的计算结果就是选中的 k 个专家网络输出的加权和。
import paddle
from paddle.nn import Layer, LayerList, Linear, Dropout
from paddle.incubate.distributed.models.moe import MoELayer
from paddle.distributed.collective import Group
from paddle.distributed import fleet
import numpy as np
# 构建一个可以正常训练的模型
num_experts = 8
d_model = 512
d_hidden = 2048
class ExpertLayer(Layer):
def __init__(self, d_model, d_hidden, name=None):
super().__init__()
self.htoh4 = Linear(d_model, d_hidden)
self.h4toh = Linear(d_hidden, d_model)
def forward(self, x):
x = self.htoh4(x)
x = self.h4toh(x)
return x
# 然后初始化分布式环境，并构建 expert 通信组 moe_group
fleet.init(is_collective=True)
moe_group = paddle.distributed.new_group(list(range(fleet.worker_num())))
# 设置门网络的 gate 策略和 top_k 机制，并将模型单层扩展为 num_expert 个相同结构的专家网络
gate_config = {
"type": "gshard",
"top_k": 2,
}
experts_list = LayerList()
for expi in range(num_experts):
exp_layer = ExpertLayer(d_model, d_hidden)
experts_list.append(exp_layer)
# 接着调用 MoELayer API 封装并创建出 MoE 模型
class Model(Layer):
def __init__(self, d_model, d_hidden, name=None):
super().__init__()
self.linear1 = Linear(d_model, d_model)
self.moe_layer = MoELayer(d_model = d_model,
experts=experts_list,
gate=gate_config,
moe_group=moe_group,
recompute_interval=0)
self.linear2 = Linear(d_model, d_model)
self.dropout = Dropout(p=0.1)
def forward(self, x):
x = self.linear1(x)
x = self.moe_layer(x)
x = self.linear2(x)
x = self.dropout(x)
return x
model = Model(d_model, d_hidden)
optim = paddle.optimizer.SGD(parameters=model.parameters())
# 最后创建数据集，开始训练
for step in range(1, 100):
x = paddle.rand([4, 256, d_model])
y = model(x)
loss = y.mean()
loss.backward()
optim.step()
optim.clear_grad()
print("=== step : {}, loss : {}".format(step, loss.numpy()))
# 运行方式：
# python -m paddle.distributed.launch --gpus=0,1,2,3,4,5,6,7 --log_dir logs train_moe.py
