分布式训练原理总结（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等。

1. Broadcast：广播行为（对应MPI_Bcast接口）。执行Broadcast时，数据从主节点0广播至其他各个指定的节点（0~3）

2. Scatter：和broadcast类似，都是一对多的通信方式。不同的是，Broadcast将0号节点将相同的信息发送给所有的节点，而Scatter则是将数据的不同部分，按需发送给所有的节点。
   

3. Reduce：规约运算，是一系列简单运算操作的统称，细分可以包括：SUM、MIN、MAX、PROD、LOR等类型的规约操作。Reduce意为减少/精简，因为其操作在每个节点上获取一个输入元素数组，通过执行操作后，将得到精简的更少的元素。
   

4. AllReduce：多对多的Reduce，即在所有的节点上都应用同样的Reduce操作。从下图中可以看出，all reduce操作可通过单节点上reduce+broadcast操作完成。
   

5. Gather：将多个sender上的数据收集到单个节点上，Gather可以理解为反向的Scatter。

6. AllGather：多对多的Gather，收集所有数据到所有节点上。从最基础的角度来看，All Gather相当于一个Gather操作之后跟着一个Broadcast操作。
   

7. 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 大模型训练的目标公式

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

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

1. 单卡速度：单卡速度既然是运算速度和数据IO的快慢来决定，那么就需要对单卡训练进行优化，于是主要的技术手段有精度训练、算子融合、梯度累加来加快单卡的训练性能。

2. 加速芯片数量：理论上，AI芯片数量越多，模型训练越快。但是，随着训练数据集规模的进一步增长，加速比的增长并不明显。如数据并行就会出现局限性，当训练资源扩大到一定规模时，由于通信瓶颈的存在，增加计算资源的边际效应并明显，甚至增加资源也没办法进行加速。这时候需要通讯拓扑进行优化，例如通过ring-all-reduce的通讯方式来优化训练模式。

3. 多卡加速比：多卡加速比既然由计算、通讯效率决定，那么就需要结合算法和集群中的网络拓扑一起优化，于是有了数据并行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)，即可实现多卡训练。

# 数据集的长度为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在训练中重复多次

1. Construction阶段：DDP需要额外的建立进程组阶段（Construction），用于确定通信协议和总进程数

   • 通信协议是DDP的底层基础，决定了并行训练的方式；
   • 总进程数（worldsize）：即有多少个独立的并行进程参与训练；

   根据需求每个进程可以占用一个或多个GPU，但并不推荐多个进程共享一个GPU，这会造成潜在的性能损失。为了便于理解，在本文的所有示例中我们假定每个进程只占用1个GPU，占用多个GPU的情况只需要简单的调整GPU映射关系就好。

   Construction只在训练开始前执行，因此不会带来额外的延迟

2. 初始状态广播
   在每个GPU上创建一个模型的副本（replica），GPU-1中模型的状态会被广播到所有其他进程，确保所有模型具有相同的初始状态。

3. 训练迭代：每个GPU接收其分配的数据批次，并在本地（每个GPU本身）计算损失函数的梯度（local gradients）。

4. 全体度求和（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。

5. 基于global gradients进行模型参数更新（weight update）

6. 重复步骤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 功能，整个测试结果为：

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对数据并行优化的三个阶段：

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

   在每个gpu中保存全部的参数和梯度，但是只保存1/Nd的优化器变量。通过将优化器状态进行分割，实现4倍的内存减少，同时保持与DP相同的通信量。

2. 梯度分割（ P o s + g P_{os+g} Pos+g）：
   每个gpu中只保存1/Nd的梯度，实现8倍的内存减少，并保持与DP相同的通信量。

3. 参数分割（ 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：

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

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

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

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

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

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

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 单节点并行化策略

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

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

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

6.2 多节点并行化策略

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