大模型基础理论学习笔记——分布式训练

8.分布式训练

8.1 为什么需要分布式训练

近年来，深度学习被广泛应用到各个领域，包括计算机视觉、语言理解、语音识别、广告推荐等。随着ChatGPT的涌现，大模型时代的帷幕被拉开，为了追求模型的能力，模型的规模也越来越大，如GPT-3的模型参数高达1750亿，即使用1024张80GB的A100，完整训练GPT-3也需要1个月的时间。

模型规模的扩大，对硬件（算力、内存）都提出了更高的要求。如下图所示，最近几年，计算机视觉(CV)，自然语言处理(NLP)和语音识别领域最新模型的训练运算量，以大约每两年翻15倍数的速度在增长。而 Transformer 类的模型运算量的增长则更为夸张，约为每两年翻 750 倍。模型的规模近两年也在快速增长，特别是 Transformer 类模型，模型大小平均每两年翻240倍（如下图所示），训练 AI 模型时候所需要的内存一般比模型参数量还要多几倍，因为训练时候需要保存中间层的输出激活值，通常需要增加3到4倍的内存占用^[1]。

然而，在应付最新 AI 模型的训练时，芯片内部、芯片间还有 AI 硬件之间的通信，都已成为不少 AI 应用的瓶颈。受限于物理定律，AI 硬件（GPU）的内存大小仅仅是以每两年翻2倍的速率在增长，持续提高芯片的集成越来越困难，难以跟上模型扩大的需求，这就形成了所谓的 “内存墙” 问题。

大模型计算量演化

大模型规模与GPU显存对比

为了解决算力增速不足的问题，研究者考虑用多节点集群进行分布式训练，将训练扩展到多个 AI 硬件（GPU）上，从而突破于单个硬件内存容量和带宽的限制，支撑更大规模模型的训练。然而这么做也会遇到内存墙的问题：AI 硬件之间会遇到通信瓶颈，甚至比片上数据搬运更慢、效率更低。因此，分布式策略的横向扩展仅在通信量和数据传输量很少的情况下，才适合解决计算密集型问题^[1]。

因此，尽管分布式训练只能在某种条件下实现对更大规模模型的训练支撑，更加有效的改进仍需围绕AI硬件的设计与提升、高效训练算法的研发开展，但不可否认截至目前，其仍为最为行之有效且可快速生效的办法。

8.2 常见的分布式训练策略

目前，深度神经网络的分布式训练主要可以分为数据分布式、模型分布式、流水分布式和混合分布式四种方式。其中，混合分布式是前面三种方法的混合运用。接下来，我们以矩阵乘法的例子对上述四种分布式训练策略进行介绍。

假设神经网络中某一层是做矩阵乘法，其中的输入 $x$ 的形状为 $4\times 5$，模型参数 $w$ 的形状为 $5\times 8$，那么，矩阵乘法输出形状为 $4\times 8$，此过程示意图如下：

单机单卡的训练中，以上矩阵乘法，先计算得到 $out$，并将 $out$ 传递给下一层，并最终计算得到 $loss$，然后在反向传播过程中，得到 $\frac{\partial loss}{\partial w}$，用于更新 $w$。下面，我们按照上述四种分布式训练方式对此过程的分布式计算处理进行介绍。

8.2.1 数据分布式

数据分布式，即是将数据 $x$ 进行切分，将模型 $w$ 复制到每个设备上，从而实使用相同的模型在不同设备上处理不同数据的计算，实现模型的分布式训练。 基于数据并行处理上述矩阵乘法前向推理的过程如下图所示。

这样，在两台设备上，分别得到的输出，都只是逻辑上输出的一半，将两个设备上的输出拼接到一起，才能得到逻辑上完整的输出。应注意的是，神经网络模型的参数是放在GPU内存中的，所得的计算结果也是放在GPU内存中。因此，按照上述计算过程进行计算，数据被分发到了不同的设备上，当采用 $\frac{\partial loss}{\partial w}$ 进行参数更新时，如果直接进行反向传播，则每个设备上的模型参数将更新为不同的参数，造成不同设备上的模型出现差异，且难以再根据参数进行统一，训练过程就出现了问题。

为解决上述问题，就需要对各个设备上的梯度进行AllReduce处理。具体而言，每个设备均基于分配到其上的数据进行梯度计算，得到 $gra{d}_i$，然后将这些梯度进行求和，再将求和后的结果分发给各个设备，如下图所示^[2]。

然而，上述过程只是给出了理论/逻辑层面的设计，在具体实现中上述梯度更新实现的最简单方式是PS（Parameter Server）模式。在这种实现下，有一个server设备用来存放和同步模型参数，有多个worker设备来计算梯度（就像上图的计算方式），所有的worker设备将梯度发送给server设备进行求和取平均，然后再由server设备将更新后的模型参数同步给worker设备。这种方式一个很大的问题就是随着机器GPU卡数的增加，server设备的通信量也是线性增长。在通信带宽确定的情况下（不考虑延迟），GPU卡数越多，通信量越大，通信时间越长， 所需要的时间会随着GPU数量增长而线性增长。

为了解决这一问题，引入了一种GPU卡之间的通信方式的优化算法，就是Ring All-Reuce，它是高性能领域中的一个经典算法，并不是为了深度学习而生。该算法的优点是通信量是恒定的，不随GPU数量的增加而增长。这个算法分两个步骤，第一步是scatter-reduce，然后是allgather。scatter-reduce操作将GPU交换数据，是的每个GPU可得到最终结果的一个块。在allgather中，GPU将交换这些块，使得所有GPU得到完整的结果。如下图所示^[3]。

scatter-reduce

allgather

那么上述过程在具体代码实现时该如何操作呢？事实上，在当前的主流深度学习框架，如Pytorch和Tensorflow中都对上述过程进行了实现。以Pytorch为例，其torch.nn.parallel库中的DistributedDataParallel类（DDP）实际上已经采用一种类似上述Ring All-Reuce过程进行了实现^[4]，并额外引入了一些优化操作：

… Then, each DDP process creates a local Reducer, which later will take care of the gradients synchronization during the backward pass. To improve communication efficiency, the Reducer organizes parameter gradients into buckets, and reduces one bucket at a time.

… Model parameters are allocated into buckets in (roughly) the reverse order of Model.parameters() from the given model. The reason for using the reverse order is because DDP expects gradients to become ready during the backward pass in approximately that order. The figure below shows an example. Note that, the grad0 and grad1 are in bucket1, and the other two gradients are in bucket0.

… Besides bucketing, the Reducer also registers autograd hooks during construction, one hook per parameter. These hooks will be triggered during the backward pass when the gradient becomes ready.)

在基于DDP的具体实现方面，如下为Pytorch官网给出的一个例子，后续的运用可以基于此例子进行修改适配（该示例还融合了模型分布式的部分）^[8]：

import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP   


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)


def cleanup():
    dist.destroy_process_group()


class ToyMpModel(nn.Module):
    def __init__(self, dev0, dev1):
        super(ToyMpModel, self).__init__()
        self.dev0 = dev0
        self.dev1 = dev1
        self.net1 = torch.nn.Linear(10, 10).to(dev0)
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to(dev1)

    def forward(self, x):
        x = x.to(self.dev0)
        x = self.relu(self.net1(x))
        x = x.to(self.dev1)
        return self.net2(x)


def demo_model_parallel(rank, world_size):
    print(f"Running DDP with model parallel example on rank {rank}.")
    setup(rank, world_size)

    # setup mp_model and devices for this process
    dev0 = rank * 2
    dev1 = rank * 2 + 1
    mp_model = ToyMpModel(dev0, dev1)
    ddp_mp_model = DDP(mp_model)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_mp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    # outputs will be on dev1
    outputs = ddp_mp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(dev1)
    loss_fn(outputs, labels).backward()
    optimizer.step()
    cleanup()


if __name__ == "__main__": 

    n_gpus = torch.cuda.device_count()
    assert n_gpus >= 2, f"Requires at least 2 GPUs to run, but got {n_gpus}"
    world_size = n_gpus
    run_demo(demo_basic, world_size)
    run_demo(demo_checkpoint, world_size)
    world_size = n_gpus//2
    run_demo(demo_model_parallel, world_size)

8.2.2 模型分布式

当神经网络非常巨大，数据并行同步梯度的代价就会很大，甚至网络可能巨大到无法存放到单一计算设备中。这时，就可以采用模型分布式策略解决问题。

模型分布式就是将模型 $w$ 被切分到了各个设备上，每个设备只拥有模型的一部分，所有计算设备上的模型拼在一起，才是完整的模型，各个设备上模型的输入输出按照并行、串行的逻辑组织传递，最后得到模型的最终输出。按照模型输入输出的组织逻辑，模型分布式可以进一步划分为三种模式：

• 并行模式：每个设备上的模型解耦合，每个设备上的数据 $x$ 是完整的、一致的，而仅在输出时需要合并。

• 串行模式：神经网络切为多个阶段，并分发到不同的设备上，每个设备之间存在输入输出耦合关系，通过“接力”的方式完成训练。

• 混合模式：并行与串行的混合组织形式。

（1）并行模式

在模型分布式并行策略中，多个设备之间的梯度不再需要AllReduce操作，但数据会在多个设备之间进行广播，产生通信代价（这里指数据不会复制多份而是通过广播来传递输入数据）。基于模型分布式并行处理上述矩阵乘法前向推理的过程如下图所示。

（2）串行模式

当神经网络过于巨大，无法在一个设备上存放时，除了上述的模型分布式并行策略外，还可以选择模型分布式串行策略。模型分布式串行策略，即为将网络切为多个阶段，并分发到不同的计算设备上，各个计算设备之间以“接力”的方式完成训练。

下图展示了一个逻辑上的4层网络（T1到T4）的模型分布式串行计算过程。 4层网络被切分到2个计算设备上，其中GPU0上进行T1与T2的运算，GPU1上进行T3与T4的计算。 GPU0上完成前两层的计算后，它的输出被当作GPU1的输入，继续进行后两层的计算。

（3）混合策略

模型分布式混合策略需要对模型的计算过程进行分析，从而定义模型计算策略，达到：

• 模型的每个部分都能够适配其所在设备（GPU）的内存

• 模型不同部分之间计算耗时的均衡，从而实现并行训练过程的高吞吐

8.2.3 混合分布式

在神经网络的训练中，也可以将多种分布式训练策略混用，以 GPT-3 为例，以下是它训练时的设备分布式训练方案：

• GPT-3模型被分为 64 个阶段，进行模型分布式串行策略。

• 每个阶段都运行在6台DGX-A100主机上。

• 在6台主机之间，进行的是数据分布式训练，模型被复制在6台主机上，数据被分割并分发至6台主机进行训练。

• 每台主机有8张GPU显卡，同一台机器上的8张 GPU 显卡之间是进行模型分布式并行策略训练。
  

上述过程是人工地对模型的训练过程进行设计，从而实现多种分布式策略的混合使用，一方面满足设备的内存限制并充分利用；另一方面增加模型训练过程的吞吐，加速模型训练。事实上，根据近期研究，在一定程度上，这一过程也可以通过自动化的方式实现^[5]，如下图所示，

该方法的实现过程主要分为三个阶段：

阶段一（原子级划分，atomic-level partitioning）：通过启发式规则方法识别原子子组件（atomic subcomponents，即不可拆分的最小组件）。该过程首先将原始模型转换为ONNX格式的任务图（task graph），形成如上图中（b）所示的由任务（tasks）和值（values）两类节点组成的图。原子子组件是上述任务图的相互连接的子图。如上图（b）中蓝色虚线圈出的 $C_1-C_8$ 所示。原子子组件的构建采用如下方法：

• 首先，我们使原子子组件尽可能细颗粒化。

• 然后，我们先顺序前向遍历整个任务图，查看每个子组件的输出是否取决于输入，如果是则其为一个非定常任务，否则则为定常任务。

• 接着，我们反向遍历整个任务图，查看所有子组件的输出，每当我们遇到一个非定常任务，就创建一个原子子组件。当遇到定常任务时，我们将其与其输出的值划入接收其输出的子组件中。

阶段二（模块级划分，block-level partitioning）：通过平衡各个子组件（subcomponents）的计算时间将上述原子子组件划分为更粗粒度的子组件块。该过程的处理依据两个核心指标：1）模块计算时间的均衡；2）通信开销。该阶段具体包括三个步骤，如下图所示：

• 粗粒度化（Coarsening）：算法迭代合并子组件，得到粗粒度子组件；

• 去粗粒度化（Uncoarsening）：通过调整子组件之间的边界削减通信开销；

• 压实（Compaction）：合并在粗粒度化过程中没能合并的子组件。

基于上述步骤，在给定预期模块数 $k$ 后，此阶段最终得到 $k$ 个模块。

阶段三（阶段级划分，stage-level partitioning）：通过使用一种新的基于动态规划的算法高效搜索上述子组件块的合适组合，得到最终划分（partition）。我们记上一阶段所得的模块为 B = { B 1 , … , B k } B=\{B_1,\dots,B_k\} B={B1​,…,Bk​}，并假设 $B$ 是依据计算依赖经过拓扑排序的。于是，当我们考虑将上述图划分为 $S$ 个阶段时，则第 $i$ 个阶段可以表示为连续模块 { B j ∣ b i − 1 < j ≤ b i } \{B_j|b_{i-1}\lt j\leq b_i\} {Bj​∣bi−1​<j≤bi​} 的连续计算，其中，$b_0=0$ 并且 $b_S=|B|$。该阶段的执行时间（包括计算时间和通信时间）表示为 $h(b_{i-1},b_i,s{d}_i)$。其中，$s{d}_i$ 表示当前阶段的副本（replicas）数量（考虑到融入数据分布式的计算过程）。于是，阶段级划分问题就转化成下列优化问题：搜索恰当的 $b_i$ 和 $s{d}_i$，使得下式表示的最长执行时间最短

f S ( { b 1 , … , b S } , { s d 1 , … , s d S } ) = m a x 1 ≤ i ≤ S   h ( b i − 1 , b i , s d i ) f_S(\{b_1,\dots,b_S\},\{sd_1,\dots,sd_S\}) = \underset{1\leq i\leq S}{\mathrm{max}}\,h(b_{i-1},b_i,sd_i) fS​({b1​,…,bS​},{sd1​,…,sdS​})=1≤i≤Smax​h(bi−1​,bi​,sdi​)

基于此优化问题，我们可以采用动态规划方法进行求解。

最后，尽管上述谈到了各种分布式策略，但在实现中仍应考虑可用的/可支撑的编程接口，分布式训练策略的选择影响着训练效率，框架对分布式训练的接口支持程度，决定了算法工程师的开发效率。

参考资料

[1] AI算力的阿喀琉斯之踵：内存墙_OneFlow一流科技有限公司

[2] Collective Operations — NCCL 2.19.3 documentation (nvidia.com)

[3] 【第一期】AI Talk：TensorFlow 分布式训练的线性加速实践 - 知乎 (zhihu.com)

[4] Distributed Data Parallel — PyTorch master documentation

[5] M. Tanaka, K. Taura, T. Hanawa and K. Torisawa, “Automatic Graph Partitioning for Very Large-scale Deep Learning,” 2021 IEEE International Parallel and Distributed Processing Symposium (IPDPS), Portland, OR, USA, 2021, pp. 1004-1013, doi: 10.1109/IPDPS49936.2021.00109.

[6] 第8章 分布式训练 (datawhalechina.github.io)

[7] 流水并行训练 - OneFlow

[8] Getting Started with Distributed Data Parallel — PyTorch Tutorials 2.2.0+cu121 documentation