Group-based Interleaved Pipeline Parallelism for Large-scale DNN Training

1 INTRODUCTION . Several recent lines of research ( Liu et al. , 2019 ; Yang et al. , 2019 ; Lan et al. , 2019 ; Raffel et al. , 2019 ; Brown et al. , 2020 ; Lin et al. , 2021 ) on various application domains have collectively demonstrated that larger DNN models can yield better performance . This creates an emerging trend in scaling up the number of model parameters , resulting in very large DNNs , the memory footprint of which will be beyond the limit of a single accelerator . To resolve the aforementioned problem , researchers and practitioners focused on the use of the modelparallel technique to partition large models over multiple accelerators/workers ; this allows to further scale up model parameters significantly . However , the conventional model parallelism , which includes inter-layer model parallelism ( Narayanan et al. , 2021 ) and intra-layer model parallelism ( Shoeybi et al. , 2019 ) , either suffers from low resource utilization or high communication overhead . Recently , some studies proposed the pipeline-parallel technique to accelerate conventional model-parallel training . The pipeline technique can effectively improve the utilization of computing resources by scheduling different workers to consume different mini-batches in parallel . However , naive pipelining can lead to two major problems : ( 1 ) staleness of weight updates ( Harlap et al. , 2018 ) , which is caused by the inconsistency of the weight versions used in a forward pass and its corresponding backward pass , and ( 2 ) excessive in-flight activations , which is caused by excessive forward passes waiting for their corresponding backward passes in pipelining . The problem of staleness of weight updates can be solved using workable pipelines with different emphases on low memory footprint and high throughput . GPipe uses the periodic pipeline flushing technique , which maintains only one weight version . This technique can achieve a low memory footprint and almost the same weight update semantics as data parallelism , but at the cost of introducing bubble overhead ( Huang et al. , 2019 ) , thus limiting throughput . PipeDream uses the weights stashing ( Harlap et al. , 2018 ) technique without introducing bubble overhead , resulting in higher throughput ; however , the multiple weight versions incur a high memory footprint . To further reduce the excessive weight versions maintained in PipeDream , PipeDream-2BW was proposed in which the double-buffered weight update technique ( Narayanan et al. , 2021 ) is used , which largely reduces the number of weight versions and enables similar weight update semantics as that in data parallelism . Moreover , GPipe and PipeDream-2BW use the activation recomputation technique ( Chen et al. , 2016 ; Jain et al. , 2019 ) to solve the problem of excessive in-flight activations . Recomputation enables the forward pass to function without retaining its generated activations . Instead , it will be recomputed in the corresponding backward pass . This effectively alleviates the stacking up of in-flight activations and greatly reduces the memory footprint . Notably , although both GPipe and PipeDream-2BW utilize the recomputation technique , PipeDream-2BW has less in-flight activations than GPipe owing to the more timely execution of the backward pass in its pipeline . Additionally , PipeDream does not take any measures for solving the activations problem , and therefore the training of very large models becomes difficult . At present , PipeDream-2BW is the best pipeline-parallel system in terms of overall performance . However , it still suffers from two main problems : ( 1 ) delayed weight updates , which is to solve the problem of weight staleness ( Narayanan et al. , 2021 ) , and ( 2 ) excessive memory redundancy , which includes two weighted versions in 2BW ( Narayanan et al. , 2021 ) and in-flight activations . In this paper , we propose a novel pipeline , named WPipe , to solve these two problems . It divides model partitions into two groups , namely G0 and G1 , which share the same workers . In this method , the forward pass of the next period is moved to the front of the backward pass of the current period in G0 , and the order in G1 is retained , and then , G0 and G1 are updated alternatively . We named the alternate update technique for the grouping weights as double-grouped weight updates ( 2GW ) . Compared to 2BW , the 2GW pipeline also does not introduce idle time , and in G1 , 2GW realizes the same weight update semantics as that in data parallelism and only maintains one weight version ; in G0 , to deal with the weight staleness problem caused by the forward pass move operation , we reserve a delay item in the semantics of weight update and maintains the same weight versions . In addition , the in-flight activations of WPipe are reduced by half , because only half of the weights are involved in the training at a time . Therefore , it reduces half of the delayed gradients and memory redundancy compared to PipeDream-2BW . Additionally , we combined the overlap of computation and communication , hybrid parallelism , and heterogeneous communication to further reduce communication overhead . Our throughput experiments show that WPipe achieves acceleration of 7.1× compared to that of PipeDream and 2.1× compared to that of GPipe for Bert192 with 1.4 billion parameters . When training models with up to 5.5 billion parameters , which can not be trained by GPipe and PipeDream , WPipe is 1.4× faster than PipeDream-2BW . The memory experiments show that compared to PipeDream-2BW , WPipe reduces the memory footprint by 35.8 % for Bert96 . 2 BACKGROUND . In this section , we briefly introduce the related techniques of model parallelism . Model parallelism . Model parallelism ( Chen et al. , 2018 ; Chilimbi et al. , 2014 ; Dean et al. , 2012 ; Jia et al. , 2018 ) partitions large models into multiple parts and assigns them to different workers . Each worker is responsible for their weight updates and sending and receiving intermediate activations and gradients . Inter-layer model parallelism ( Narayanan et al. , 2021 ) suffers from underutilized resources as only one worker is active at a time , as shown in Figure 1a . Intra-layer model parallelism ( Shoeybi et al. , 2019 ) requires all-to-all aggregation with intermediate activations for each layer . The all-to-all communication overhead , which grows in proportion to the number of layers , makes it difficult to expand to deeper layers , especially in heterogeneous network interconnections . Pipeline Parallelism . Pipeline parallelism can effectively improve the utilization of computing resources of inter-layer model parallelism by scheduling different workers to consume different minibatches simultaneously . However , naive pipeline parallelism suffers from staleness issue for weight updates . As shown in Figure 1b , as different workers process different mini-batches simultaneously , when the backward pass of the latter mini-batch is executed , its corresponding weights could be updated by the prior mini-batch . For example , for mini-batch 2 , the weights used in the backward pass were updated by mini-batch 1 . Existing solutions tackle the issue by trading off throughput and memory footprint . GPipe proposes the periodic pipeline flushing technique ( Huang et al. , 2019 ) . As shown in Figure 2a , it splits a mini-batch into multiple micro-batches , which are injected in the model continuously to achieve pipelining . Within each pipeline period , GPipe can accumulate gradients across micro-batches and flush synchronously . In this way , it achieves the same weight update semantics as data parallelism and only maintains a single weight version . However , in GPipe , bubble overhead ( Huang et al. , 2019 ) is introduced , which limits throughput , especially when a large mini-batch can not be supported . PipeDream-flush is essentially the same as GPipe ( Narayanan et al. , 2021 ) . Compared to GPipe , the main improvement is to move forward the backward pass . This can release in-flight activation as soon as possible , reducing memory footprint . PipeDream uses the weight stashing technique to tackle the staleness issue without introducing the bubble overhead , and thus achieves higher throughput ( Harlap et al. , 2018 ) . However , it must maintain a large number of weight versions . From Figure 2b worker i ( i = 0 , 1 , ... ) comprises N − i weight versions , where N is the number of model partitions . This results in a high memory footprint , especially when N is large . Moreover , in PipeDream , the weight update semantics has different delay terms at different stages . PipeDream-2BW ( Narayanan et al. , 2021 ) , as an upgraded version of PipeDream , has higher throughput and more memory efficiency . As shown in Figure 2c , it uses double-buffered weight updates ( 2BW ) , which is combined with gradient accumulation , to reduce effectively the number of weight versions maintained on each worker from N − i to 2 . Notably , in 2BW , its weight update semantics has only one weight delay term across all stages . 3 SYSTEM DESIGN OF WPIPE . In this section , we first analyze the problems of 2BW . Then , we propose double-grouped weight updates ( 2GW ) , and analyze 2GW from three aspects : model partitions grouping , effective learning , and memory footprint . Finally , we introduce several communication optimization techniques . 3.1 INTRODUCTION . We denote naive model-parallel timeline as V2 ( where ” 2 ” represents the number of model partitions ) , as shown in Figure 1a . V2s are cross-spliced to form pipeline , as shown in Figure 1b . We denoted this process as P and its result as P ( nV2 ) ( where n represents the number of V2 ) . As shown in Figure 1b , P can ensure that the backward pass is executed immediately after the forward pass is completed so that the forward activations can be released in time to alleviate the stacking up of the activations . However , the raw pipeline parallelism formed by P and Vn suffers from the inconsistency issue discussed above . To address the issue , 2BW uses double weight versions and combines gradient accumulation . Where is the other weight version in 2BW used ? It is at the intersection of adjacent periods . As shown in Figure 2c , the forward pass for micro-batch 3 is executed before the backward pass for micro-batch 2 , but its backward pass is executed after the backward pass for micro-batch 2 . At this time , it is necessary to continue to use the old weight version in the period of micro-batches 3 and 4 , and the new weight version ( updated by micro-batch 2 ) will be used in the next period ( micro-batches 5 and 6 ) . Thus , if we want to further reduce the number of weight versions , we need to eliminate the intersections of micro-batches 2 and 3 , 4 and 5 , etc . To achieve this , we temporarily introduce idle time blocks , as shown in Figure 3 ( its pipelining is the same as PipeDream-flush ) . If these idle time blocks can not be eliminated , can we fill them ? Through the statistical analysis , we found that for a pipeline of PipeDream-2BW with n model partitions , at least the idle time blocks shaped like P ( xVn ) , x > = ( n− 1 ) are required to eliminate the intersections . In addition , a pipeline accumulation period has at least execution time blocks shaped like P ( xVn ) , x > = n. Obviously , when x > = n , they have the same shape . Thus , can we fill these idle time blocks with the pipelining of another part of the weights of a model ?

Summary: The paper addresses the problem of very large scale deep neural network training through model parallelism. A new model parallel pipeline called WPipe is proposed, that builds on previously existing schemes while targeting (1) lower memory redundancy and (2) fresher weight updates. Experimental evaluation shows that WPipe can achieve significant acceleration (1.4x) and lower memory footprint (36% smaller) when compared to state-of-the-art PipeDream-2BW.