分布式训练原理总结（DP、PP、TP 、ZeRO）

文章目录

- 一、分布式训练基础知识
- - [1.1 集合通信、集合通信库](#1.1 集合通信、集合通信库)
  - [1.2 通信模式](#1.2 通信模式)
  - - [1.2.1 Parameter Server（2014）](#1.2.1 Parameter Server（2014）)
    - [1.2.2 Ring-AllReduce（2017）](#1.2.2 Ring-AllReduce（2017）)
  - [1.3 同步范式](#1.3 同步范式)
  - [1.4 大模型训练的目标公式](#1.4 大模型训练的目标公式)
- 二、数据并行
- - [2.1 DataParallel（DP)](#2.1 DataParallel（DP))
  - [2.2 DistributedDataParallel（DDP)](#2.2 DistributedDataParallel（DDP))
  - - [2.2.1 原理](#2.2.1 原理)
    - [2.2.2 `all-reduce`优化](#2.2.2 all-reduce优化)
    - [2.2.3 对比DP](#2.2.3 对比DP)
  - [2.3 ZeRO Data Parallelism](#2.3 ZeRO Data Parallelism)
- [三、管道并行（Pipeline Parallelism）](#三、管道并行（Pipeline Parallelism）)
- - [3.1 模型并行（MP,Naive Model Parallelism)](#3.1 模型并行（MP,Naive Model Parallelism))
  - [3.2 GPipe](#3.2 GPipe)
  - - [3.2.1 工作原理](#3.2.1 工作原理)
    - [3.2.2 active checkpoint](#3.2.2 active checkpoint)
    - [3.2.3 实验结果](#3.2.3 实验结果)
  - [3.3 PipeDream](#3.3 PipeDream)
- [四、张量并行（Tensor Parallelism）](#四、张量并行（Tensor Parallelism）)
- - [4.1 1D张量并行](#4.1 1D张量并行)
  - [4.2 MLP并行](#4.2 MLP并行)
  - [4.3 多头注意力并行](#4.3 多头注意力并行)
  - [4.4 实现](#4.4 实现)
- 五、3D并行
- - [5.1 Data Parallelism + Pipeline Parallelism](#5.1 Data Parallelism + Pipeline Parallelism)
  - [5.2 Data Parallelism + Pipeline Parallelism + Tensor Parallelism](#5.2 Data Parallelism + Pipeline Parallelism + Tensor Parallelism)
- 六、总结：如何选择并行策略
- - [6.1 单节点并行化策略](#6.1 单节点并行化策略)
  - [6.2 多节点并行化策略](#6.2 多节点并行化策略)

参考资料：

HF文档《Efficient Training on Multiple GPUs》、《全网最全-超大模型+分布式训练架构和经典论文》

《一文捋顺千亿模型训练技术：流水线并行、张量并行和3D并行》、《大模型-LLM分布式训练框架总结》、《一文了解预训练相关加速技巧》、《Framework(二)：分布式训练》

如果在单个 GPU 上训练模型的速度太慢，或者模型的权重在单个 GPU 上装不下，此时就需要多GPU的分布式训练。首先，PyTorch 内置的DataParallel （DP）和 DistributedDataParallel （DDP）都可用于多 GPU 训练，相比DP，DDP更值得推荐。

一、分布式训练基础知识

1.1 集合通信、集合通信库

参考《分布式训练硬核技术------通信原语》、《NCCL文档》

在计算机科学中，多个处理单元（计算机节点、线程、进程或其他通信实体）之间信息传递的模式有两种：

Point-to-point communication：点对点通信（P2P）：一对一的通信方式，通常用于在两个不同的进程或处理单元之间传递数据，是最基本的通信模式；
Collective communication：集合通信，一对多或多对多或的通信方式，用于在多个处理单元之间进行协同操作，通常在并行计算环境中使用。

在分布式系统中，各个节点间往往存在大量的集合通信需求，而我们可以用消息传递接口(Message Passing Interface, MPI)来定义一些比较底层的消息通信行为，譬如Broadcast、Reduce、Allreduce、Scatter、Gather、Allgather等。

Broadcast：广播行为（对应MPI_Bcast接口）。执行Broadcast时，数据从主节点0广播至其他各个指定的节点（0~3）
Scatter：和broadcast类似，都是一对多的通信方式。不同的是，Broadcast将0号节点将相同的信息发送给所有的节点，而Scatter则是将数据的不同部分，按需发送给所有的节点。
Reduce：规约运算，是一系列简单运算操作的统称，细分可以包括：SUM、MIN、MAX、PROD、LOR等类型的规约操作。Reduce意为减少/精简，因为其操作在每个节点上获取一个输入元素数组，通过执行操作后，将得到精简的更少的元素。
AllReduce：多对多的Reduce，即在所有的节点上都应用同样的Reduce操作。从下图中可以看出，all reduce操作可通过单节点上reduce+broadcast操作完成。
Gather：将多个sender上的数据收集到单个节点上，Gather可以理解为反向的Scatter。
AllGather：多对多的Gather，收集所有数据到所有节点上。从最基础的角度来看，All Gather相当于一个Gather操作之后跟着一个Broadcast操作。
ReduceScatter：将各个节点的输入先进行求和，然后在第0维度按卡数切分，将数据分发到对应的卡上。

上世纪90年代针对HPC领域定义了一套集合通信相关的接口标准，称为 MPI，主要是被应用在科学计算，尤其是超算领域。由于容错性一般，故在机器学习场景下使用较少。

目前流行的集合通信库 Open MPI、NCCL、Gloo等，都在MPI的基础上，对各种集合通信的模式和算法作了各自的实现。例如NCCL（NVIDIA Collective Communications Library）是基于NCIDIA-GPU的一套开源的集合通信库，针对NVIDIA GPU进行了性能优化，实现多GPU和多节点集体通信原语，所以在在英伟达硬件上能带来更低的延迟和更高的带宽。

1.2 通信模式

参考《Framework(二)：分布式训练》、《MXNet之ps-lite及parameter server原理》、《Ring Allreduce》

为了保证分布式训练的结果与单机训练的结果是一致的，需要某种机制在多个机器之间同步信息。目前最流行的模式有两种：

参数服务器模式（Parameter Server,PS）：基于参数服务器的中心化架构模式，
集合通讯模式（Collective Communication,CC）：基于规约(Reduce)的去中心化架构模式。

1.2.1 Parameter Server（2014）

Parameter Server架构

Parameter Server架构具有星状的拓扑结构，有一个或一组服务器来存储模型参数，众多worker服务器负责读取数据，执行前向和反向并计算梯度。通过网络连接，这些worker把自己的梯度上传（push）到参数服务器，参数服务器收集所有的worker的梯度并进行计算之后，各worker再下拉（pull）模型参数。

假设每个节点的数据量是M，在一次梯度同步过程中，N台worker节点都需要和中心PS进行一次通信，则PS节点总通信量为 N×M 。可以看出，这种架构总通信量与集群规模成线性关系。因此，当集群规模较大或模型较大时，参数服务器的带宽可能会成为瓶颈。

1.2.2 Ring-AllReduce（2017）

Ring-AllReduce架构

参数服务器的主要问题是多个 Worker 同时跟 PS 通信，PS 本身有可能成为瓶颈，随着 Worker 数量的增加，整体的通信量也线性增加，加速比可能会停滞在某个点位上。

基于规约的架构是一种去中心化的架构，典型的有Tree All-reduce和Ring All-Reduce（百度于2017年提出）。在Ring All-Reduce架构下，多个worker通过网络组成了一个环，每个worker依次把自己计算出的梯度同步给下一个worker。假设集群规模为N（GPU数量），经过至多 2×(N-1) 轮同步，就可以完成所有worker的更新。

Ring AllReduce 分为3个步骤：Split, ScatterReduce, AllGather。

Split 阶段：根据集群的规模N，把需要同步的数据平均分成N份。

节点数据平均分为5个分块（Split）
ScatterReduce阶段：进行N-1次 ScatterReduce 迭代，每次迭代中，GPU将向其右邻居发送一个块，并从其左邻居接收一个块并累积到该块中。各个节点依次交换数据，使得每个节点只包含最终结果的一部分（1/N）。每个GPU发送和接收的块在每次迭代中都是不同的；第n个GPU从发送块N和接收块N - 1开始，然后从那里向后进行，每次迭代都发送它在前一次迭代中接收到的块。

第一次 ScatterReduce 操作

第二次 ScatterReduce 操作

第三次 ScatterReduce 操作

最后一次 ScatterReduce 操作

ScatterReduce完成时，每个节点都有一个数据块，它累加了其他所有worker节点相应数据块的数据

AllGather阶段：各个节点再次交换数据，最终得到完整的结果。此过程与scatter-reduce是相同的(发送和接收的N-1次迭代)，只是gpu接收的值没有累加，而是简单地覆盖块。第n个GPU首先发送第n+1个块并接收第n个块，然后在以后的迭代中总是发送它刚刚接收到的块。

AllGather 操作操作完成，所有worker 节点的所有数据块就都包含了来自于其他worker节点的数据

最终，经过N-1次 ScatterReduce 操作和 N-1次 AllGather 操作，整个集群就完成了数据同步。假设每个节点的数据量是M，每一次操作各个worker 节点发送的数据量是 M/N，接收的数据量也是 M/N，则总传输量为：

所以在Ring All-Reduce中，通信成本是恒定的，与系统中gpu的数量无关，完全由系统中gpu之间最慢的连接决定。

所有传输都是在离散迭代中同步进行的，因此所有传输的速度受到环中相邻GPU之间最慢(最低带宽)连接的限制.一般来说，如果一个节点上的所有GPU在环中彼此相邻，则该算法的功能最佳。

1.3 同步范式

在基于参数服务器的架构下，多个worker 间的梯度同步需要通过一个中心节点参数服务器来完成，这不可避免要涉及多个 worker 间的合作，一般来说，模型更新有同步（sync）、异步（async）和混合三种模式。

同步模式是指，当所有worker都完成一次梯度计算和参数更新后，才开始下一轮的迭代。这种模式会出现木桶效应，使得整个集群的速度上限受限于最慢的机器。

异步模式则相反，每个worker只关心自己的进度，完成梯度计算后就尝试更新，至于能跟其他多少个worker"互通有无"则完全随机，其过程不可控，有可能出现无法收敛的问题。

混合模式综合了上述两种方式，各个worker都会等待其他worker完成梯度计算和参数更新，但不是永远等待，而是通过一个超时机制来完成。混合模式虽然也带来了一定的不确定性，但影响并不大，因此其应用也最为普遍。

1.4 大模型训练的目标公式

参考《全网最全-超大模型+分布式训练架构和经典论文》

超大模型训练的总体目标就是提升总的训练速度，减少大模型的训练时间，其总的训练速度的公式为：

单卡速度：单卡速度既然是运算速度和数据IO的快慢来决定，那么就需要对单卡训练进行优化，于是主要的技术手段有精度训练、算子融合、梯度累加来加快单卡的训练性能。
加速芯片数量：理论上，AI芯片数量越多，模型训练越快。但是，随着训练数据集规模的进一步增长，加速比的增长并不明显。如数据并行就会出现局限性，当训练资源扩大到一定规模时，由于通信瓶颈的存在，增加计算资源的边际效应并明显，甚至增加资源也没办法进行加速。这时候需要通讯拓扑进行优化，例如通过ring-all-reduce的通讯方式来优化训练模式。
多卡加速比：多卡加速比既然由计算、通讯效率决定，那么就需要结合算法和集群中的网络拓扑一起优化，于是有了数据并行DP、模型并行MP、流水线并行PP相互结合的多维度混合并行策略，来增加多卡训练的效率。

总的来说呢，超大模型训练的目标就是优化上面的公式，提升总训练速度。核心思想是将数据和计算有关的图/算子切分到不同设备上，同时尽可能降低设备间通信所需的代价，合理使用多台设备的计算资源，实现高效的并发调度训练，最大化提升训练速度。

二、数据并行

2.1 DataParallel（DP)

博客《Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups》、PyTorch文档DATAPARALLEL

知乎：《PyTorch的分布式》、《上手Distributed Data Parallel的详尽教程》、《DataParallel & DistributedDataParallel分布式训练》

最常见的并行训练方式是数据并行（DataParallel），这种方法把输入数据split到各个 workers中（每个worker拥有全部模型）做并行计算，以解决batch size过大的问题。因为求导以及加和都是线性的，所以数据并行在数学上是等价的。

DataParallel是Pytorch中最容易实现的并行方案，只需要增加一行代码 model = nn.DataParallel(model)，即可实现多卡训练。

python 复制代码

# 数据集的长度为100，batch size为32。模型是一个简单的fc层，输入长度是5，输出是2
input_size,output_size = 5,2
batch_size,data_size = 32,100

model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
    print("Gemfield have ", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),batch_size=batch_size, shuffle=True)

for data in rand_loader:
    input = data.to(device)
    output = model(input)

上述代码中，batch_size=32。但是由于使用了DataParallel，在有2个GPU时，一个batch被划分成了2份，也就是tensor.split(16)，分别送往两个GPU上。

另外，在第一次调用model.to(device)时，模型被加载到了第一个GPU设备上，而在第一次调用output = model(input)也就是第一次执行前向传播时，模型被复制到了其余的GPU上。DataParallel一次迭代的过程总结如下：

主GPU初始化模型，然后复制到每个 GPU上
主GPU读取数据批次，然后将一个batch的数据切分成多个更小的batch，分发给不同的GPU；
在每个GPU上完成前向计算，输出被gather到主GPU上进行loss计算
loss被scatter到每个GPU上，每个GPU通过BP计算得到梯度；
每个GPU上的梯度被reduce到主GPU上，然后模型权重在主GPU上获得更新；
在下一次迭代之前，主GPU将模型参数broadcast到其它GPU上，完成权重参数值的同步。

DataParallel采用的是Parameter Server并行架构，有以下局限性：

采用 PS 架构通信开销大（每个worker都要与参数服务器多次通信）
负载不均衡，主GPU负载大，从而导致 GPU 利用率不足
仅支持单机多卡模式，无法实现多机多卡训练

2.2 DistributedDataParallel（DDP)

详细原理：《Distributed data parallel training using Pytorch on AWS》、DDP 文档《DISTRIBUTED DATA PARALLEL》、

实践：《GETTING STARTED WITH DISTRIBUTED DATA PARALLEL》、《上手Distributed Data Parallel的详尽教程》、《PyTorch分布式训练简明教程(2022更新版)》

2.2.1 原理

在DDP中，不再有主 GPU，每个 GPU 执行相同的任务。推理，损失函数计算，梯度计算等都可以在各个GPU上并行独立地完成，提高了训练的效率和速度。

DDP实现并行训练的核心在于模型间的梯度同步，这是通过all-reduce通信操作实现的，保证每个GPU会得到完全相同的梯度。每个GPU也会建立独立的优化器。由于模型具有同样的初始状态和后续相同的梯度，因此每轮迭代后不同进程间的模型是完全相同的，这保证了DDP的数理一致性。下面根据DDP的流程图，介绍其主要步骤：

DDP的流程图。其中Construction只在训练开始前执行一次，仅Forward和Backward在训练中重复多次

Construction阶段：DDP需要额外的建立进程组阶段（Construction），用于确定通信协议和总进程数
- 通信协议是DDP的底层基础，决定了并行训练的方式；
- 总进程数（worldsize）：即有多少个独立的并行进程参与训练；
根据需求每个进程可以占用一个或多个GPU，但并不推荐多个进程共享一个GPU，这会造成潜在的性能损失。为了便于理解，在本文的所有示例中我们假定每个进程只占用1个GPU，占用多个GPU的情况只需要简单的调整GPU映射关系就好。

Construction只在训练开始前执行，因此不会带来额外的延迟
初始状态广播
在每个GPU上创建一个模型的副本（replica），GPU-1中模型的状态会被广播到所有其他进程，确保所有模型具有相同的初始状态。
训练迭代：每个GPU接收其分配的数据批次，并在本地（每个GPU本身）计算损失函数的梯度（local gradients）。
全体度求和（All-Reduce）：为了保持模型权重的一致性并允许模型更新，需要将来自各个GPU的local gradients汇总为一个global gradients，这个操作称为all-reduce，其包含两个操作，reduce-scatter和all-gather：
- Reduce-Scatter：
  - 每个节点（replica）的local gradients都被分成不同的块或分片（blocks or shards，通常是均匀分配的），然后在N个节点之进行 N-1 轮数据交换，这使得每个节点获得了其他节点的一部分梯度信息。
  - 每一轮的交换都会将部分梯度数据合并，以获得一部分全局梯度信息。最终每个 replica 都会持有来自所有其他 replica 的梯度分片的汇总（fully reduced data）
  reduced是指每一轮的合并操作，使得shards减少。fully reduced就是指最终所有的梯度shards都被合并成一个值。
- All-gather：
  - 每个 replica 将在 Reduce-Scatter 阶段中获得的 fully reduced data再次分成多个小块，然后将其广播其它节点。同样经过N-1轮数据交换，每个节点都获得了其他节点的全部梯度信息。
  - 每个节点根据收集到的全部梯度信息汇总为global gradients。
基于global gradients进行模型参数更新（weight update）
重复步骤3至5，直到训练完成。

在 Reduce-Scatter阶段，之所以要先将梯度数据被分成不同的shards，再分N-1次进行广播，而不是在一轮之内广播完成，是为了减少通信开销和内存占用，同时提高训练速度和可扩展性。因为在大规模训练中，GPU的数量可能非常大，一次性将不同GPU的local gradients广播给所有GPU可能导致大量的数据流通信，也会占用大量的GPU内存。通过分桶操作，每次只传递local gradients的部分数据，可以减轻网络负载和内存压力，这样训练可以扩展到更多的GPU或更大的模型规模。

2.2.2 `all-reduce`优化

为了优化性能，DDP中针对all-reduce操作进行了更深入的设计。梯度的计算过程和进程间的通信过程分别需要消耗一定量的时间。等待模型所有的参数都计算完梯度再进行通信显然不是最优的。如下图所示，DDP中的设计是通过将全部模型参数划分为无数个小的bucket，在bucket级别建立all-reduce。当所有进程中bucket0的梯度计算完成后就立刻开始通信，此时bucket1中梯度还在计算。这样可以实现计算和通信过程的时间重叠，使得DDP的训练更高效。

DDP后端的通信由多种CPP编写的协议支持，不同协议具有不同的通信算子的支持，在开发中可以根据需求选择。

2.2.3 对比DP

DDP的优点是负载分散在每个gpu节点上，而且通信成本是恒定的，与 GPU 数量无关，所以训练更快；另外还支持单机多卡和多机多卡。其局限性在于，当模型无法加载在单个GPU时，无法处理。下面通过一个实验来说明 DP 和 DDP 之间的区别。

Hardware: 2x TITAN RTX 24GB each + NVlink with 2 NVLinks (NV2 in nvidia-smi topo -m).
Software: pytorch-1.8-to-be + cuda-11.0 / transformers==4.3.0.dev0.

其中一个实验使用NCCL_P2P_DISABLE=1禁用 NVLink 功能，整个测试结果为：

python 复制代码

Type	NVlink	Time
2:DP	 Y	  110s
2:DDP	 Y	  101s
2:DDP	 N	  131s

可以看到，DP 比使用 NVlink 的 DDP 慢 ~10%，但比不使用 NVlink 的 DDP 快 ~15%。真正的区别将取决于每个 GPU 需要与其他 GPU 同步多少数据------要同步的数据越多，slow link就越会阻碍整体runtime。更多内容详见《Efficient Training on Multiple GPUs》

NVLink（NVIDIA NVLink）是由NVIDIA开发的一种高速互连技术，允许不同GPU之间以非常高的带宽进行直接通信，而无需经过较慢的PCI-E总线。

2.3 ZeRO Data Parallelism

参考《ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters》、《大模型分布式训练策略：ZeRO、FSDP》、《大模型-LLM分布式训练框架总结》

数据并行和模型并行都保持了整个训练过程中所需的所有模型状态，但并不是所有时候这都是必需的。例如，仅在某个层的正向传播和反向传播期间才需要与每个层对应的参数。

ZeRO-DP是一种改进的数据并行性方法，它通过对参数（包括优化器状态、梯度和模型参数）进行分片来消除内存冗余，使得每个GPU仅保存部分参数及相关状态，提高了内存效率；同时还通过在训练过程中使用动态通信来保持计算和通信效率。

图1：使用Adam优化器进行混合精度训练时，ZeRO-DP优化的三个阶段中，每个设备的model states内存消耗。

参数解释：

Baseline：未优化的基线

Ψ：模型大小，上图假设模型参数为Ψ=75亿

K：存储优化器状态要消耗的内存倍数，上一节讲过，对于混合精度的Adam优化器而言，K=12

N d N_d Nd：数据并行度。基于Adam优化器的混合精度训练，数据并行度为Nd=64（即64个GPU）

上图展示了ZeRO-DP对数据并行优化的三个阶段：

优化器状态分割（ P o s P_{os} Pos）：

在每个gpu中保存全部的参数和梯度，但是只保存1/Nd的优化器变量。通过将优化器状态进行分割，实现4倍的内存减少，同时保持与DP相同的通信量。
梯度分割（ P o s + g P_{os+g} Pos+g）：
每个gpu中只保存1/Nd的梯度，实现8倍的内存减少，并保持与DP相同的通信量。
参数分割（ P o s + g + p P_{os+g+p} Pos+g+p）：
每个gpu中只保存1/Nd的参数，实现64倍的内存减少，通信量会略微增加50%。作者通过用少量的计算的成本和通信成本换来了大幅的内存节省。

Zero DP与DataParallel类似，不同之处在于，每个 GPU 不是复制完整的模型参数、梯度和优化器状态，而是只存储其中的一部分。假设有一个具有 3 层（La、Lb 和 Lc）的简单模型，其中每层有 3 个参数。例如，图层 La 的权重为 a0、a1 和 a2：

python 复制代码

La | Lb | Lc
---|----|---
a0 | b0 | c0
a1 | b1 | c1
a2 | b2 | c2

如果我们有 3 个 GPU，ZeRO-DP 会将模型拆分为 3 个 GPU，在某种程度上，这与张量并行性的水平切片相同，与垂直切片相反（垂直切片将整个层组放在不同的 GPU 上），如下所示：

python 复制代码

GPU0:
La | Lb | Lc
---|----|---
a0 | b0 | c0

GPU1:
La | Lb | Lc
---|----|---
a1 | b1 | c1

GPU2:
La | Lb | Lc
---|----|---
a2 | b2 | c2

因为采用数据并行，每个GPU会获得一个mini-batch：

python 复制代码

x0 => GPU0
x1 => GPU1
x2 => GPU2

La层：3 个 GPU 中的每一个都可以重建完整的张量，并使用自己的小批量进行前向传递。一旦计算完成，不再需要的数据就会被删除（只在计算过程中使用）

在 GPU0 上，x0 批次需要 a0、a1、a2 参数来执行其通过层的正向路径，但 GPU0 只有 a0。它将从 GPU1 获得 a1，从 GPU2 获得 a2，将模型的所有部分组合在一起。
GPU1将从 GPU0 和 GPU2获取参数a0，a2来计算x1 批次；
GPU2将从 GPU0 和 GPU1获取参数a0，a1来计算x2 批次；

然后对 Lb 层、Lc层重复整个过程完成前向计算，再反过来完成后向计算。完整工作流程如下：

另外Zero除了Zero DP还包括Zero-R，后续改进工作还有ZeRO-Offload和ZeRO-Infinity，分别实现将GPU的数据和计算卸载到CPU和NVMe内存，能在有限资源下能够训练前所未有规模的模型，而无需对模型代码进行重构。详情见另一篇博客《大模型分布式训练策略：ZeRO、FSDP》。

ZeRO实现 ：DeepSpeed（实现Zero DP、Zero-R、ZeRO-Offload和ZeRO-Infinity的全部功能、Accelerate集成、transformers集成

三、管道并行（Pipeline Parallelism）

参考：《Efficient Training on Multiple GPUs》、《一文捋顺千亿模型训练技术：流水线并行、张量并行和3D并行》

3.1 模型并行（MP,Naive Model Parallelism)

当一个模型大到单个GPU无法训练时，最朴素（Naive）的想法是对模型的层进行划分。Naive Model Parallelism(Naive MP)通过使用.to()来将特定层分配到特定的GPU上。当数据通过这些层传递时，它会被移动到与该层相同的GPU上，而其他层保持不变。

Naive MP也被称作Vertical MP，因为模型通常是以垂直方式切分的。下面以一个4层的序列模型为例进行介绍：

o u t p u t = L 4 ( L 3 ( L 2 ( L 1 ( i n p u t ) ) ) ) output=L4(L3(L2(L1(input)))) output=L4(L3(L2(L1(input))))

将其按层划分至两个GPU上：

GPU1负责计算： i n t e r m e d i a t e = L 2 ( L 1 ( i n p u t ) ) intermediate=L_2(L_1(input)) intermediate=L2(L1(input))；
GPU2负责计算： o u t p u t = L 4 ( L 3 ( i n t e r m e d i a t e ) ) output=L_4(L_3(intermediate)) output=L4(L3(intermediate))；

整个朴素层并行前向传播和后向传播的过程如上图所示。GPU1执行前向传播，并将激活(activations)缓存下来。然后将 L 2 L_2 L2 层的输出intermediate发送给GPU2，GPU2完成前向传播和loss计算后，开始反向传播。当GPU2完成反向传播后，会将 L 3 L_3 L3的梯度返还给GPU1。GPU1完成最终的反向传播。

通过以上过程可以发现Naive MP存在一些缺点：

低GPU利用率：在任意时刻，有且仅有一个GPU在工作，其他GPU都是空闲的。
计算和通信没有重叠。在发送前向传播的中间结果(FWD)或者反向传播的中间结果(BWD)时，GPU也是空闲的。
高显存占用。GPU1需要保存整个batch的所有激活，直至最后完成参数更新。另外Shared embeddings 也需要在不同GPU之间来回复制。如果batch size很大，这将对显存带来巨大的挑战。

Pipeline Parallelism (PP)几乎与Naive MP相同，但通过将传入的batch分成micro-batches并人工创建一个pipeline来解决了GPU空闲问题，允许不同的GPU同时参与计算过程。

3.2 GPipe

论文《GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism》、《图解大模型训练之：流水线并行Gpipe》

3.2.1 工作原理

假设有4个GPU，并将模型按层划分为4个部分。朴素层并行的过程为：

如图，阴影部分所表示的时间段里，总有GPU在空转。在Gpipe中，将阴影部分定义为bubble。假设有 K块GPU，而单块GPU上做一次forward和backward的时间为： t f b = ( t f + t b ) t_{fb} = (t_{f} + t_{b}) tfb=(tf+tb) 。则：

图中灰色长方形的整体面积为： K ∗ K t f b K*Kt_{fb} K∗Ktfb
图中实际在做forward和backward的面积为： K t f b Kt_{fb} Ktfb
图中阴影部分的面积为： ( K − 1 ) K t f b (K-1)Kt_{fb} (K−1)Ktfb
图像阴影部分的占比为： ( K − 1 ) / K (K-1)/K (K−1)/K

则我们定义出bubble部分的时间复杂度为： O ( K − 1 K ) O(\frac{K-1}{K}) O(KK−1)。当K越大，即GPU的数量越多时，空置的比例接近1，即GPU的资源都被浪费掉了。因此这个问题肯定需要解决。

流水线并行的核心思想是：在模型并行的基础上，进一步引入数据并行的办法，即把原先的数据再划分成若干个batch，送入GPU进行训练。未划分前的数据，叫mini-batch。在mini-batch上再划分的数据，叫micro-batch。

图c中，第一个下标表示GPU编号，第二个下标表示micro-batch编号。假设我们将mini-batch划分为M个。假设单个GPU上完成前向传播或者后向传播的面积为1(也就是上图中的单个小方块面积为1)。上图中的总长度为2(M+K-1)，宽度为K，则总面积为2K(M+K-1)。其中，彩色小方块占用的面积表示GPU执行的时间，为 2KM；空白处bubble面积的占比为 1 − 2 K M 2 K ( M + K − 1 ) 1-\frac{2KM}{2K(M+K-1)} 1−2K(M+K−1)2KM= K − 1 K + M − 1 \frac{K-1}{K+M-1} K+M−1K−1

PP的目标是平衡计算强度（即每个micro-batch的大小）和最小化Pipeline中的空闲时间，以提高效率。Gpipe通过实验证明，当 M > = 4 K M>=4K M>=4K 时，bubble产生的空转时间占比对最终训练时长影响是微小的，可以忽略不计，但这可能会导致更多的通信开销。

将batch切好，并逐一送入GPU的过程，就像一个流水生产线一样（类似于CPU里的流水线），因此也被称为Pipeline Parallelism。

3.2.2 active checkpoint

增大batch size就会线性增大需要被缓存激活的显存需求。在GPipe中，GPU需要在前向传播至反向传播这段时间内缓存激活(activations)。以GPU0为例，micro-batch1的激活需要从timestep 0保存至timestep 13。随着模型的增加，每块GPU中存储的中间结果也会越大。

GPipe为了解决显存的问题，使用了为re-materalization技术，后称为active checkpoint。该技术不需要缓存所有的激活，而是在反向传播的过程中重新计算激活。这降低了对显存的需求，但是增加了计算代价（时间换空间），图例如下：

每块GPU上，我们只保存来自上一块的最后一层输入z，其余的中间结果我们算完就废。等到backward的时候再由保存下来的z重新进行forward来算出activations。

如果你使用Pytorch提供的pipeline接口，其中有一个参数叫checkpoint，就是用来做这一项的。

3.2.3 实验结果

Gpipe分别在AmoebaNet（图像）和Transformer（自然语言）两个大模型上做了实验。

Naive：单卡
Pipeline-N：re-materalization + N卡。
AmeobaNet-D和Trasformer-L：超参数的量
# of Model Parameter表示模型的参数量
Total Model Parameter Memory：模型参数所占内存大小
Peak Activation Memory：峰值时中间结果大小。可以发现，中间结果占据的内存大小是相当可观的。

从实验结果里，我们可以发现：

在Transformer上，Gpipe基本实现了模型大小（参数量）和GPU个数之间的线性关系。例如从32卡增到128卡时，模型的大小也从21.08B增加到82.9B，约扩4倍
对AmoebaNet而言，却没有完全实现线性增长。例如从4卡到8卡，模型大小从1.05B到1.8B，不满足2倍的关系。本质原因是AmoebaNet模型在切割时，没有办法像Transformer一样切得匀称，保证每一块GPU上的内存使用率是差不多的。因此对于AmoebaNet，当GPU个数上升时，某一块GPU可能成为木桶的短板。

上图是使用不同数量的GPU数 K 和不同micro-batches数M 在TPUs上进行GPipe Normalized training throughput（吞吐量）。性能随M提高。当 M ≥ K 时，Transformer模型的加速器数量几乎呈线性加速。如果必要，批次大小会根据内存进行调整。

3.3 PipeDream

论文《PipeDream: Generalized Pipeline Parallelism for DNN Training》、《PipeDream: Fast and Efficient Pipeline Parallel DNN Training》、《PipeDream: 数据并行+流水线》

Model parallel training
GPipe's inter-batch parallelism，频繁的管道刷新导致了 idle time 的增加

GPipe使用 active checkpoint技术，在反向传播的过程中重新计算激活，降低了对显存的需求。在将mini-batch划分为M个micro-batch后，如果 M 较小，则由于重新计算开销和频繁的管道刷新，其在硬件效率方面可能会受到影响。

GPipe需要等所有的microbatch前向传播完成后，才会开始反向传播。PipeDream则是当一个microbatch的前向传播完成后，立即进入反向传播阶段。理论上，反向传播完成后就可以丢弃掉对应microbatch缓存的激活。由于PipeDream的反向传播完成的要比GPipe早，因此也会减少显存的需求。

下图是PipeDream的调度图，4个GPU和8个microbatchs。蓝色的方块表示前向传播，绿色表示反向传播，数字则是microbatch的id。

An example PipeDream pipeline with 4 workers

PipeDream在bubbles上与GPipe没有区别，但是由于PipeDream释放显存的时间更早，因此会降低对显存的需求。

四、张量并行（Tensor Parallelism）

参考论文：《Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM》、《一文捋顺千亿模型训练技术：流水线并行、张量并行和3D并行》

除了模型并行和数据并行外，还有一种更小尺度的并行方式：张量并行。Transformer中的主要部件是全连接层和注意力机制，其核心都是矩阵乘法。张量并行的核心就是将矩阵乘法进行拆分，解决单层参数过大的问题。

在 Tensor Parallelism 中，每个 GPU 处理张量的一个切片，并且只聚合完整的张量以进行需要它的操作。本章概念详见 Megatron-LM 论文《Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM》。

4.1 1D张量并行

对于全连接层，最朴素的思想就是利用分块矩阵计算法则，获得结果的一致性。下面以矩阵乘法的方式辅助理解1D张量并行。

列并行：将权重矩阵行列划分为n份，那么矩阵乘法表示为
X A = X [ A 1 , A 2 , ... , A n ] = [ X A 1 , X A 2 , ... , X A n ] XA=X[A_1,A_2,\dots,A_n]=[XA_1,XA_2,\dots,XA_n] XA=X[A1,A2,...,An]=[XA1,XA2,...,XAn]
行并行：对权重矩阵和输入矩阵都进行划分，那么矩阵乘法表示为
X A = [ X 1 , X 2 , ... , X n ] [ A 1 A 2 ... A n ] = X 1 A 1 + X 2 A 2 + ⋯ + X n A n XA=[X_1,X_2,\dots,X_n] \left[ \begin{array}{l} A_1 \\ A_2 \\ \dots \\ A_n \end{array} \right] = X_1A_1+X_2A_2+\dots+X_nA_n XA=[X1,X2,...,Xn] A1A2...An =X1A1+X2A2+⋯+XnAn

可见，无论是行并行还是列并行，都只需要在各个部分计算完后进行一次通信。只不过列并行将通信的结果进行拼接，而行并行则是对通信结果相加。

4.2 MLP并行

transformer的一个主要部分是MLP层。下面并模拟两层的全链接层。设X、Y是输入和输出，A和B则是两个全链接层的权重，则有：

Y = G e L U ( G e L U ( X A ) B ) Y={GeLU(GeLU(XA)B)} Y=GeLU(GeLU(XA)B)

使用列并行可写作：
Y = G e L U ( G e L U ( X A ) B ) = GeLU ( GeLU ( [ X A 1 , X A 2 , ... , X A n ] ) [ B 1 B 2 ... B n ] ) = GeLU ( [ GeLU ( X A 1 ) B 1 , GeLU ( X A 2 ) B 2 , ... , GeLU ( X A n ) B n ] ) = [ GeLU ( GeLU ( X A 1 ) B 1 ) , ... , GeLU ( GeLU ( X A n ) B n ) ] \begin{aligned} \text Y={GeLU(GeLU(XA)B)}&=\text{GeLU}\Big(\text{GeLU}([XA_1,XA_2,\dots,XA_n]) \left[ \begin{array}{l} B_1 \\ B_2 \\ \dots \\ B_n \end{array} \right] \Big) \\ &=\text{GeLU}([\text{GeLU}(XA_1)B_1,\text{GeLU}(XA_2)B_2,\dots,\text{GeLU}(XA_n)B_n]) \\ &=[\text{GeLU}(\text{GeLU}(XA_1)B_1),\dots,\text{GeLU}(\text{GeLU}(XA_n)B_n)] \end{aligned} Y=GeLU(GeLU(XA)B)=GeLU(GeLU([XA1,XA2,...,XAn]) B1B2...Bn )=GeLU([GeLU(XA1)B1,GeLU(XA2)B2,...,GeLU(XAn)Bn])=[GeLU(GeLU(XA1)B1),...,GeLU(GeLU(XAn)Bn)]

通过上面的公式可以看到。当我们将A和B提前划分好后，就可以独立进行计算，在计算出 G e L U ( GeLU ( X A i ) B i ) {GeLU}(\text{GeLU}(XA_i)B_i) GeLU(GeLU(XAi)Bi) 后再进行通信。也就是说，这个例子中虽然有两个全链接层，但是仅需要在得到最终结果前进行通信即可。 **有多个全链接层堆叠时，仅需要在最终输出时进行一次通信即可。**利用这一原理，我们可以更新任意深度的MLP。

4.3 多头注意力并行

由于多头注意力的各个头之间本质上就是独立的，因此各个头完全可以并行运算。

需要注意的是，TP 需要非常快的网络，因此不建议跨多个节点执行 TP。实际上，如果一个节点有 4 个 GPU，则最高 TP 度数为 4。另外，在DeepSpeed中，将TP称之为tensor slicing。

4.4 实现

Megatron-LM
Parallelformers（目前仅提供推理）
SageMaker ：将 TP 与 DP 相结合，以实现更高效的处理，但只能在 AWS 上使用
OSLO 具有基于 Transformer 的张量并行实现
Accelerate：集成了 Megatron-LM 的 T张量并行。

五、3D并行

《Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model》

5.1 Data Parallelism + Pipeline Parallelism

DeepSpeed 流水线教程中的演示了如何将 DP 与 PP 结合使用。如下图所示，DP rank 0 看不到GPU2， DP rank 1看不到GPU3，对于 DP，只有 GPU 0 和 1 提供数据，就好像只有 2 个 GPU 一样。GPU0 使用 PP 将其部分负载卸载到 GPU2，同样的GPU1 使用 PP 将其部分负载卸载到 GPU3。由于每个维度至少需要 2 个 GPU，因此这里至少需要 4 个 GPU。

5.2 Data Parallelism + Pipeline Parallelism + Tensor Parallelism

实现：Megatron-DeepSpeed、OSLO

3D并行是由数据并行(DP)、张量并行(TP)和流水线并行(PP)组成。下图参考《 3D parallelism: Scaling to trillion-parameter models,》，由于每个维度至少需要 2 个 GPU，因此这里至少需要 8 个 GPU。此功能可通过DeepSpeed 、Megatron-LM 、Varuna、SageMaker 、OSLO 来实现。

DeepSpeed 的主要功能之一是 ZeRO，它是DP 的超级扩展，这是一个独立的功能，不需要 PP 或 TP。但它可以与PP和TP结合使用。

当 ZeRO-DP 与 PP（或TP）结合使用时，它通常仅启用 ZeRO stage 1（优化器分片）。理论上此时可以启用 ZeRO stage 2（梯度分片），但这会对性能产生负面影响。

DP本身已经将一个batch的数据进行划分，PP会将mini-batch进一步划分为micro-batch，这样每一个micro-batch都需要一个额外的 reduce-scatter操作来传递梯度信息，划分的太多会增加通信开销，降低计算效率；
PP本身已经减小了梯度的大小，因为每个阶段只需要处理一部分梯度信息。在这种情况下，进一步采用梯度分片（gradient sharding）来减小梯度的大小可能不会带来显著的内存节省，因为梯度已经相对较小。

同样的道理，此时也不宜启用ZeRO stage 3（参数分片），因为会导致更多的节点间通信。

六、总结：如何选择并行策略

6.1 单节点并行化策略

单个GPU可以装入整个模型
- DDP - Distributed DataParallel
- ZeRO：根据使用的情况和配置，这种方法可能会也可能不会更快，但是，值得尝试一下
单个GPU无法装入整个模型
- PipelineParallel (PP)
- TensorParallel (TP)
- ZeRO

对于非常快速的节点间连接（例如 NVLINK 或 NVSwitch），所有三种策略（PP、ZeRO、TP）都应该产生相似的性能。但是，如果没有这些，PP 将比 TP 或 ZeRO 更快。另外TP 适合在单节点内使用，即TP size <= GPUs per node。

单个GPU无法装入模型最大的layer
- 使用 TensorParallel （TP），因为仅靠 PipelineParallel （PP）不足以容纳大型层。
- 使用ZeRO，同时参考 Methods and tools for efficient training on a single GPU中的方法，以便在单个 GPU 上进行高效训练。

6.2 多节点并行化策略

具有快速的节点间连接（fast inter-node connectivity），例如NVLINK 或 NVSwitch，考虑使用：
- ZeRO ：几乎不需要对模型进行任何修改
- 组合使用PP、TP 和DP，这种方法将减少通信，但需要对模型进行重大更改
节点间连接速度较慢且 GPU 内存仍然不足时，组合使用PP、TP 、DP和ZeRO。

分布式训练原理总结（DP、PP、TP 、ZeRO）

文章目录

一、分布式训练基础知识

1.1 集合通信、集合通信库

1.2 通信模式

1.2.1 Parameter Server（2014）

1.2.2 Ring-AllReduce（2017）

1.3 同步范式

1.4 大模型训练的目标公式

二、数据并行

2.1 DataParallel（DP)

2.2 DistributedDataParallel（DDP)

2.2.1 原理

2.2.2 all-reduce优化

2.2.3 对比DP

2.3 ZeRO Data Parallelism

三、管道并行（Pipeline Parallelism）

3.1 模型并行（MP,Naive Model Parallelism)

3.2 GPipe

3.2.1 工作原理

3.2.2 active checkpoint

3.2.3 实验结果

3.3 PipeDream

四、张量并行（Tensor Parallelism）

4.1 1D张量并行

4.2 MLP并行

4.3 多头注意力并行

4.4 实现

五、3D并行

5.1 Data Parallelism + Pipeline Parallelism

5.2 Data Parallelism + Pipeline Parallelism + Tensor Parallelism

六、总结：如何选择并行策略

6.1 单节点并行化策略

6.2 多节点并行化策略

2.2.2 `all-reduce`优化