导言
ICLR，全称为 International Conference on Learning Representations (国际学习表征会议) 是三大机器学习领域顶会之一 (另外两个是ICML和NeuriPS)。该会议的主要创办者就包含了深度学习三大巨头的YoShua Bengio和Yann LeCun。ICLR的第一届是2013年举行，短短几年就得到的快速发展，逐步成为机器学习领域的顶会。我们这项工作就是ICLR 2022收录论文。
近年来，使用大规模深度神经网络 (DNN) 来提高性能的趋势，推动了用于高效 DNN 训练的并行流水线技术的发展，也催生出了一系列优秀的的Pipeline Parallelism(PP)技术，例如 GPipe、PipeDream 和 PipeDream- 2BW。然而，目前领先的 PipeDream-2BW 技术仍然存在两个主要缺点，即过多的内存冗余和权重更新延迟。在这个方向上，我们提出了一种名为 WPipe 的新的PP技术，它可以实现更好的内存效率和更新鲜的更新梯度。
WPipe 使用一种新颖的流水线方案，它将模型分为两组，并通过巧妙的移动操作，使得两个分组既可以无缝更新，也可以消除一半的更新延迟和额外内存占用。通过大型语言模型Bert和视觉模型ResNeXt的训练实验表明，与当前SOTA技术PipeDream-2BW 相比，WPipe 实现了 1.4x 的加速且内存占用减少了 36%，同时保障最终的模型准确度不受影响。
介 绍
模型容量越大，表达能力就越强，拥有更好表现效果的潜力就越大，因此为了实现更好的效果，模型的参数规模越来越大。最近有越来越多的工作从这个思路出发，他们使用拥有大规模参数的模型来提升模型的效果，最终也收获了SOTA的结果。相关工作有OpenAI的GPT-3，Microsoft的T5，达摩院的M6等。
然而，大规模模型带来好的效果的同时也给训练提出了新的挑战，这些模型的显存占用要远大于目前最高规格硬件加速设备的显存。以GPT-3为例，它的最大参数量为1750亿，有效训练所需要的显存占用大概是6000G左右，而目前NVIDIA- A100的最大显存只有80G，远满足不了需求。
为了解决这个问题，业内人士开始聚焦到模型并行技术上。传统的模型并行可以分为两类：层内模型并行(intra-layer model parallelism)和层间模型并行(inter-layer model parallelism)。层内模型并行主要的代表工作有Megatron-LM，deepspeed的Zero1，Zero2，Zero3，Zero-Infinity等。层内模型并行每一层都需要两次全员参入的All-gather通信，对通信带宽的要求很高，尤其在训练Transformer这类稠密模型时。
一般在单机8卡的机型下，由于卡间存在NvLink，这个问题并不明显。但是在扩展到多机训练时，这种全员参入的All-gather通信消耗将成为吞吐进一步提升的瓶颈，尤其在机间通信时 (即使有层次化通信优化技术，但是也无法根本解决这个问题)。而层间模型并行，在做完partition后，同一个partition内的各层并不需要通信，其通信的频次可以远低于层内并行。
然而，传统的层间并行由于资源利用率低，不会直接被使用，而是需要通过流水线并行 (Pipeline Parallelism, PP) 来加速，提升资源利用率。PP可以细分为同步流水线并行(Sync-PP)和异步流水线并行 (Async-PP)。Sync-PP的代表工作有GPipe，PipeDream-flush等；Async-PP的代表工作有PipeDream，PipeDream-2BW等。同步方法与数据并行具有相同的权值更新语意，但是需要引入bubble（空闲等待时间），会降低训练吞吐。
异步方法彻底消除的训练timeline中的bubble，但是需要引入不同的权值版本来解决权值过期的问题。WPipe可以使用更少的权值版本来解决权值过期问题，理论上，运行时WPipe相比PP领域的SOTA方法PipeDream-2BW可以节省一半运行时内存占用，并且拥有更少的权值切换成本，相同batch下的吞吐也能提升1.4x。
论文详细信息如下：
论文题目：Group-based Interleaved Pipeline Parallelism for Large-scale DNN Training
论文地址和论文的代码，均可在openreview.net的官网上获取。
相关技术及问题
模型并行是一种将一个大模型按照一定的规则划分为多个partitions，并将其分配到不同的加速设备上，每一个加速设备负责其拥有权值的前向推理和反向求导，并将产生的中间激活/梯度发送到其他加速设备上，同时也接受其他加速设备产生的激活和梯度。按照划分方向的不同，分为层内模型并行和层间模型并行，如图1所示。本小节主要介绍层间并行相关技术。
图- 1模型并行的划分
流水线并行 (Pipeline Parallelism, PP)是一种层间并行技术，通过让不同的workers在同一时刻执行不同的mini-batches的数据来实现计算利用率的提升，如图2所示，从(a)到(b)可以有效地提升资源利用率。然而原始的PP依然不能直接使用，如图2(b)所示，worker0上，mini-batch 2的反向过程所使用的权值，被mini-batch 1的反向过程更新，导致mini-batch 2的前反向过程权值版本不统一。
图- 2传统层内并行和流水线并行
当前主流的方法通过权衡吞吐和显存占用来解决这个权值过期问题。Sync-PP方法，代表有GPipe和PipeDream-flush，前者将一个mini-batch划分为多个micro-batch，并在micro-batch内做流水线并行，这个方法彻底解决的过期问题问题，但是需要引入空闲区域 (bubbles)，牺牲了吞吐。
后者PipeDream-flush本身并没有对GPipe做本质的修改，主要改进是将micro-batch的反向过程前移了，从而可以使激活释放的更加迅速，并且到稳定状态后 只有 个micro-batch的反向过程激活没有被释放，而不是M(micro-batch的数目，远大于worker的数目)。所以PipeDream-flush的内存占用更小。
图- 3 (a) GPipe，(b) PipeDream-flush，(c) PipeDream 和(d) PipeDream-2BW的执行时序图
Async-PP方法，代表有PipeDream和PipeDream-2BW。PipeDream默认使用Weight stashing技术， 需要维持个 权值版本来保证同一个worker的前向和反向所使用的权值版本一致，虽然这个方案可有效的消除同步方法中的空闲区域，但是需要大量的显存空间来保存多余的权值版本。
针对这个问题，PipeDream-2BW提出了进一步的改进方案--2BW。2BW利用梯度累积技术有效减少了PipeDream里的权值版本，权值更新语意中的梯度延迟项也减少到常数1。尽管如此PipeDream-2BW依然有可以改进的地方。
WPipe的解决方案
● 问题引入
如图4所示是插入空闲时间块后的PipeDream-2BW的执行示意图。相邻两个周期被idle时间块分隔，相互之间没有交叉，类似PipeDream-flush，可以彻底解决权值过期的问题，但是这些idle时间块影响吞吐的进一步提升。通过统计，我们发现这些idle的时间块，刚好是一个模型并行的执行流水线图，所以能不能用另外执行流水线来填充这些idle时间块，成为我们的优化思路。
图- 4 插入Bubble的PipeDream-2BW
● 2GW
在上文的猜想的指导下，我们做了下面的尝试，首先(1)我们将划分好的模型partitions分为两组 和 ， 如图5(a)所示；（2）然后我们将下一个更新周期的 前向过程移动到当前更新周期的 的反向过程的前面，如图5(b)所示；（3）最后我们交叉每一个训练周期就得到了无缝更新pipeline，如图5(c)所示。
图- 5 WPipe的推理过程
● 内存分析
● 实验分析
我们从收敛性，吞吐，和内存占用三个方面来分析WPipe的性能。我们分别选取了CV领域的ResNeXt系列模型和NLP领域的Bert系列模型。
● 收敛性
从表2中的各项试验数据可以看出，WPipe，DP，和PD-2BW三者的最终准确率相差不大，说明WPipe和PD-2BW一样都对收敛性无损。但是分析他们的训练过程中我们发现，WPipe的loss曲线和准确率曲线更加接近数据并行（DP），如图6所示，更多的loss曲线和准确率曲线可以看论文附录。
表格 2 收敛性分析表
图- 6 训练loss变化曲线和一个从头开始训练的Accuracy变化曲线
● 吞吐
由于训练的batch-sze很大程度影响深度模型的最终的训练效果，batchsize太大或者太小都影响收敛性。所以对于一个实际使用的训练系统，需要一定的batch范围内有更高的吞吐。常用的batchsize的范围是 ，对于大模型而言，很多时候是达不到这个batch范围的，所以训练大模型首先需要节省显存。
下图所示在 范围内，WPipe均有更高的吞吐。同时我们发现（1）在大规模模型训练重计算可以有效提高batchsize，且这种提升效果在WPipe-R中尤其的明显，如图7(e)WPipe的batchsize提升了44倍，而PipeDream-BW只有16倍；(2) WPipe结合数据并行和模型的变换空间更大，它的解空间比PipeDream-2BW更广，如图7(f)。
图- 7 吞吐分析的各组实验
● 显存占用
减少显存占用方面，WPipe效果明显，尤其是WPipe-R。并且随着batch-size的增大，WPipe-R的效果越明显。这个得益于WPipe-R的运行时激活可以减少一半的优点。
图- 8 在ResNeXt200上随batchsize增加的条件下，显存占用上，与主流方法的比较分析
图- 9 在Bert96上随batchsize增加的条件下，显存占用上，与主流方法的比较分析
● 结论
在这项工作中，我们提出并实现了 WPipe，这是一种基于分组的交错流水线并行训练系统。 与最先进的方法 PipeDream-2BW 相比，WPipe 通过双分组权重更新实现了更好的内存效率、更高的吞吐量和更新鲜的权重更新。 具体来说，（1）吞吐量：WPipe实现1.4倍加速； （2）内存占用：WPipe-R减少了36%的内存占用； (3) 收敛性：虽然 WPipe 和 PipeDream-2BW 在训练时具有相似的最终准确率，但 WPipe 具有更接近数据并行性的权重更新语义。
参考文献
[1] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020.[2] Chi-Chung Chen, Chia-Lin Yang, and Hsiang-Yun Cheng. Efficient and robust parallel dnn training through model parallelism on multi-gpu platform. arXiv preprint arXiv:1809.02839, 2018.[3] Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory-efficient pipeline-parallel dnn training. In International Conference on Machine Learning, pp. 7937–7947. PMLR, 2021.[4] Jeffrey Dean, Greg S Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Quoc V Le, Mark Z Mao, Marc’Aurelio Ranzato, Andrew Senior, Paul Tucker, et al. Large scale distributed deep networks. 2012.[5] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.[6] Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Ganger, and Phil Gibbons. Pipedream: Fast and efficient pipeline parallel dnn training. arXiv preprintarXiv:1806.03377, 2018.[7] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Advances in neural information processing systems, pp. 103–112, 2019.[8] Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Kurt Keutzer, IonStoica, and Joseph E Gonzalez. Checkmate: Breaking the memory wall with optimal tensor rematerialization. arXiv preprint arXiv:1910.02653, 2019.[9] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. arXiv preprintarXiv:1909.11942, 2019.
[10] Junyang Lin, Rui Men, An Yang, Chang Zhou, Ming Ding, Yichang Zhang, Peng Wang, Ang Wang, Le Jiang, Xianyan Jia, et al. M6: A chinese multimodal pretrainer. arXiv preprint arXiv:2103.00823, 2021.
