分布式训练

原创已于 2024-07-03 18:59:59 修改 · 865 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#算法 #nlp #分布式 #python

于 2024-06-17 09:47:47 首次发布

部署运行你感兴趣的模型镜像

文章目录

- - 分布式训练

分布式训练

数据并行(DP & DDP)

DataParallel

DP 是较简单的一种数据并行方式，直接将模型复制到多个 GPU 上并行计算，每个 GPU 计算 batch 中的一部分数据，各自完成前向和反向后，将梯度汇总到主 GPU 上。其基本流程：

加载模型、数据至内存；
创建 DP 模型；
DP 模型的 forward 过程：
1. 一个 batch 的数据均分到不同 device 上；
2. 为每个 device 复制一份模型；
3. 至此，每个 device 上有模型和一份数据，并行进行前向传播；
4. 收集各个 device 上的输出；
每个 device 上的模型反向传播后，收集梯度到主 device 上，更新主 device 上的模型，将模型广播到其他 device 上；
3-4 循环。

在 DP 中，只有一个主进程，主进程下有多个线程，每个线程管理一个 device 的训练。因此，DP 中内存中只存在一份数据，各个线程间是共享这份数据的。DP 和 Parameter Server 的方式很像。

DistributedDataParallel

准备阶段
- 环境初始化：在各张卡上初始化进程并建立进程间通信，对应代码：init_process_group。
- 模型广播：将模型 parameter、buffer 广播到各节点，对应代码：model = DDP(model).to(local_rank)。
- 创建管理器 reducer，给每个参数注册梯度平均 hook。
准备数据
- 加载数据集，创建适用于分布式场景的数据采样器，以防不同节点使用的数据重叠。
训练阶段
- 前向传播
  - 同步各进程状态（parameter 和 buffer）；
  - 当 DDP 参数 find_unused_parameter 为 true 时，其会在 forward 结束时，启动一个回溯，标记未用到的参数，提前将这些设置为 ready。
- 计算梯度
  - reducer 外面：各进程各自开始反向计算梯度；
  - reducer 外面：当某个参数的梯度计算好了，其之前注册的 grad hook 就会触发，在 reducer 里把这个参数的状态标记为 ready；
  - reducer 里面：当某个 bucket 的所有参数都是 ready 时，reducer 开始对这个 bucket 的所有参数开始一个异步的 all-reduce 梯度平均操作；
  - reducer 里面：当所有 bucket 的梯度平均都结束后，reducer 把得到的平均梯度正式写入到 parameter.grad 里。
- 优化器应用梯度更新参数。

DDP 与 DP 的区别

DP DDP
多线程
1. 受到 GIL 的限制
2. 单机工作 多进程
1. 多机多卡
迭代更新 传输数据包括 梯度和参数
1. 全程维护 一个 optimizer
2 梯度汇总到主 GPU, 主 GPU 进行参数更新
3. 主 GPU Broadcast 参数给其他的 GPU 传输数据包括梯度
1. 每个进程具有 自己的 optimizer
2. 各进程自己计算梯度
3. Ring All-Reduce 将梯度汇总平均
4. 各进程用梯度来独立的更新参数
通信效率 通信成本随着 GPU 数量线性增长 Ring All-Reduce 通信成本恒定，与 GPU 数量无关

DDP 中由于各进程中的模型，初始参数一致 (初始时刻进行一次 broadcast)，而每次用于更新参数的梯度也一致，因此，各进程的模型参数始终保持一致。

	DP	DDP
	多线程 1. 受到 GIL 的限制 2. 单机工作	多进程 1. 多机多卡
迭代更新	传输数据包括梯度和参数 1. 全程维护一个 optimizer 2 梯度汇总到主 GPU, 主 GPU 进行参数更新 3. 主 GPU Broadcast 参数给其他的 GPU	传输数据包括梯度 1. 每个进程具有自己的 optimizer 2. 各进程自己计算梯度 3. Ring All-Reduce 将梯度汇总平均 4. 各进程用梯度来独立的更新参数
通信效率	通信成本随着 GPU 数量线性增长	Ring All-Reduce 通信成本恒定，与 GPU 数量无关

TP (Tensor Parallelism)

每个张量都被水平分成多个块，因此张量的每个分片都位于其指定的 GPU 上，而不是让整个张量驻留在单个 GPU 上。在处理过程中，每个分片在不同的 GPU 上分别并行处理，结果在步骤结束时同步。

MLP 的并行化：权重 A 矩阵竖着切 B 矩阵横着切最后 MERGE

对于输入 $X$ ，它的行数是批量大小 $B$ 乘以序列长度 $L$ ，列数是隐藏层的宽度即 $D$
其隐藏层的模块里面其实就是两个全连接层
假定第一个隐藏层的权重是 $A$ ( $[D, D^{'}]$ ， $D^{'}$ 一般是 $D$ 的 4 倍)，则先做矩阵乘法，然后再接一个激活函数比如 GELU (GELU 类似把 ReLU 的拐点顺滑下)
假定第二个隐藏层的权重是 $B$ ( $[D^{'}, D]$ )，最终 $\sigma(X \cdot A) B= Y$
如果是输入数据比较大，则优先选择做数据并行，即对输入做拆分
如果是模型本身比较大，则优先选择做模型并行，即对矩阵做拆分:
1. 对 $A$ 按行拆 ( $=\left[\begin{matrix} A_1 \\ A_2\end{matrix}\right]$ ) , 则相应的 $X$ 按列拆，导致的结果是两个 GPU 之间需要通讯。
2. 对 $A$ 按列拆 ( $=\left[\begin{matrix} A_1 , A_2\end{matrix}\right]$ ) ，则相应的 $X$ 按行拆，或者在两个 GPU 上都得有一份才行，此时则无需任何额外的通信。
使用第 2 种拆分方式，通过执行矩阵乘法得到 $XA_1$ 到 $XA_n$ ，它们可以独立输入 GeLU： $\left[Y_{1}, Y_{2}\right]=\left[\operatorname{GeLU}\left(X A_{1}\right), \operatorname{GeLU}\left(X A_{2}\right)\right]$ , 然后再得到 $n$ 个输出向量 $Y_1, Y_2, ⋯, Y_n$
同理，对 $B$ 按行拆分 ( $=\left[\begin{matrix} B_1 \\ B_2\end{matrix}\right]$ ) ，通过执行矩阵乘法得到 $Y_1B_1$ 到 $Y_nB_n$ ，即 $n$ 个输出向量 $Z_1, Z_2, ⋯, Z_n$ ，最终 merge 得到完整的 $Z$ 。
通过上述操作，我们可以更新任意深度的 MLP，只需在每个 拆列-拆行 序列之后同步 GPU

Self-Attention 的并行化：各个头各自计算

对于输入 $X$ ，它的行数是批量大小 $B$ 乘以序列长度 $L$ ，列数是隐藏层的宽度即 $D$ ，在自注意力机制中，输入 $X$ 会被复制成三份，分别对应为 $X$ 的 $Q$ 、 $K$ 、 $V$ 向量矩阵。
对于多头注意力，头的维度为 $D / h$ , 假定 $h = 2$ ，之后针对每个头中 $X$ 输入矩阵中各个单词的 $Q$ 向量，会与各自上下文的 $K$ 向量做缩放点积然后做 softmax 得到一个注意力分数或权重，之后再与 $V$ 相乘，得到一个 $[L, D / h]$ 的输出
每个头的计算是各自独立并行的，那意味着一个头可以放在 GPU 0 上，另一个头可以放在 GPU 1 上，最后 all reduce 每个头的结果

由于前向和后向传播中每层都有两个 all reduce，因此 TP 需要设备间有非常快速的互联。因此，除非你有一个非常快的网络，否则不建议跨多个节点进行 TP。实际上，如果节点有 4 个 GPU，则最高 TP 度设为 4 比较好。如果需要 TP 度为 8，则需要使用至少有 8 个 GPU 的节点

PP (Pipeline Parallelism)

模型在多个 GPU 上 垂直 (即按层) 拆分，因此只有一个或多个模型层放置在单个 GPU 上。每个 GPU 并行处理流水线的不同阶段，并处理 batch 的一部分数据

把网络分成 4 块，每一块放在一个 GPU 上，不同的颜色表示不同的 GPU)，于是就有了 F0、F1、F2、F3 这 4 个管级的前向路径，然后是 B3、B2、B1、B0 的逆序后向路径。
b 部分表示朴素 PP 方案： 在每个时间点只有一台设备在处理计算逻辑，完成计算后将结果发送给下一台设备。
c 部分是 PP 方法：
1. PP 引入了一个新的超参数来调整，称为 块 (chunks)。它定义了通过同一管级按顺序发送多少数据块。图中 $\text{chunks} = 4$ .
2. GPU 0 在 chunk 0、1、2 和 3 (F0,0、F0,1、F0,2、F0,3) 上执行相同的前向路径，然后等待，等其他 GPU 完成工作后，GPU 0 会再次开始工作，为块 3、2、1 和 0 (B0,3、B0,2、B0,1、B0,0) 执行后向路径。
  
  请注意，从概念上讲，这与梯度累积 (gradient accumulation steps，GAS) 的意思相同。PyTorch 叫它 块，而 DeepSpeed 叫它 GAS
3. 因为 块 (chunks），PP 引入了 micro-batches (MBS) 的概念。
  
  DP 将全局 batch size 拆分为小 batch size，因此如果 DP 度为 4，则全局 batch size 1024 将拆分为 4 个小 batch size，每个小 batch size 为 256 (1024/4)。而如果 块 (或 GAS) 的数量为 32，我们最终得到的 micro batch size 为 8 (256/32)。每个管级一次处理一个 micro batch。
  计算 DP + PP 设置的全局批量大小的公式为: $dp_degree ( 8 ∗ 32 ∗ 4 = 1024 ) \text{mbs}*\text{chunks}*\text{dp\_degree }(8*32*4=1024)$
4. 将 mini-batch 进一步划分成更小的 micro-batch，同时利用 pipipline 方案，每次处理一个 micro-batch 的数据，得到结果后，将该 micro-batch 的结果发送给下游设备，同时开始处理后一个 micro-batch 的数据，通过这套方案减小设备中的 Bubble(设备空闲的时间称为 Bubble)