GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

paper reading

Posted by Dawei on March 8, 2019


这篇论文利用流水线并行化大规模神经网络训练过程,不仅加速模型训练过程,而且能够训练更大的模型。主要做法包括:1.将模型划分为多个模块,分散在不同硬件上训练。2.将一个mini batch数据分解为T个micro batch,方便流水线化处理。3.每个硬件(accelerator,下面简称acc)上在反向传播时候,重新计算前向传播值,节省存储空间。通过这3项措施,论文显示在使用8个acc时能训练25倍大小的单acc模型,同时保证梯度计算一致性。在使用4个acc时,能够达到3.5倍的加速比。


现有的Tensorflow、Pytorch等框架并行化操作无非两种,第一个数据并行化,将一个batch分解成几个部分,送进具有相同模型的acc里,但是这样每个acc都需要存储模型所有的参数,现有实验已经证明越大的模型(参数量越多)表现越好,这种模式满足不了这种扩展性;第二个模型并行化,将模型分解为可以并行的几部分放到不同的acc里,这样它们可以同时处理,但是神经网络大多或者大部分是序列化模型,可以并行化处理的部分不多。我的理解是论文参考cpu执行指令的流水线式操作,提高处理速度。


现在介绍具体方法。首先将要处理的模型定义为L layers序列模型,每一层有pure function(operator)f,参数w,估计计算量c。首先将模型化为K个部分,每个部分计算量保持相近(启发式算法划分),每个部分放置在一个acc上,相邻的acc能够由通信组件进行数据传输。然后将大小为N的mini batch划分成T份,每一块大小为N/T个数据,我们定义$F_{k,t}$为前向计算里第t个micro batch在第k个acc的计算,$B_{k,t}$表示反向传播。于是第$k+1$ acc在完成第$t-1$份数据的计算后,立即接收来自$F_{k,t}$的结果,与此同时第$k$ acc开始处理$F_{k,t+1}$,在处理一个mini-batch数据时,每个acc都会进行T次计算。最后一个acc负责将输出整合,给出最终loss。


在反向传播时,每个micro batch的梯度都是基于前向传播的参数,在mini batch的最后进行整合更新,所以保证了梯度下降更新的同步一致性,这和parameter server不一样,论文说这是因为神经网络对学习率更新策略和dropout概率很敏感。另外的优化是在反向传播时不会保存所有前向计算的中间结果,我们知道保存中间结果能够直接在反向计算中使用,从而节省大量计算量,但是从论文中结果显示出,相比于计算能力,存储空间才是更大的瓶颈所在,论文显示只保存每个acc的输出,这样在反向传播时候,每个acc内部需要重新计算前向传播值,这样每个acc只需要存储N(mini batch size)个数据在此acc的输入值,在每个micro batch计算时,在保存acc内部layer的中间值即可,用完删掉,再来下一个micro batch就行了。


论文在实验部分探讨resnet和AmoebaNet(我不是做cv的,不太了解)的表现效果,显示在流水线按只有一个组件(就是单GPU训练),吞吐量会有25%的损失,这是重新计算带来的后果,但是在8组件流水线下,AmoebaNet的参数量达到25倍,加速比也很吸引人,具体数据在论文Fig. 3,论文显示在组件越多的情况下,load imbalance带来的损失就越大,所以寻找有效的算法将模型计算量分布在acc上非常重要,论文结果也显示越大的模型会带来越优秀的效果,现有的单硬件水平已经不能满足日益膨胀的数据集和与之对应的大模型了,Gpipe大有用武之地。论文最后说明该成果在其它领域的模型如BIGGAN和BERT同样适用。