模型并行分布式训练Megatron (5) --Pipedream Flush-CSDN博客

本文链接：https://blog.csdn.net/luoganttcc/article/details/144645326

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

在流水线训练之中，如何给流水线各个阶段安排执行执行序列是一个关键，所以这里我们看看如何做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_interleaving
        else:
            forward_backward_func = forward_backward_pipelining_without_interleaving
    else:
        forward_backward_func = forward_backward_no_pipelining
    return 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 上的划分方式。这具有挑战性，原因有三：

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