[源码解析] 模型并行分布式训练Megatron (5) --Pipedream Flush

文章目录

[源码解析] 模型并行分布式训练Megatron (5) --Pipedream Flush
- 0x00 摘要
- 0x01 背景
- 0x02 论文
- - 2.1 引论
  - 2.2 背景
  - 2.3 流水线权重问题
  - - 2.3.1 问题1
    - 2.3.2 问题2
    - 2.3.3 问题3
  - 2.4 PipeDream-2BW 系统设计
  - - 2.4.1 GPipe
    - 2.4.2 Double-Buffered Weight Updates (2BW)
    - 2.4.2. Weight Updates with Flushes (PipeDream-Flush)
- 0x03 PipeDream-Flush 实现
- - 3.1 总体思路
  - - 3.1.1 缺省计划
    - 3.1.2 PipeDream计划
  - 3.2 启动阶段
  - 3.3 热身阶段
  - 3.4 通信模块
  - - 3.4.1 基础通信方法
    - 3.4.2 API
    - 3.4.3 流水线上下游
  - 3.5 稳定阶段
  - - 3.5.1 逻辑
    - 3.5.2 串行
  - 3.6 冷却阶段
  - 3.7 Flush 体现在哪里？
- 0xEE 个人信息
- 0xFF 参考

0x00 摘要

NVIDIA Megatron 是一个基于 PyTorch 的分布式训练框架，用来训练超大Transformer语言模型，其通过综合应用了数据并行，Tensor并行和Pipeline并行来复现 GPT3，值得我们深入分析其背后机理。本系列有 5 篇文章，通过论文和源码和大家一起学习研究。本文将看看 Megatron 如何给流水线各个阶段安排执行执行序列。

本系列其他文章为：

[源码解析] 模型并行分布式训练Megatron (1) — 论文 & 基础

[源码解析] 模型并行分布式训练Megatron (2) — 整体架构

[源码解析] 模型并行分布式训练 Megatron (3) —模型并行实现

[源码解析] 模型并行分布式训练 Megatron (4) — 如何设置各种并行

0x01 背景

在流水线训练之中，如何给流水线各个阶段安排执行执行序列是一个关键，所以这里我们看看如何做schedule。

对于 Megatron 来说，在训练时候，get_forward_backward_func 获取pipeline 的schedule，这里分为 flush 和 interleaving 两种，因为时间所限，我们只分析 flush 的schedule，有兴趣的读者可以自行研究 interleaving。

def get_forward_backward_func():args = get_args()if mpu.get_pipeline_model_parallel_world_size() > 1:if args.virtual_pipeline_model_parallel_size is not None:forward_backward_func = forward_backward_pipelining_with_interleavingelse:forward_backward_func = forward_backward_pipelining_without_interleavingelse:forward_backward_func = forward_backward_no_pipeliningreturn forward_backward_func

概括来说，Megatron 是基于 PipeDream-2BW 之上实现了定期刷新。

PipeDream-2BW 在流水线之中维护了两个版本的模型权重，“2BW” 是双缓冲权重（double-buffered weights）”，PipeDream-2BW 会为每个微批次生成一个新的模型版本K（K>d），但是因为有些剩余后向传递仍然依赖于旧版本模型，所以新的模型版本无法立即取代旧版本，但是由于只保存了两个版本，所以极大降低了内存占用。
PipeDream-flush 则在 PipeDream-2BW 之上添加了一个全局同步的流水线更新刷新操作，思路类似 GPipe。这种方法通过牺牲吞吐量的能力部分下降的代价来减少了内存占用（即只维护一个版本的模型权重）。

0x02 论文

Memory-Efficient Pipeline-Parallel DNN Training 和 Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM 是 Megatron 对应的相关论文，我们就从论文开始研究。注：下面论述内容是基于原论文发表时间，因为各个开源系统也在演进，所以其针对其他开源系统的论述在今天看来不一定完全正确。

2.1 引论

近来，一些工作提出了流水线模型并行以加速模型并行训练。例如 GPipe（Huang等人，2019年）和PipeDream（Harlap等人 2018年；Narayanan等人，2019年）把多个输入顺序推送到一系列worker之中来训练，每个worker负责一个模型分区，这允许不同worker并行处理不同的输入。

由于特定输入的向前和向后传播之间的权重版本不一致，Native 流水线可能会造成模型不收敛，现有技术权衡了内存占用和吞吐量，以不同的方式来避免这种情况。
GPipe维护单一权重版本，但是会定期进行流水线刷新（图1a），具体刷新时间是在流水线训练完输入要更新权重时候，由于资源空闲，这些刷新限制了总体吞吐量。
PipeDream不会定期Flush流水线，但会存储多个权重版本，这增加了吞吐量，但也增加了内存占用，由于内存限制，无法训练大型模型。

所以，有效地训练大型模型需要一种同时具有高吞吐量和低内存占用的方法。此外，流水线并行系统的性能取决于DNN模型operators在 worker 上的划分方式。这具有挑战性，原因有三：

内存容量限制：与模型分区相关的参数和中间激活需要能够放置在加速器的主设备内存之中。
异构网络互连：如今的训练部署具有异构网络拓扑性，同一服务器上的设备之间具有更高的带宽链路。
运算符如何放置的大搜索空间：随着模型尺寸的增加，拆分运算符图在计算上变得非常昂贵，因为不同分区方式数量是指数级的。

图1a。不同流水线并行执行的timeline。后向传播的时间假定为向前传播的两倍；向前传播以蓝色显示，后向传播以绿色显示。数字表示微批次ID，时间沿x轴显示，每个worker的利用率沿y轴显示。GPipe维护单一权重版本，但会定期刷新flush流水线。PipeDream不引入周期性流水线刷新，但维护多个权重版本。

在论文中，作者介绍了PipeDream-2BW，一个高效的DNN模型流水线并行训练系统。PipeDream-2BW通过两个关键贡献实现了高吞吐和低内存占用。首先，作者提出了双缓冲权重更新（2BW），这是一种在避免流水线刷新的同时减少训练内存占用的技术。

作者利用了这样一个事实，即每个输入生成的梯度不需要立即应用于权重，而是可以累积为“合并（coalesced）”梯度，以限制保留的权重版本的数量。2BW没有在使用最近更新的权重之前刷新流水线，而是将新权重用于新进入流水线的输入，同时将以前的权重版本（称为阴影版本）用于已在训练中的输入（in-flight inputs）。

每个 worker 的权重双缓冲产生了一种流水线方案，其吞吐量高于GPipe（无流水线刷新），内存效率高于PipeDream（这里是2个权重版本，而在PipeDream 的一个depth-d流水线中，最差的情况是d个权重版本）。

作者还介绍了2BW的一种变体（称为PipeDream Flush），它在吞吐量上进行权衡，以获得更低的内存占用和更高的性能。

2.2 背景

在本节中，作者简要概述DNN模型分布式训练的相关技术。

数据并行。
- 数据并行用于扩展模型训练。利用数据并行性（Xing等人，2015），每个worker都有整个模型的副本，输入数据集在worker之间分片。Worker 定期聚合他们的梯度，以确保所有 Worker 都看到一致版本的权重。数据并行性不能训练没法放入单个worker的大型模型，但可以用于较小的模型分区（model partitions）。
- 数据并行scale-out通常工作良好，但存在两个限制：a）超出某一点，每个GPU的batch size变得太小，降低了GPU利用率并增加了通信成本；b）可以使用的最大设备数量是batch size，这限制了可用于训练的加速器数量。于是人们提出了各种模型并行技术来解决这两个挑战。
模型并行。对于不适合单个worker的大型模型，一般来说使用模型并行训练。
- 利用模型并行性（Dean等人，2012年；Chilimbi等人，2014年），模型中的权重参数在可用worker上进行分割（每个transformer层内的矩阵乘法在多个GPU上分割），worker之间交流中间激活和梯度。层间模型并行性未充分利用资源，因为在任何时间点最多只有一个工作进程处于活动状态。
- Tensor（intra-layer）模型并行性（Shoeybi et al.，2019）容易导致关键路径中all-to-all通信过于昂贵。因为 tensor 并行所需的all-reduce通信需要通过服务器间链路，这比多GPU服务器中可用的高带宽NVLink要慢，于是容易将模型分区的数量限制为单个服务器中的GPU数量。而且高度的模型并行会创建大量的小矩阵乘法 (GEMMs)，这可能会降低GPU的利用率。
- FlexFlow（Jia et al.，2018）展示了如何使用模型和数据并行性拆分模型图，但在使用模型并行性时仍然存在资源利用率低的问题。
流水线并行。为了解决模型并行性的缺点，最近的工作如PipeDream和GPipe提出了流水线并行性。
- 通过流水线并行，将多个输入（而不是1个）注入到由层间（inter-layer）模型分区组成的训练中。这确保了计算资源得到更好的利用。
- 一个批（batch）被分割成更小的微批（microbatches），并在这些微批之间以流水线方式执行。可以用各种方式将层分配给worker，并且输入的向前和向后传播可以使用各种不同计划。
- 然而，简单的流水线可能会导致特定输入的前后传递之间的权重版本不匹配。具体来说，如果立刻用最新的权重版本来进行权重更新，那么在流水线之中，一个输入可能会看到的是向后传播更新的权重，而不是它在向前传播时候看到的权重，从而导致不正确的梯度计算。
- 层分配和调度策略导致不同的性能权衡。不管计划如何，为了保持严格的优化器语义，优化器步骤需要跨设备同步，从而在每个批处理结束时进行流水线刷新，允许微批处理完成执行（此时不加入新的微批）用来Flush流水线的时间最高可以达到50%，这取决于注入流水线的微批次数量。微批次数量与流水线尺寸的比率越大，流水线Flush所花费的时间越短。因此，为了实现高效率，通常需要更大的batch size。

用户可以使用各种技术训练他们的大型模型，每种技术都有不同的权衡。此外，这些技术可以结合使用。然而，结合这些技术会产生复杂（non-trivial）的交互，为了获得良好的性能，需要仔细地进行推理，才能做到在保持严格的优化器语义的同时最大化给定batch size的大型模型训练吞吐量。

要实现大规模的吞吐量，需要沿着多个轴进行创新和精心设计：高效的内核实现使大部分训练受计算限制，而不是内存限制；应该在设备上对计算图进行智能分区，以减少通过网络链路发送的字节数，同时限制设备空闲时间；使用特定领域的通信优化和快速硬件（最先进的GPU以及相同和不同服务器上GPU之间的高带宽链路）。

此外，论文作者还研究了影响吞吐量的各种成分之间的相互作用，包括经验和分析。基于这些研究，论文作者就如何配置分布式培训提供以下指导原则：

不同形式的并行以复杂的方式相互作用：并行化策略影响通信量、执行内核的计算效率，以及 Worker 因流水线刷新（流水线气泡）而等待计算的空闲时间。例如，张量和流水线模型并行性的次优组合可以导致高达2×更低的吞吐量，即使服务器之间的网络链路带宽较高；张量模型并行性在多GPU服务器中是有效的，但流水线模型并行性必须用于更大的模型。
用于流水线并行性的计划会影响通信量、流水线气泡大小以及用于存储激活的内存。
- 超参数的值（如微批次大小）会影响内存占用、在辅助进程上执行的内核的算术效率以及流水线气泡大小。微批次大小的最佳值取决于具体问题，一个合适取值可以将吞吐量提高15%。
- 分布式培训是通信密集型的。如果节点间互连较慢或更多的通信密集型分区将阻碍性能的扩展。
- 论文没有研究如何自动探索并行策略的搜索空间（如FlexFlow、PipeDream、Tarnawski 和DAPPLE），而是建议使用那些在实践中效果良好的启发式方法。

2.3 流水线权重问题

我们这里回顾一下流水线权重问题。下图是朴素流水线执行情况，本质上是一种 async SGD。

2.3.1 问题1

遇到的第一个问题是：在一般情况下，当计算第二个迭代时候，我们需要基于第一个迭代更新之后的模型来计算。但是如下图所示，对于机器 1，当第二轮迭代开始时候（红色圆圈的深蓝色2号），第一轮迭代的反向传播（浅绿色1号格）还没有开始。

2.3.2 问题2

第二个问题是：对于机器2，当它进行第5个mini-batch的前向传播时候（第二行蓝色5），它基于更新两次的权重来进行前向计算（第二行蓝色5之前有两个绿色格子，意味着权重被更新了两次）。

但进行第5个mini-batch的反向传播（第二行浅绿色5）时候，用到的权重是更新了4次的（第二行前面浅绿色的1,2,3,4一共会更新权重4次）。这与单节点深度学习假设冲突，会导致训练效果下降。

PipeDream 作者为了解决这些问题，提出了 Weight Stashing，以确保相同输入的向前和后向传播中使用相同的权重版本（原论文图1b）。具体就是每个机器多备份几个版本的权重，前向传播用哪个权重计算，反向传播还用这个权重计算。

就上图来说，机器1需要保存4个版本的权重，机器2需要保存3个版本的权重，机器3需要保存2个版本的权重，机器4需要保存1个版本的权重。在最坏的情况下，储存的权重版本总数是d，其中d是流水线深度，这对于大型模型来说内存占用太高了。使用PipeDream的默认权重更新语义，每个阶段的权重更新都有不同的延迟项，并且不会在流水线内执行累积。

2.3.3 问题3

另外一个问题是：现在做前向传播时候，每个机器计算时候，其基于的权重被更新的不同次数，比如第5个mini-batch（深蓝色的5），在机器 1 计算 5 时候，基于的权重是更新一次的（其前面有一个绿色），但是机器 2 计算 5 时候，基于的权重是更新两次的（其前面有两个绿色）。

解决思路是：每次前向传播时候，每个机器基于更新最少的权重来计算，比如对于机器2，就忽略绿色2更新的权重，对于机器3，就忽略绿色2,3两次更新之后的权重，它们都使用被绿色1更新一次之后的权重（图上矩形框黄色 1 ）。

2.4 PipeDream-2BW 系统设计

PipeDream-2BW使用内存高效的流水线并行性来训练不适合单个加速器的大型模型。它的双缓冲权重更新（2BW）和刷新机制确保了高吞吐量、低内存占用和类似于数据并行的权重更新语义。PipeDream-2BW将模型拆分为多个Worker上的多个阶段，并对每个阶段进行相同次数的复制（在同一阶段的副本之间进行数据并行更新）。这种平行流水线适用于每层重复固定次数的模型（例如transformer模型）。

2.4.1 GPipe

GPipe维护模型权重的单一版本。输入批次被分成更小的微批次。权重梯度是累积的，不会立即应用，并且定期flush 流水线，以确保不需要保持多个权重版本。GPipe提供了类似于数据并行的权重更新语义。原论文图1a显示了GPipe执行的时间线。周期性流水线Flush可能会很昂贵，从而限制吞吐量。缓解这一开销的一种方法是在流水线内进行额外的累积，但这并不总是切实可行的：a）在large scale factors下，能支持的最小batch size较大（与scale factor成比例），且大批量会影响所有模型的收敛性，b）GPipe需要保持与批大小成比例的激活存储。

2.4.2 Double-Buffered Weight Updates (2BW)

PipeDream-2BW结合1F1B调度（Narayanan等人，2019年）使用了一种新颖的双缓冲权重更新（2BW）方案，其中每个 worker 在不同输入的向前和向后传递之间交替，以确保在特定输入的向前和向后传递中使用相同的权重版本（论文原图2）。2BW的内存占用比PipeDream和GPipe低，并且避免了GPipe昂贵的流水线刷新。

梯度是以较小的mi-crobatches粒度计算的。对于任何输入微批次，PipeDream-2BW对输入的向前和向后传播使用相同的权重版本。在以批的粒度应用更新之前，会在多个微批次上累积更新，从而限制生成和维护的权重版本的数量。图2显示了2BW的时间线示例。

PipeDream-2BW 为每m个微批次生成一个新的权重版本（m≥ d, d是流水线深度）。为了简单起见，作者首先假设m=d（图2中的d=4）。新权重版本不能立即使用。特别是，进行中的输入（in-flight）不能使用最新的权重版本进行向后传播（例如，在t=21时， worker 3上的输入7），因为这些输入的向前传递已在不同阶段使用较旧的权重版本启动。

因此，新生成的权重版本需要缓冲以备将来使用。但是，需要维护的权重版本总数最多为2，因为用于生成新权重版本的权重版本可以立即丢弃（通过该阶段的未来输入不再使用旧权重版本）。例如，在图2中，每个 worker 在处理完输入8 的 backward pass后都可以丢弃W(0)，因为所有后续输入的前向传递和后向传递都使用更高的权重版本。

给定输入微批次k（基于1开始的索引）使用的权重版本为 $m a x (⌊ (k - 1) / m ⌋ - 1, 0)$ ，其中m是批次中的微批次数（图2中的4）。对于输入k的向前和向后传播，此权重版本相同。m可以是任何 ≥ d 的数字，额外的梯度累积（较大的m）会增加全局 batch size。

论文原图2。时间轴显示PipeDream-2BW的双缓冲权重更新 (2BW) 方案，时间轴沿x轴进行。在不丧失通用性的情况下，假设向后传播的时间是向前传播的两倍。PipeDream-2BW在每个worker上只存储两个权重版本，减少了总内存占用，同时不再需要昂贵的流水线暂停。 $W_i^{(v)}$ 表示worker i上的具有版本v的权重（包含从输入v生成的权重梯度）。在方格绿色框中会生成新的权重版本； $W_4^{(4)}$ 首先用在输入9的向前传播之中。

上图中的 Before 意思是丢弃版本之前系统的两个权重buffer，After 意思是做了丢弃动作之后系统的两个权重buffer。

2.4.2. Weight Updates with Flushes (PipeDream-Flush)

作者还提出了第二个内存高效的流水线计划，称为PipeDream Flush。它的内存占用比2BW和vanilla优化器语义更低，但以较低的吞吐量为代价。该计划重用PipeDream（Narayanan等人，2019年）的1F1B计划，但保持单一权重版本，并引入定期流水线刷新，以确保权重更新期间的一致权重版本。具有两个流水线阶段的PipeDream-Flush和GPipe的时间表如图3所示。

为何要选择 1F1B？因为它将n-flight microbatches 数量缩减到流水线深度 d，而不是GPipe的微批次数目 m，所以 1F1B 是memory-efficient。为了降低bubble time，一般来说，m >> d。

内存占用。使用PipeDream Flush，in-flight “活动”输入激活的总数小于或等于流水线深度，这使其内存占用比GPipe低，GPipe必须保持输入激活与梯度平均（m）的微批次数量成比例。PipeDream Flush的内存占用也低于PipeDream-2BW，因为它只需要维护一个权重版本（而PipeDream-2BW只需要2个）。

Semantics。定期流水线刷新确保可以使用最新权重版本计算的梯度执行权重更新。这将使权重更新用如下方式进行： $W^{(t+1)} = W^{(t)} − ν · ∇f(W^{(t)})$ 。

论文原图3。GPipe和PipeDream的时间表将分为两个阶段。GPipe和PipeDream Flush都使用管道刷新；PipeDream-Flush在稳定状态下交替进行向前和向后传播，通过限制仅保留进行中（in-flight）微批次的激活来保持较低的内存占用。

0x03 PipeDream-Flush 实现

我们前面提到，当没有设置 virtual_pipeline_model_parallel_size 时候，就是使用 Flush 方法得到流水线schedule，具体函数是 forward_backward_pipelining_without_interleaving。

def get_forward_backward_func():args = get_args()if mpu.get_pipeline_model_parallel_world_size() > 1:if args.virtual_pipeline_model_parallel_size is not None:forward_backward_func = forward_backward_pipelining_with_interleavingelse:# Flush scheduleforward_backward_func = forward_backward_pipelining_without_interleavingelse:forward_backward_func = forward_backward_no_pipeliningreturn forward_backward_func

为何要选择 1F1B？论文作者提到，因为它将in-flight microbatches 数量缩减到流水线深度 d，而不是GPipe的微批次数目 m，所以 1F1B 是memory-efficient。为了降低bubble time，一般来说，m >> d。

3.1 总体思路

3.1.1 缺省计划

GPipe提出了一个执行计划，其中首先执行一个批次中所有微批次的正向传播，然后执行所有微批次的反向传播（如图3所示）。我们可以量化GPipe流水线气泡的大小(