1.1. 流水线并行

1.1.1. 数据并行与模型并行的对比

在分布式训练中,数据并行和模型并行是两种常见的并行技术。 数据并行训练中,每个 GPU 持有整个模型权重的副本,这虽然简单,却带来了参数冗余的问题。 为了解决这一问题,模型并行技术应运而生,它将模型分割并分布在一个设备阵列上,每个设备只保存模型的一部分参数。 模型并行进一步分为张量并行流水线并行,前者是层内并行,后者是层间并行,针对的是模型的不同 Transformer 层间进行分割。

1.1.2. 流水线并行的简介

流水线并行技术是为了应对单个模型过大,无法放置在单张 GPU 卡中的问题。 通过将模型的不同层放置到不同的计算设备上,可以降低单个计算设备的显存消耗,实现超大规模模型的训练。 在前向计算过程中,输入数据在设备间传递并逐层处理,最终得到结果。 反向传播过程类似,但传递的是梯度信息。 由于仅传输相邻设备间的输出张量,通信量相对较小。

1.1.3. 朴素流水线并行

朴素流水线并行是实现流水线并行的最直接方法,它将模型按层间切分成多个部分,并将每个部分分配给一个 GPU。 然而,这种方法存在明显的缺陷,

  • 在任意给定时刻,除了一个 GPU 之外的其他所有GPU都是空闲的,导致 GPU 使用率极低。
  • 在设备之间复制数据的额外通信开销,这进一步降低了效率。
  • 通信和计算没有交错的问题:当通过网络进行数据通信时,没有 GPU 执行任何操作。
  • 高内存需求:先执行前向传播的GPU(如:GPU1)将保留整个小批量缓存的所有激活,直到最后。如果批量大小很大,可能会产生内存问题。

1.1.4. 微批次流水线并行

为了解决 GPU 空闲问题,微批次流水线并行通过将传入的小批次分块为微批次,并人为创建流水线来提高设备利用率。 这种方法允许不同的 GPU 同时参与计算过程,显著提升了流水线并行设备的利用率,并减少了设备空闲状态的时间。

1.1.5. GPipe

GPipe 是谷歌提出的一种流水线并行方案,它基于 TensorFlow 和 PyTorch 实现。 GPipe 通过微批次流水线并行提高了模型训练的并行度,并利用重计算技术降低了显存消耗。 这种方法通过纵向对模型进行切分解决了单个设备无法训练大模型的问题,同时增加了多设备上的并行程度。

1.1.6. 流水线并行策略

流水线并行策略可以根据执行的策略分为 F-then-B 和 1F1B 两种模式。 F-then-B 模式先进行前向计算,再进行反向计算,可能会导致内存占用很高。 而 1F1B 模式,即 One Forward pass followed by One Backward pass,是一种前向计算和反向计算交叉进行的方式,可以及时释放不必要的中间变量,从而节省显存。

1.1.7. PipeDream

PipeDream 是微软 DeepSpeed 提出的 1F1B 策略,它通过合理安排前向和反向过程的顺序来解决内存过高的问题。 这种策略可以解决缓存 activation 的份数问题,使得 activation 的缓存数量只跟 stage 数相关,进一步节省显存,训练更大的模型。

1.1.8. PipeDream-2BW

PipeDream-2BW 是 PipeDream 的一个变体,它在流水线中只维护两个版本的模型权重,称为双缓冲权重。 这种方法减少了需要维护的权重版本总数,从而极大降低了内存的占用。

1.1.9. PipeDream-Flush(1F1B)

PipeDream-Flush 是 PipeDream 的一个变体,它维护单个权重版本并引入定期流水线刷新,以确保权重更新期间的权重版本保持一致。 虽然这种方法以执行性能为代价降低了峰值内存,但吞吐量较低。

1.1.10. 1F1B 调度模式

在使用 1F1B 策略时,存在两种调度模式:非交错式和交错式调度。 非交错式调度分为热身阶段、稳定阶段和完成后向计算阶段。 而交错式调度中,每个设备可以对多个层的子集进行计算,这种模式既节省内存又节省时间,但要求 micro-batch 的数量是流水线阶段的整数倍。

1.1.11. PipeDream(交错式1F1B)-Megatron-LM

Megatron-LM 基于 PipeDream-Flush 提出了交错式 1F1B 调度,即虚拟流水线。 这种方案通过在设备数量不变的情况下分出更多的流水线阶段,以更多的通信量换取流水线 Bubble 比率降低。 这种交错式调度需要额外的通信,但可以通过高速网络带宽来减少这种额外通信的影响。

1.1.12. 分布式训练框架流水线并行方案

在分布式训练框架中,流水线并行方案可以细分为同步流水线并行(Sync-PP)和异步流水线并行(Async-PP)。 PyTorch 采用的是 GPipe 方案,而 DeepSpeed 采用的是 PipeDream-Flush。 Megatron-LM 基于 PipeDream-Flush 进行了改进,提供了一种交错式 1F1B 方案。 Colossal-AI 基于 Megatron-LM 的交错式 1F1B 方案,提供了非交错和交错调度策略。


Ref: Pipeline Parallel

Copyright © 版权信息 all right reserved,powered by Gitbook该文件修订时间: 2024-12-20 14:33:58

results matching ""

    No results matching ""