Press "Enter" to skip to content

AI顶会ICLR 2022 | WPipe 蚂蚁集团大规模 DNN 训练的流水线并行技术

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

导言

 

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的推理过程

 

 

上述操作之后,
的执行pipeline将被
的执行pipeline隔离开,使得
相邻更新周期没有交叉,从而对
而言就没有更新延迟。对于
由于一个更新周期的内部都存在交叉,所以依然存在更新延迟。尽管如此,整体的更新延迟也被减少一半,用于解决梯度过期问题的权值版本也会减少一半。同时由于同一时刻只有一半的权值参与运算,所以运行时的中间激活也会减少一半。

的权值更新公式如下。

 

●  内存分析

 

如表1所示,是WPipe相比GPipe,PipeDream,PipeDream-2BW的内存占用分析,其中
单个macro-batch的运行时激活占用,
表示总的参数量,
表示总的pipeline stage的数目,
表示梯度累积的周期。从表中我们可以清楚的看到,WPipe相比PipeDream-2BW减少了一半的参数缓存。重计算时也减少了一半的运行时激活,这点很重要,在大规模模型训练时往往运行时激活的显存占用要远高于模型参数,在大batchsize时尤其明显。

 

表格 1主流的Pipeline方法的内存消耗分析

 

 

●  实验分析

 

我们从收敛性,吞吐,和内存占用三个方面来分析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.

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注