分布式训练概述

分布式训练

开始之前首先贴几个链接：

背景

算力不足

单处理器的算力不足是促使人们设计分布式训练系统的一个主要原因。一个处理器的算力可以用每秒钟浮点数操作（Floating Point Operations Per Second，FLOPS）来衡量。根据摩尔定律（Moore’s Law），中央处理器的算力每18个月增长2倍,并且中央处理器的算力目前正处在瓶颈期，这一定律也面临着失效的风险。虽然计算加速卡(如GPU和TPU)针对机器学习计算提供了大量的算力，但机器学习模型正在快速发展，模型对于算力需求每18个月增长了5、6倍。 解决处理器性能和算力需求之间鸿沟的关键就在于利用分布式计算。通过大型数据中心和云计算设施，可以快速获取大量的处理器。通过分布式训练系统有效管理这些处理器，可以实现算力的快速增长，从而持续满足模型的需求。

内存不足

另外，训练机器学习模型需要使用到大量内存。训练机器学习模型需要大量内存。

举个例子：假设一个大型神经网络模型具有1000亿的参数，每个参数都由一个32位浮点数（4个字节）表达，存储模型参数就需要400GB的内存。在实际中，我们需要更多内存来存储激活值和梯度。假设激活值和梯度也用32位浮点数表达，那么其各自至少需要400GB内存，总的内存需求就会超过1200GB（即1.2TB）。而如今的硬件加速卡（如NVIDIA A100）仅能提供最高80GB的内存。

单卡内存空间的增长受到硬件规格、散热和成本等诸多因素的影响，难以进一步快速增长。因此，我们需要分布式训练系统来同时使用数百个训练加速卡，从而为千亿级别的模型提供所需的TB级别的内存。

系统架构

分布式训练的载体一般是大量用于分布式训练的服务器，并且依靠数据中心来进行管理。一个数据中心管理数百个集群，每个集群存在几百或者上千个服务器。 为了确保分布式训练系统的高效运行，需要首先估计系统计算任务的计算和内存用量。

一个模型训练任务（Model Training Job）往往会有一组数据（如训练样本）或者任务（如算子）作为输入，利用一个计算节点（如GPU）生成一组输出（如梯度）。

分布式执行一般具有三个步骤：

第一步将输入进行切分；
第二步将每个输入部分会分发给不同的计算节点，实现并行计算；
第三步将每个计算节点的输出进行合并，最终得到和单节点等价的计算结果。

这种首先切分，然后并行，最后合并的模式，本质上实现了分而治之（Divide-and-Conquer）的方法：由于每个计算节点只需要负责更小的子任务，因此其可以更快速地完成计算，最终实现对整个计算过程的加速。

实现方法

分布式训练系统的设计目标是：将单节点训练系统转换成等价的并行训练系统，从而在不影响模型精度的条件下完成训练过程的加速。。
single_node 一个单节点训练系统往往如图所示。一个训练过程会由多个数据小批次（mini-batch）完成。训练系统会利用数据小批次生成梯度，提升模型精度。这个过程由一个训练程序实现。在实际中，这个程序往往实现了一个多层神经网络的执行过程。该神经网络的执行由一个计算图（Computational Graph）表示。这个图有多个相互连接的算子（Operator），每个算子会拥有计算参数。每个算子往往会实现一个神经网络层（Neural Network Layer），而参数则代表了这个层在训练中所更新的的权重（Weights）。

在上篇文章中已经知道，为了更新参数，计算图的执行分为前向计算和反向计算两个阶段。前向计算的第一步会将数据读入第一个算子，该算子会根据当前的参数，计算出计算给下一个算子的数据。算子依次重复这个前向计算的过程（执行顺序：算子1，算子2，算子3），直到最后一个算子结束。最后的算子随之马上开始反向计算。反向计算中，每个算子依次计算出梯度（执行顺序：梯度3，梯度2，梯度1），并利用梯度更新本地的参数。反向计算最终在第一个算子结束。反向计算的结束也标志本次数据小批次的结束，系统随之读取下一个数据小批次，继续更新模型。

给定一个模型训练任务，人们会对数据和程序切分（Partition），从而完成并行加速。单节点训练系统可以被归类于单程序单数据模式。

而假如用户希望使用更多的设备实现并行计算，

一种方式是单程序多数据模式，也就是数据并行（data parallelism），首先可以选择对数据进行分区，并将同一个程序复制到多个设备上并行执行。
另一种并行方式是对程序进行分区（模型中的算子会被分发给多个设备分别完成）。这种模式是多程序单数据模式，常被称为模型并行（Model Parallelism）。
当训练超大型智能模型时，开发人员往往要同时对数据和程序进行切分，从而实现最高程度的并行。这种模式是多程序多数据模式，常被称为混合并行（Hybrid Parallelism）。

数据并行

数据并行是最常见的并行形式。在数据并行训练中，数据集被分割成几个碎片，每个碎片被分配到一个设备上。这相当于沿批次（Batch）维度对训练过程进行并行化。每个设备将持有一个完整的模型副本，并在分配的数据集碎片上进行训练。在反向传播之后，模型的梯度将会聚合（All Reduce），以便在不同设备上的模型参数能够保持同步。典型的数据并行实现：TensorFlow DistributedStrategy、PyTorch Distributed、Horovod DistributedOptimizer。

假定用户给定一个训练批次大小为\(N\),并且希望使用\(M\)个并行设备来加速训练。那么，该训练批次大小会被分为\(M\)个分区，每个设别会分配到\(\frac{N}{M}\)个训练样本。这些设备共享一个训练程序的副本，在不同数据分区上独立执行、计算梯度。不同的设备（假设设备编号为\(i\)）会根据本地的训练样本计算出梯度\(G_i\)。最后为了确保程序参数的一致性，本地梯度\(G_i\)需要进行聚合，并计算出平均梯度（\(\sum_{i=1}^N{G_i}/N\)）。 data parallelism

TIPS: 细节上，会有一个设备作为参数服务器来对每个训练设备中的梯度来进行累加，最后再广播到其他节点上。也可以将参数服务器分布在所有的节点上，每个训练设备只更新一部分梯度。

以下以pytorch中数据并行方法举例子。

数据并行(torch.nn.DataParallel)

这是Pytorch最早提供的一种数据并行方式，它基于单进程多线程进行实现的，它使用一个进程来计算模型权重，在每个批处理期间将数据分发到每个GPU。

DataParallel 的计算过程如下所示：

将 inputs 从主 GPU 分发到所有 GPU 上。
将 model 从主 GPU 分发到所有 GPU 上。
每个 GPU 分别独立进行前向传播，得到 outputs。
将每个 GPU 的 outputs 发回主 GPU。
在主 GPU 上，通过 loss function 计算出 loss，对 loss function 求导，求出损失梯度。
计算得到的梯度分发到所有 GPU 上。
反向传播计算参数梯度。
将所有梯度回传到主 GPU，通过梯度更新模型权重。不断重复上面的过程。但是它的缺点也很明显：

单进程多线程带来的问题：DataParallel使用单进程多线程进行实现的，方便了信息的交换，但受困于 GIL，会带来性能开销，速度很慢。而且，只能在单台服务器（单机多卡）上使用（不支持分布式）。

The Python Global Interpreter Lock or GIL, in simple words, is a mutex (or a lock) that allows only one thread to hold the control of the Python interpreter.

效率问题，主卡性能和通信开销容易成为瓶颈，GPU 利用率通常很低：数据集需要先拷贝到主进程，然后再分片（split）到每个设备上；权重参数只在主卡上更新，需要每次迭代前向所有设备做一次同步；每次迭代的网络输出需要聚集到主卡上。因此，通信很快成为一个瓶颈。除此之外，这将导致主卡和其他卡之间，GPU利用率严重不均。

分布式数据并行（PyTorch DDP）

分布式数据并行(torch.nn.DistributedDataParallel)，是基于多进程进行实现的，每个进程都有独立的优化器，执行自己的更新过程。每个进程都执行相同的任务，并且每个进程都与所有其他进程通信。进程之间只传递梯度，因此网络通信就不再是瓶颈。分布式数据并行

具体流程为：

首先将 rank=0 进程中的模型参数广播到进程组中的其他进程；
然后，每个 DDP 进程都会创建一个 local Reducer 来负责梯度同步。
在训练过程中，每个进程从磁盘加载 batch 数据，并将它们传递到其 GPU。每个 GPU 都有自己的前向过程，完成前向传播后，梯度在各个 GPUs 间- 进行 All-Reduce，每个 GPU 都收到其他 GPU 的梯度，从而可以独自进行反向传播和参数更新。
同时，每一层的梯度不依赖于前一层，所以梯度的 All-Reduce 和后向过程同时计算，以进一步缓解网络瓶颈。
在后向过程的最后，每个节点都得到了平均梯度，这样各个 GPU 中的模型参数保持同步。

DistributedDataParallel方式可以更好地进行多机多卡运算，更好的进行负载均衡，运行效率也更高，虽然使用起来较为麻烦，但对于追求性能来讲是一个更好的选择。

完全分片数据并行(Pytorch FSDP)

通常来说，在模型训练的过程中，GPU上需要进行存储的参数包括了模型本身的参数、优化器状态、激活函数的输出值、梯度以及一些临时的Buffer。分布图
可以看到模型参数仅占模型训练过程中所有数据的一部分，当进行混合精度运算时，其中模型状态参数（优化器状态 + 梯度+ 模型参数）占到了一大半以上。

针对模型状态的存储优化（去除冗余），DeepSpeed 提出了 ZeRO，ZeRO 使用的方法是分片，即每张卡只存 1/N 的模型状态量，这样系统内只维护一份模型状态参数。

ZeRO对模型状态（Model States）参数进行不同程度的分割，主要有三个不同级别： 1. ZeRO-1 : 优化器状态分片（ Optimizer States Sharding） 2. ZeRO-2 : 优化器状态与梯度分片（Optimizer States & Gradients Sharding） 3. ZeRO-3 : 优化器状态、梯度和模型权重参数分片（Optimizer States & Gradients & Parameters Sharding）

完全分片数据并行(torch.distributed.fsdp.FullyShardedDataParallel)，是Pytorch最新的数据并行方案。Pytorch DDP用起来简单方便，但是要求整个模型加载到一个GPU上，这使得大模型的训练需要使用额外复杂的设置进行模型分片。为了打破模型分片的障碍（包括模型参数，梯度，优化器状态），同时仍然保持数据并行的简单性，FSDP因此被提出。

FSDP是一种新型数据并行训练方法，但与传统的数据并行不同，传统的数据并行维护模型参数、梯度和优化器状态的每个 GPU 副本，而 FSDP 将所有这些状态跨数据并行工作线程进行分片，并且可以选择将模型参数分片卸载到 CPU。 FSDP工作流程

Model shard：每个GPU上仅存在模型的分片。
All-gather：每个GPU通过all-gather从其他GPU收集所有权重，以在本地计算前向传播。
Forward（local）：在本地进行前向操作。前向计算和后向计算都是利用完整模型。
All-gather：然后在后向传播之前再次执行此权重收集。
Backward（local）：本地进行后向操作。前向计算和后向计算都是利用完整模型，此时每个GPU上也都是全部梯度。
Reduce-Scatter：在向后传播之后，局部梯度被聚合并且通过 Reduce-Scatter 在各个GPU上分片，每个分片上的梯度是聚合之后本分片对应的那部分。
Update Weight（local）：每个GPU更新其局部权重分片。

通常，模型层以嵌套方式用 FSDP 包装，因此，只有单个 FSDP 实例中的层需要在前向或后向计算期间将完整参数收集到单个设备。计算完成后，收集到的完整参数将立即释放，释放的内存可用于下一层的计算。当实例在计算中不活动时，FSDP 可以将参数、梯度和优化器状态卸载到 CPU。

解释ZeRO/FSDP的关键是我们可以把DDP之中的All-Reduce操作分解为独立的 Reduce-Scatter 和 All-Gather 操作。 all_reduece

All-Reduce 是 Reduce-Scatter 和 All-Gather 的组合。标准 All-Reduce 操作可以分解为两个单独的阶段。

Reduce-Scatter 阶段，在每个GPU上，会基于 rank 索引对 rank 之间相等的块进行求和。
All-Gather 阶段，每个GPU上的聚合梯度分片可供所有GPU使用。

通过重新整理 Reduce-Scatter 和 All-Gather，每个 DDP worker只需要存储一个参数分片和优化器状态。

模型并行

模型并行分为张量并行和流水线并行，张量并行为层内并行，对模型 Transformer 层内进行分割；流水线并行为层间并行，对模型不同的 Transformer 层间进行分割。两种方式

模型并行往往用于解决单节点内存不足的问题。一个常见的内存不足场景是模型中含有大型算子，例如深度神经网络中需要计算大量分类的全连接层。完成这种大型算子计算所需的内存可能超过单设备的内存容量。那么需要对这个大型算子进行切分。

假设这个算子具有\(P\)个参数，而系统拥有\(N\)个设备，那么可以将\(P\)个参数平均分配给\(N\)个设备（每个设备分配\(P/N\)个参数），从而让每个设备负责更少的计算量，能够在内存容量的限制下完成前向计算和反向计算。这也就是层内并行的方式，也叫算子内并行/流水线并行。

假设一个神经网络具有两个算子，算子1的计算（包含正向和反向计算）需要预留16GB的内存，算子2的计算需要预留1GB的内存。而本例中的设备最多可以提供10GB的内存。为了完成这个神经网络的训练，需要对算子1实现并行。具体做法是，将算子1的参数平均分区，设备1和设备2各负责其中部分算子1的参数。由于设备1和设备2的参数不同，因此它们各自负责程序分区1和程序分区2。在训练这个神经网络的过程中，训练数据（按照一个小批次的数量）会首先传给算子1。由于算子1的参数分别由两个设备负责，因此数据会被广播（Broadcast）给这两个设备。不同设备根据本地的参数分区完成前向计算，生成的本地计算结果需要进一步合并，发送给下游的算子2。在反向计算中，算子2的数据会被广播给设备1和设备2，这些设备根据本地的算子1分区各自完成局部的反向计算。计算结果进一步合并计算回数据，最终完成反向计算。

另一种内存不足的场景是：模型的总内存需求超过了单设备的内存容量。在这种场景下，假设总共有\(N\)个算子和\(M\)个设备，可以将算子平摊给这\(M\)个设备，让每个设备仅需负责\(N/M\)个算子的前向和反向计算，降低设备的内存开销。这也就是层间并行，也称为张量并行/算子间并行。

流水线并行

所谓流水线并行，就是将模型的不同层放置到不同的计算设备，降低单个计算设备的显存消耗，从而实现超大规模模型训练。

如图所示，模型共包含四个模型层（如：Transformer层），被切分为三个部分，分别放置到三个不同的计算设备。即第 1 层放置到设备 0，第 2 层和第三 3 层放置到设备 1，第 4 层放置到设备 2。
流水线并行

具体地讲，前向计算过程中，输入数据首先在设备 0 上通过第 1 层的计算得到中间结果，并将中间结果传输到设备 1，然后在设备 1 上计算得到第 2 层和第 3 层的输出，并将模型第 3 层的输出结果传输到设备 2，在设备 2 上经由最后一层的计算得到前向计算结果。反向传播过程类似。最后，各个设备上的网络层会使用反向传播过程计算得到的梯度更新参数。由于各个设备间传输的仅是相邻设备间的输出张量，而不是梯度信息，因此通信量较小。

根据流水线的设计不同，又可以进一步分为朴素流水线并行和微批次流水线并行等。 ##### 朴素流水线朴素流水线并行是实现流水线并行训练的最直接的方法。我们将模型按照层间切分成多个部分（Stage），并将每个部分（Stage）分配给一个 GPU。然后，我们对小批量数据进行常规的训练，在模型切分成多个部分的边界处进行通信。

朴素流水线并行朴素流水线存在最大的问题就是其会产生特别多的空泡，主要是因为该方案在任意给定时刻，除了一个 GPU 之外的其他所有 GPU 都是空闲的。因此，如果使用 4 个 GPU，则几乎等同于将单个 GPU 的内存量增加四倍，而其他资源 (如计算) 相当于没用上。朴素的流水线并行将会导致GPU使用率过低。 ##### 微批次流水线微批次（MicroBatch）流水线并行与朴素流水线几乎相同，但它通过将传入的小批次（minibatch）分块为微批次（microbatch），并人为创建流水线来解决 GPU 空闲问题，从而允许不同的 GPU 同时参与计算过程，可以显著提升流水线并行设备利用率，减小设备空闲状态的时间。
微批次流水线

Gpipe就是谷歌以其为基础提出的一种流水线并行方案，可以通过纵向对模型进行切分解决了单个设备无法训练大模型的问题；同时，又通过微批量流水线增加了多设备上的并行程度，除此之外，还使用re-materialization（重计算）降低了单设备上的显存峰值 #### 流水线并行策略流水线并行根据执行的策略，又可以分为两种模式：F-then-B和1F1B模式。 ##### F-then-B 策略 F-then-B模式，即先进行前向计算，再进行反向计算。但由于其缓存了多个micro-batch的中间变量和梯度，显存的实际利用率并不高。 F-then-B策略 ##### 1F1B策略 1F1B（One Forward pass followed by One Backward pass）模式，是一种前向计算和反向计算交叉进行的方式。在 1F1B 模式下，前向计算和反向计算交叉进行，可以及时释放不必要的中间变量。

1F1B策略 1F1B 方式相比于 F-then-B 方式，峰值显存可以节省 3/8，对比朴素流水线并行峰值显存明显下降，设备资源利用率显著提升。 #### 张量并行将计算图中的层内的参数（张量）切分到不同设备（即层内并行），每个设备只拥有模型的一部分，以减少内存负荷，也就是所谓的张量模型并行。

从数学原理上来看就是对于linear层就是把矩阵分块进行计算，然后把结果合并；对于非linear层，则不做额外设计。

张量切分方式可以分为按行切分和按列切分，对应于行并行（row parallelism）和列并行（column parallelism）.

行并行就是把权重 A 按照行分割成两部分。为了保证运算，同时我们也把 X 按照列来分割为两部分，具体如下所示： \[ X A = \begin{bmatrix} X_1&X2 \end{bmatrix}\begin{bmatrix} A_1\\ A_2 \end{bmatrix} =X_1A_1+X_2A_2=Y_1+Y_2=Y \] 这样,将 X1 和 A1 就可以放到 GPU0 之上计算得出 Y1，X2 和 A2 可以放到第二个 GPU1 之上计算得出 Y2，然后，把Y1和Y2结果相加，得到最终的输出Y。

列并行就是把 A按照列来分割，具体示例如下： \[ X A=\begin{bmatrix} X \end{bmatrix}\begin{bmatrix} A_1 & A_2 \end{bmatrix}=\begin{bmatrix} XA_1 & XA_2\end{bmatrix} =\begin{bmatrix} Y_1 & Y_2\end{bmatrix} =Y \] 这样，将 X 分别放置在GPU0 和GPU1，将 A1 放置在 GPU0，将 A2 放置在 GPU1，然后分别进行矩阵运行，最终将2个GPU上面的矩阵拼接在一起，得到最终的输出Y。