G ROUP - BASED I NTERLEAVED P IPELINE P ARALLELISM FOR L ARGE - SCALE DNN T RAINING
A BSTRACT
The recent trend of using large-scale deep neural networks (DNN) to boost performance has propelled the development of the parallel pipelining technique for efﬁcient DNN training, which has resulted in the development of several prominent pipelines such as GPipe, PipeDream, and PipeDream- 2BW .
However, the current leading pipeline PipeDream- 2BW still suffers from two major drawbacks, i.e., the excessive memory redundancy and the delayed weight updates across all stages.
In this work, we propose a novel pipeline named WPipe, which achieves better memory efﬁciency and fresher weight updates.
WPipe uses a novel pipelining scheme that divides model partitions into two groups.
It moves the forward pass of the next period of weight updates to the front of the backward pass of the current period of weight updates in the ﬁrst group, retains the order in the second group, and updates each group alternatively.
This scheme can eliminate half of the delayed gradients and memory redundancy compared to PipeDream- 2BW .
The experiments, which train large BERT language models, show that compared to PipeDream- 2BW ,
WPipe achieves 1 .
4 × acceleration and reduces the memory footprint by 36%, without nearly sacriﬁcing any ﬁnal model accuracy.
I NTRODUCTION
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 have focused on model parallelism, which allows to further scale up model parameters signiﬁcantly by partitioning a large model over multiple accelerators/workers.
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 pipelining technique can effectively improve the utilization of computing resources by scheduling different workers to consume different mini-batches in parallel.
However, naive pipelining cannot be directly used for DNN training due to two 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-ﬂight activations, which are heaped up by continuous forward activations waiting for their corresponding backward passes in pipelining.
To address the staleness problem, some workable pipelines have been proposed with different emphases on low memory footprint and high throughput.
GPipe proposes the periodic pipeline ﬂushing 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 proposes the weights stashing (Harlap et al., 2018) technique to eliminate bubble overhead, resulting in higher throughput; however, the multiple weight versions incur a high memory footprint.
To further reduce the weight versions maintained in PipeDream, PipeDream- 2BW was proposed in which the doublebuffered 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.
To reduce in-ﬂight activations, GPipe and PipeDream- 2BW adopt activation recomputation (Chen et al., 2016; Jain et al., 2019).
Recomputation enables the forward pass to execute without retaining the internal activations, which are recomputed in the corresponding backward pass.
This method of being recomputed when needed can effectively alleviate the stacking up of in-ﬂight activations and greatly reduce the memory footprint.
Although both GPipe and PipeDream- 2BW utilize the recomputation technique, PipeDream- 2BW has fewer in-ﬂight activations than GPipe owing to the more timely execution of the backward pass in its pipeline.
Regarding Pipedream, unfortunately, it doesn’t take effective measures on this issue, which prevents it from training a larger model.
Currently, PipeDream- 2BW is the pipeline-parallel system with the best overall performance.
However, it still suffers from two main problems: delayed weight updates and excessive memory redundancy on weight versions and activations (Narayanan et al., 2021).
In this paper, we propose a novel pipeline, named WPipe, to solve these two problems.
In WPipe, model partitions are divided into two groups, namely G 0 and G 1 , which share the same workers.
In the execution pipeline, the forward pass of the next update period is moved to the front of the backward pass of the current update period in
G 0 , and the execution order in G 1 is retained, and then, G 0 and G 1 are updated alternatively.
We named the alternate update technique for the grouping weights as double-grouped weight updates ( 2GW ).
Compared to 2BW , 2GW has two main improvements: (1) 2GW realizes the same weight update semantics as that in data parallelism and only maintains one weight version without introducing bubble overhead in G 1 .
(2) 2GW reduces half of the in-ﬂight activations since only half of the weights are involved in the training at a time.
Additionally, we incorporate some conventional but effective communication optimization techniques including the overlap of computation and communication, hybrid parallelism, and heterogeneous communication to further reduce communication overhead.
Finally, 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.
billion parameters, which cannot be trained by GPipe and PipeDream, WPipe is 1 . 4 × faster than
PipeDream- 2BW .
Memory experiments show that compared to PipeDream- 2BW , WPipe reduces the memory footprint by 35.8% for Bert96.
Convergence experiments show that WPipe can achieve similar ﬁnal accuracy to data parallelism.
B ACKGROUND
In this section, we brieﬂy 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 a large model into multiple parts and assigns them to different workers.
Each worker is responsible for their own weight updates and sending and receiving intermediate activations and gradients.
The conventional model parallelism includes inter-layer model parallelism and intra-layer model parallelism (Shoeybi et al., 2019).
The former suffers from underutilized resources as only one worker is active at a time, as shown in Figure 1a.
The latter requires all-to-all aggregation with intermediate activations/gradients for each layer.
The all-to-all communication overhead, which grows in proportion to the number of layers, makes it difﬁcult to expand to deeper layers, especially in heterogeneous network interconnections.
[Figure]
execution timelines of inter-layer model parallelism and naive pipeline parallelism.
Pipeline Parallelism.
Pipeline parallelism can effectively improve the utilization of computing resources of inter-layer model parallelism by scheduling different workers to consume different mini-batches simultaneously.
However, naive pipeline parallelism suffers from a staleness issue for weight updates.
As shown in Figure 1b, 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 ﬂushing technique (Huang et al., 2019).
As shown in Figure 2a, it splits a minibatch into multiple smaller micro-batches, which are injected into the model continuously to achieve pipelining.
Within each pipeline period, GPipe can accumulate gradients across micro-batches and ﬂush synchronously.
In this way, it achieves the same weight update semantics as data parallelism and only maintains one weight version.
However, in GPipe, bubble overhead (Huang et al., 2019) limits throughput, especially when a large mini-batch cannot be supported.
PipeDream-ﬂush is essentially the same as GPipe (Narayanan et al., 2021), and its main improvement is to move forward the backward pass, thereby releasing in-ﬂight activations as soon as possible and reducing memory footprint.
[Figure]
[Figure]
Update 𝑊 !!
𝑊 !" 𝑊 "!
𝑊 !!
𝑊 !" 𝑊 "!
𝑊 ""
Weight stashing
Forward Pass
Backward Pass
Idle
[Figure]
Time
Time 2BW
[Figure]
Update
[Figure]
[Figure]
Update
[Figure]
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 efﬁciency.
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 at each stage (Narayanan et al., 2021).
S YSTEM D ESIGN OF WP IPE
In this section, we ﬁrst 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. I NTRODUCTION
To simplify the description, we denote the naive model-parallel timeline as V 2 (where “2” represents the number of model partitions), as shown in Figure 1a.
Then we denote the process of cross-splicing
V 2 s to form a pipeline as P ( nV 2 ) (where n represents the number of V 2 ), as shown in Figure 1b.
Analyzing the execution regular of the pipeline in Figure 1b, we can ﬁnd that P ( nV 2 ) can ensure that
[Figure]
Idle
Introduced
[Figure]
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 P ( nV 2 ) cannot be used directly due to the inconsistency issue discussed above.
To address the issue, 2BW uses two weight versions and combines gradient accumulation.
Why does 2BW work effectively?
The point is that 2BW effectively handles the intersection of adjacent update 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-batchesand 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 pipeline is the same as PipeDream-ﬂush).
If these idle time blocks cannot be eliminated, can we ﬁll 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 ( xV n ) , x > = ( n − 1) are required to eliminate the intersections.
In addition, a pipeline accumulation period has at least execution time blocks shaped like P ( xV n ) , x > = n .
Obviously, when x > = n , they have the same shape.
Thus, can we ﬁll these idle time blocks with the pipelining of another part of the weights of the same model?
D OUBLE -
G ROUPED W EIGHT U PDATES ( 2GW )
In this subsection, we introduce the 2GW technique from two points: model partitions grouping and weight update semantics.
(a) (b) (c)
G 0 ,F
G 1 , F&B
G 0 ,B
[Figure]
B
[Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure] [Figure]
[Figure]
Model Partitions Grouping.
As shown in Figure 4a, based on Figure 1a, we further split each model partition into two and obtain twice the number of model partitions.
Then, we divide them into two groups, G 0 and G 1 , and train them on the same devices.
For G 1 , it is a typical V 2 structure, which can realize continuous pipeline expansion through P .
But for G 0 , since its forward pass and backward pass are separated, a continuous pipeline cannot be directly achieved by cross-splicing ( P ).
Therefore, we ﬁrst need to convert it into a V 2 structure.
By analyzing the execution rules of the pipeline after grouping, we found that moving the forward pass of G 0 of the next step to the front of the backward pass of G 0 of the current step can form a V 2 , as shown in Figure 4b.
After doing this, we can cross-splice G 0 and G 1 respectively to obtain P ( G 0 ) and P ( G 1 ) , and then continue to cross-splice their results to obtain P ( P ( G 0 ) , P ( G 1 )) .
For G 1 , its micro-batches 2 and 3 can be well
[Figure]
separated by the pipelining of G 0 ( P ( G 0 ) ), as shown in Figure 4c.
This also answers the question mentioned in section 3.1.
The pipelining of G 0 can ﬁll in the idle time blocks in Figure 3.
[Figure]
Weight Update Semantics.
As shown in Figure 4c, G 1 only maintains one weight version and has the same weight update semantics as data parallelism.
For G 0 , due to the movement operation, the intersections naturally exist, such as micro-batches 3, 4, and 2 in Figure 4c.
Nevertheless, G 0 only needs to maintain two weight versions and has only one delay term in the weight update semantics, which is the same as 2BW .
If using SGD optimizer, the weight update semantics of 2GW are as follows:
[Equation]
Other optimizers can be similarly analyzed.
Figure 5 shows a more detailed pipeline-parallel timeline of 2GW .
In each pipeline stage, the model partitions of G 0 and G 1 are executed alternately without introducing bubbles after the system is warmed up.
3. M EMORY F OOTPRINT A NALYSIS
The memory footprint of most pipelines is mainly divided into four parts, i.e., naive model weights, weight buffers (saving history weight versions), in-ﬂight activations, and optimizers.
Since the naive model weights and optimizers are roughly the same for each pipeline method, we mainly analyze weight buffers and in-ﬂight activations.
In WPipe, G 1 has no weight buffer, and the size of G 0 ’s weight buffer is 2.
We use S p to represent the size of the entire model parameters, M G 0 and
M G 1 represent the sizes of G 0 and G 1 , respectively.
Supposing that the model is evenly divided ( M G 1 = M G 0 ), the size of weight buffer of WPipe is M = 2 M G 1 =
M G 0 + M G 1 = S p .
For in-ﬂight activations, we use S a to indicate the activations of the entire model for a micro-batch.
From
Figure 5, we statistically conclude that the in-ﬂight activation amount of stage i is equal to M ∗ S a 2 N for G 0 and ( M − i ) ∗ S a 2 N for G 1 , and of all stages is equal to:
[Equation]
where the N and M indicate the number of pipeline stages and micro-batches, respectively, and
M > = N .
The same statistical method is used for GPipe, PipeDream, and PipeDream- 2BW to obtain the size of all weight buffers and in-ﬂight activations, respectively, as shown in Table 1.
With Activation Recomputation.
From Table 1, the in-ﬂight activations increase linearly with the number of pipeline stages, which will become a memory bottleneck when training a large model.
Thus, a pipeline system must take measures to reduce in-ﬂight activations.
The activation recomputation is an effective method widely accepted, especially for 2GW .
It can eliminate the stacking up of historical in-ﬂight activations, leaving only the activations in an active state in the system, as shown in Table 1.
From the table, we can summarize the following points: (a) compared to PipeDream- 2BW , the weight buffer size of WPipe is reduced by half, and the superiority is more obvious compared to PipeDream.
Although GPipe has no weight buffer, it has to introduce
N − 1 M + N − 1 bubble overhead; (b) recomputation can signiﬁcantly reduce in-ﬂight activations for GPipe, PipeDream- 2BW , and especially WPipe.
3. C OMMUNICATION O PTIMIZATION
We use C a to represent the cost of intermediate activations/gradients in one communication, and then the model-parallel communication overheads of different pipeline systems are C WPipe =
N s −
) C a , C 2BW
=
( N s − 1) C a , and C GPipe =
( N s − 1) C a , respectively, where N s indicates the number of stages.
They show that with the same number of stages, the communication overhead of
WPipe is twice that of GPipe and PipeDream- 2BW .
Nevertheless, in most cases, C a is small and can be ignored, especially when training some language models, such as Transformer (Vaswani et al., 2017).
When C a is large, We have the following countermeasures:
[Figure]
Combination with Data Parallelism.
Normally, too large C a is caused by too large a micro-batch.
we can proportionally reduce the depth of the pipeline while reducing the size of the micro-batch, and then we proportionally increase the width of data parallelism to maintain the same global batch size, as shown in Figure 6.
As a result, the size of the micro-batch becomes smaller, and the C a also decreases, while the number of accelerators and the global batch size remain unchanged.
Of course, we need to weigh the communication overhead between model parallelism ( C WPipe ) and data parallelism ( C DP ) to choose the appropriate ratio of depth ( d ) and width ( w ).
When C WPipe is greater than C DP , we reduce the value of d : w, d ∗ w = N GPU .
Otherwise, we increase the value of d : w .
Overlap of Computation and Communication & Heterogeneous Network Communication.
The former can almost offset the overhead of activation recomputation.
The latter can effectively use heterogeneous networks to balance communication overhead.
Please refer to the appendix for details.
In addition, some communication compression techniques are also shown in the appendix.
E XPERIMENTS
WPipe is implemented with PyTorch-1.4 (Edward Z. Yang, 2021) and executes on two environments, i.e., a single machine with eight 16-GB V100 GPUs (Env-1) and a private cluster with 8 × 8 V
GPUs (Env-2).
4. Q UALITY OF C ONVERGENCE
In this section, we compare the convergences of WPipe and PipeDream- 2BW by comparing the accuracy when training the same model on the same dataset with the same hyperparameters.
Text Classiﬁcation.
We ﬁnetuned BERT BASE (Devlin et al., 2018) and BERT LARGE (Devlin et al., 2018) for WPipe, PipeDream- 2BW , and data parallelism on the QQP and MNLI tasks (Wang et al., 2018).
We used respectively bert-base-uncase and bert-large-uncase pre-training weights from transformers-3.5.0 (Wolf et al., 2020).
We used Adam optimizer, a learning rate of 8 × 10 − 5 ( ν = 8 × 10 − 5 ) with 1000 steps warmup( ws = 1000 ) and a mini-batch size of 256( b = 256 ) for BERT BASE and the same optimizer, ν = 4 × 10 − 5 with ws = 2000 and b =for BERT LARGE .
From Table 2, WPipe, PipeDream- 2BW , and data parallelism have similar ﬁnal accuracy.
For further analysis, we continue to perform image classiﬁcation experiments.
[Table]
Image Classiﬁcation.
We ﬁnetuned, respectively, the ResNeXt50 (32x4d) (Xie et al., 2017) and
ResNeXt101 (32x8d)
(Xie et al., 2017) for WPipe, PipeDream- 2BW , and data parallelism on the three datasets of CIFAR-10 (Krizhevsky et al., 2009), CIFAR-100 (Krizhevsky et al., 2009), and
Oxford 102 Flowers (Nilsback & Zisserman, 2008).
We used the pre-training weights from the torchvision (Francisco Massa, 2021), SDG optimizer, ν = 1 × 10 − 2 with 0.05 warmup ratio and b = 256 .
From Table 2, there is still no obvious difference in their ﬁnal accuracy.
Thus, we continue to analyze the changes in loss of the three methods during training, while supplementing the training experiments from scratch.
Intermediate Metrics Analysis.
As shown in Figures 7a-7c, although WPipe, PipeDream- 2BW , and data parallelism have similar ﬁnal loss and accuracy, the curve of WPipe is closer to data parallelism.
Among them, the loss of data parallelism drops the fastest, and its accuracy rises the fastest, and
WPipe is second only to data parallelism, which beneﬁts from the fact that the weight update delay of
WPipe is reduced by half, compared to PipeDream- 2BW .
For more loss curves and accuracy curves, please refer to the appendix.
[Figure]
4. T HROUGHPUT
In this section, we measure the throughputs of WPipe through training the large-scale DNN models on Env-1 and Env-2 and compare them with other pipeline systems.
The main factors that affect throughput are batch size ( B ), the ratio of model parallelism to data parallelism ( d : w ), and the possibility to recompute ( S a ) as presented below:
[Equation]
where T is the optimal result in an optimization space S a × S d : w × S M .
H ( d : w ) represents all isomers in the same d : w .
S M represents the gradient accumulation period.
With regards to the optimal solution of T , the study in (Narayanan et al., 2021) gives a workable solution, we do not do further research.
Note that since B directly affects the accuracy of the model, it is often a given value.
4.2. C OMPARISON WITH P IPELINED A PPROACHES
As shown in Figure 8, throughput experiments are carried out through training different layers of
BERT and ResNeXt models on the Env-1 and Env-2.
For more throughput experiment data, please refer to the appendix.
In general, the communication overhead of data parallelism is larger for BERT, and model parallelism is larger for ResNeXt.
Comparison to PipeDream- 2BW .
In general, the suitable batch size is between 2 9 and 2 11 .
If it is too large or too small, it is not conducive to the convergence of the model.
However, in most cases, when training large models, some methods cannot reach this batch size, as shown in Figure 8a.
When in this range, as shown in Figure 8c, when batch = 2 9 , WPipe is 1.4 times faster than PipeDream- 2BW for
Bert768 with 5.5 billion parameters.
Compared to 2BW , WPipe not only reduces the cost of switching between weighted versions by half but also has a larger solution space when the batch is larger.
As shown in Figure 8d, when batch = 2 9 , WPipe can run in the conﬁguration of d : w = 4 : 16 in
Env-2, and 2BW due to memory limitations, the model-parallel depth is at least 8, that is, d > = 8 , so in this case, the throughput of WPipe is higher (for convolutional networks, in general, C a is larger and C DP is smaller, so d : w needs to be smaller to improve throughput).
From Figures 8b and 8d, it can be concluded that after communication optimization, even on a convolutional network with a larger C a , WPipe has a throughput not lower than PiprDream- 2BW .
Comparison to GPipe.
For GPipe, its biggest problem is that there are idle time blocks on its execution timeline.
In addition, its forward activations cannot be released in time.
These factors directly or indirectly affect its throughput.
Compared to GPipe, from Figure 8a- 8d, WPipe is 2 .
1 × faster for Bert192 and 2 .
2 × faster for ResNeXt302.
Comparison to PipeDream.
Due to poor memory efﬁciency, the larger models cannot be trained using PipeDream, such as Bert384 and Bert768.
Practice shows that for ResNeXt200, ResNeXt302,
ResNeXt500, and ResNeXt800, PipeDream cannot automatically split them.
According to existing data, WPipe is 7 . 1 × faster for Bert192.
[Figure] [Figure]
( s m
[Figure]
4.2. C OMMUNICATION O PTIMIZATION A NALYSIS
Recomputation.
The activation recomputation can greatly reduce memory footprint and expand mini-batch size for pipeline systems.
As shown in Figure 8e, PipeDream- 2BW can expand the maximum mini-batch by 16 times, and WPipe can even expand by 44 times.
WPipe has a more obvious memory advantage with activation recomputation.
Although the batch size is not as large as possible, usually when training large models, the batch size must be within a certain range to have better convergence.
In many cases, due to memory limitations, the effective range of the batch size cannot be reached.
At this time the advantage of WPipe is crucial.
[Table]
The Ratio between Model and Data parallelism.
As shown in Figure 8f, as M : N decreases, the corresponding throughput increases for ResNeXt500.
This happens because the communication overhead of data parallelism is smaller than that of model parallelism in the convolutional network.
In this case,
M : N can be adjusted more widely, which is more conducive to improving throughput, e.g. batch = 2 11 in Figure 8f.
Heterogeneous Network Settings.
We change the communication placement of model parallelism and data parallelism in the same M : N to achieve higher throughput.
A means that data-parallel communication takes place between machines and model-parallel communication takes place inside machines, while B is the opposite.
As shown in Table 3, in the case of A , the throughput of WPipe is higher, because the model-parallel communication is the bottleneck for ResNeXt.
Compared with the inter-machine, the bandwidth inside the machine is higher and the communication speed is faster.
The model-parallel communication bottleneck has been better alleviated
. 4. M EMORY O PTIMIZATION A NALYSIS
[Figure]
M e m o r y f oo t p r i n t (
G B )
DP
GPipe
GPipe-R PipeDream
PipeDream-2BW
PipeDream-2BW-R WPipe
WPipe-R
OO M
OO M OO M OO M OO M OO M
OO M OO M
OO M OO M OO M OO M OO M OO M
In the throughput experiment, WPipe can execute on a larger batch, which also reﬂects the advantage of the lower memory footprint of WPipe.
Here, we conduct a more detailed analysis: (1)
The results of experiments show that the memory footprint of various pipeline-parallel methods is signiﬁcantly reduced with activations recomputation, especially for WPipe.
As shown in
Figure 9, WPipe reduces the memory footprint by 65% for Bert96 on the per-GPU micro-batch of 64, and 56% for RestNeXt200 on the perGPU micro-batch of 32.
They are 47% and 29% respectively for PipeDream- 2BW ; (2) with the increase of batch size, the effect of recomputation on WPipe-R is more signiﬁcant.
From Figure 9, for Bert96, WPipe-R is reduced by -14% compared to GPipe-R on the per-GPU micro-batch of 8 (because GPipe has no weight buffer, and its initial memory footprint is lower), and on the per-GPU micro-batch of 64, it increases to 28%.
Compared to PipeDream- 2BW -R, under the same conditions, WPipe-R increases from 21% to 36%.
C ONCLUSIONS
In this work, we proposed and implemented WPipe, a system for group-based interleaved pipelineparallel training.
Compared to the state-of-the-art approach, PipeDream- 2BW , WPipe achieves better memory efﬁciency, higher throughput, and fresher weight updates through the double-grouped weight updates.
Speciﬁcally, (1) throughput:
WPipe achieves 1 .
4 × acceleration; (2) memory footprint: WPipe-R reduces the memory footprint by 36%; and (3) convergence: although WPipe and
PipeDream- 2BW have similar ﬁnal accuracy when training, WPipe has a weight update semantics closer to data parallelism.
A CKNOWLEDGMENTS
This work is supported by the Computational Intelligence Department of Ant Group.
We thank
Jiyan Jiang for his helpful discussions.
We would like to thank the Aliyun EFLOPS team for their substantial support in designing the industry-leading training platform to facilitate fast trials in this work.
We also thank anonymous reviewers for their insightful and valuable comments
A A PPENDIX A. O THER E XPERIMENTAL D ATA OF C ONVERGENCE .
Figures 10a-10h show the training loss curves of the remaining experiment in Table 2.
Figure 10i shows the F1 curve of BERT BASE when training from scratch.
Analyzing the trend of the curve, we found that WPipe and data parallelism can converge normally
[Figure]
but PipeDream- 2BW does not.
[Figure]
A. T HROUGHPUT
In the experiment, our settings of models are hidden size=768 , num attention heads=12 , seq len=128 of all Bert models, and groups=32 , width per group=4 of all ResNeXt models.
Regarding cluster conﬁguration, there are 8 machines in our private cluster, and each machine has 8 GPUs with a memory size of 16G, Intel(R)Xeon(R)
Platinum 8163 CPU, 512GB of RAM with a 25Gbps Ethernet interface, and 300GBps NVLink (nvl), which is 96 times the Ethernet bandwidth.
In addition, the version of PyTorch we used was 1.4 .
As shown in Figure 11a-11d, we continue to train Bert96 and ResNeXt200 on a single machine with
WPipe, PipeDream- 2BW , PipeDream, and GPipe, and train Bert384 and ResNeXt500 on multiple machines with WPipe and PipeDream- 2BW .
The overall conclusion is the same as the above.
WPipe has a more obvious acceleration effect on Transformer series networks (Bert, GDP, etc.), but has a slight acceleration on convolutional networks, compared to PipeDream- 2BW .
[Figure]
r o u g h p u t
[Figure]
m
[Figure]
l e s / s e c o n
[Figure] [Figure]
T
[Figure]
r o u g h p u t ( s a m p l e s / s e c
o n d )
[Figure]
show the effect of compressed model-parallel communication on different models.
A. C OMMUNICATION O PTIMIZATION
In this section, we use a speciﬁc example to analyze the effectiveness of the model-parallel communication compression.
In addition, we add a detailed description of overlap and heterogeneous network communication.
[Figure]
Communication Compression.
Referring to the automatic mixed precision algorithm of the apex (Christian Sarofeen, 2021), we implemented the automatic precision compression technique for intermediate activations/gradients, which reduced the communication overhead by nearly half across the pipeline.
This offsets the communication overhead introduced by further splitting model partitions.
As shown in Figure 12, the sender divides the intermediate activations/gradients by appropriate scaler s (powers of 2) and converts them into half-precision tensors.
If the conversion succeeds, the sender will transmit the half-precision tensors together with the scaler s to the receiver.
After receiving the data, the receiver restores their original precision by the scaler .
If the conversion fails, the sender will transmit the full precision tensors.
[Table]
Automatic Precision Compression (APC).
As shown in
Table 4, WPipe is 1 .
17 × faster with APC when training
ResNeXt500 and the batch size equals 3328.
The loss (the negative value) of accuracy, which is brought by APC, does not exceed 0.1%, as shown in Table 5.
However, when the communication overhead is not large, the use of APC will not speed up but will be slower.
As shown in Figure 11f, when batch < = 2 11 , C a is small, and the acceleration beneﬁt cannot exceed the compression overhead.
When batch > 2 11 , there is positive feedback.
For convolutional networks, C a is large, and APC has always had a positive effect, but the effect is only obvious when the batch size is large enough, as shown in Figure 11e.
In summary, when
[Table]
the batch is large, we use the communication compression technique to have a positive beneﬁt, but in many cases, we do not need to use such a large batch size to train the model, so in most cases, model-parallel communication will not be a bottleneck.
Overlap of Computation and Communication.
As shown in Figure 13, the activations or gradients communication can overlap with the forward pass, backward pass, and activation recomputation.
Especially for activation recomputation, its time overhead can be offset mostly.
[Figure]
Heterogeneous Network Communication.
Generally, the intra-machine bandwidth (NVLink) is higher than the inter-machine bandwidth.
In WPipe, the communication tasks mainly include model-parallel communication and data-parallel communication.
Thus, WPipe allows tasks with large communication overhead to use the higher intra-machine bandwidth, and tasks with small communication overhead to use inter-machine bandwidth, to balance communication overhead.
[Figure]
P
P
P
P
Stage Stage
A. E XPANSION
GPipe Grouping.
As shown in Figure 15, the pipeline grouping technique can also be applied to
GPipe, thereby reducing GPipe’s bubbles.
In Figure 15, the two model partitions are further divided into four and divided into two groups.
At this time, the bubble is reduced by half.
For further discussion, when the number of groups is N , the bubble will be reduced to 1 N , but the model-parallel communication overhead will increase by N times.
However, if NVLink can be used, the impact of increased communication overhead will be greatly reduced.
F
F F F
B B B B
F
F F F F
F F F
B B
B
B
B B
B
B
Bubble Bubble