G ROUP - BASED I NTERLEAVED P IPELINE P ARALLELISM FOR L ARGE - SCALE DNN T RAINING
Anonymous authors
Paper under double-blind review
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, namely 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, which 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 will eliminate half of the delayed gradients and memory redundancy compared to PipeDream-2BW.
The results of our experiments, in which large BERT language models were trained, show that compared to PipeDream-2BW, WPipe achieves 1 . 4 × acceleration, reduces the memory footprint by 36% with similar ﬁ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 focused on the use of the modelparallel technique to partition large models over multiple accelerators/workers; this allows to further scale up model parameters signiﬁcantly.
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-ﬂight 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 ﬂ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 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-ﬂight 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-ﬂight activations and greatly reduces the memory footprint.
Notably, although both GPipe and PipeDream-2BW utilize the recomputation technique, PipeDream-2BW has less in-ﬂight 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 difﬁcult.
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-ﬂight activations.
In this paper, we propose a novel pipeline, named WPipe, to solve these two problems.
It divides model partitions into two groups, namely G 0 and G 1 , 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 G 0 , and the 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, the 2GW pipeline also does not introduce idle time, and in G 1 , 2GW realizes the same weight update semantics as that in data parallelism and only maintains one weight version; in
G 0 , 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-ﬂight 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 cannot 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.
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 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 difﬁcult to expand to deeper layers, especially in heterogeneous network interconnections.
[Table]
Worker
[Table] [Table]
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 ﬂushing 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 ﬂush 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 cannot be supported.
PipeDream-ﬂush 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-ﬂight activation as soon as possible, reducing memory footprint.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
2 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 across all stages.
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
We denote naive model-parallel timeline as V 2 (where ”2” represents the number of model partitions), as shown in Figure 1a.
V 2 s are cross-spliced to form pipeline, as shown in Figure 1b.
We denoted this process as P and its result as P ( nV 2 ) ( where n represents the number of V 2 ).
As shown in
[Figure]
Pass
Backward Pass
Idle
Introduced Idle
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 V n 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-ﬂ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 a model?
3. D OUBLE -G ROUPED W EIGHT U PDATES (2GW)
In this subsection, we introduce the 2GW technique from model partitions grouping and weight update semantics.
(a) (b) (c)
G 0 ,F
G 1 , F&B
G 0 ,B
G 1 ,F 1 &B
G 0 ,F 2 &B𝐺 !
Forward Pass 𝐺 !
Backward Pass
Idle 𝐺 "
Forward Pass 𝐺 " Backward Pass
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 with the same number of workers.
Then, we divide them into two groups, G 0 and G 1 .
For G 1 , it is a standard V 2 , which can realize continuous pipeline expansion through P .
But for G 0 , the forward pass and backward pass are separated, which needs to be converted into a V 2 structure before P can be processed.
To achieve the transformation, we move the forward pass of G 0 of the next step to the front of the backward pass of G 0 of the current step to form a V 2 , as shown in Figure 4b.
Then we cross-splice G 0 and Grespectively 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 separated by the pipelining of
[Figure]
[Table]
[Figure]
[Table]
[Figure]
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
[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 4.
Nevertheless, G 0 only needs to maintain two weighted versions, and there is only one delay term in the weight update semantics, which is the same as 2BW.
If using SGD optimizer, 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, where in each pipeline stage, the model partitions of G 0 and G 1 is 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 optimizer.
Since the naive model weights and optimizer are roughly the same for each pipeline method, we mainly analyze weight buffers and in-ﬂight activations.
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 their 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 pay the price of introducing
N − 1 M + N − 1 bubble overhead; (b) recomputation can signiﬁcantly reduce in-ﬂight activations for GPipe, PipeDream-2BW, and WPipe, especially for WPipe.
3. C OMMUNICATION O PTIMIZATION
We use C a to represent the cost of intermediate activations or 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 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 microbatch.
we can proportionally reduce the depth of the pipeline while reducing the size of the microbatch, 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 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 the 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 gradient accumulation period.
With regards to the optimal solution of T , the study in (Narayanan et al., 2021) gives a workable solution, the paper does 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 of 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 PipeDream2BW 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]
[Equation]
[Figure] [Figure]
[Figure]
( s m
[Figure] [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]
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 case 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.
Therefore, the throughput of A is higher.
4. M EMORY O PTIMIZATION A NALYSIS
[Figure]
Bert96(8)
Bert
[Figure]
ResNeXt
[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 cute on a larger the vantage of the lower footprint
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 microbatch 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 state-of-the-art approach, PipeDream-2BW, it 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 A PPENDIX A. O THER E XPERIMENTAL D ATA OF C ONVERGENCE .
Figures 10a-10h are the training loss of the remaining part of the experiment in Table 2.
Figure 10i shows the F1 of BERT BASE training from scratch.
It can be seen that WPipe and data parallelism can converge normally, but PipeDream-2BW does not.
A. T HROUGHPUT
In the experiment, our settings of models are hidden size=768, seq len=128, num attention heads= 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 ma-chine 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]
A. C OMMUNICATION C OMPRESSION
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 about overlap and heterogeneous network communication.
[Figure]
Communication Compression.
With reference to the automatic mixed precision algorithm of apex (Christian Sarofeen, 2021), we implemented the automatic precision compression technique for intermediate activations and 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 and 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.
[Figure]
g h p u t
[Figure]
l e s / s
[Figure]
o n
[Figure] [Figure]
T
[Figure]
r o u g h p u t (
[Figure]
p l e s / s e c o n d )
[Figure]
show the effect of compressed model-parallel communication on different models.
[Figure]
[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 the batch is large, we use 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.
[Table]
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.
Heterogeneous Network Communication.
Generally, the intra-machine bandwidth is higher than the inter-machine bandwidth (NVLink).
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.
In WPipe, the communication tasks mainly include model-parallel communication and data-parallel communication.
[Figure]
Forward
Activation Recv
Gradient
Send
[Figure]
P
P P P P
P P P4 G G
[Figure] P P P
P P P Stage
Stage A. E XPANSION GPipe Grouping. As shown in Figure 15, 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