Pipelined Training with Stale Weights of Deep Convolutional Neural Networks

The growth in the complexity of Convolutional Neural Networks ( CNNs ) is increasing interest in partitioning a network across multiple accelerators during training and pipelining the backpropagation computations over the accelerators . Existing approaches avoid or limit the use of stale weights through techniques such as micro-batching or weight stashing . These techniques either underutilize of accelerators or increase memory footprint . We explore the impact of stale weights on the statistical efficiency and performance in a pipelined backpropagation scheme that maximizes accelerator utilization and keeps memory overhead modest . We use 4 CNNs ( LeNet-5 , AlexNet , VGG and ResNet ) and show that when pipelining is limited to early layers in a network , training with stale weights converges and results in models with comparable inference accuracies to those resulting from non-pipelined training on MNIST and CIFAR-10 datasets ; a drop in accuracy of 0.4 % , 4 % , 0.83 % and 1.45 % for the 4 networks , respectively . However , when pipelining is deeper in the network , inference accuracies drop significantly . We propose combining pipelined and non-pipelined training in a hybrid scheme to address this drop . We demonstrate the implementation and performance of our pipelined backpropagation in PyTorch on 2 GPUs using ResNet , achieving speedups of up to 1.8X over a 1-GPU baseline , with a small drop in inference accuracy . 1 INTRODUCTION . Modern Convolutional Neural Networks ( CNNs ) have grown in size and complexity to demand considerable memory and computational resources , particularly for training . This growth makes it sometimes difficult to train an entire network with a single accelerator ( Huang et al. , 2018 ; Harlap et al. , 2018 ; Chen et al. , 2012 ) . Instead , the network is partitioned among multiple accelerators , typically by partitioning its layers among the available accelerators , as shown in Figure 1 for an example 8-layer network . The 8 layers are divided into 4 computationally-balanced partitions , P0 ... P3 and each partition is mapped to one of the 4 accelerators , A0 ... A3 . Each accelerator is responsible for the computations associated with the layers mapped to it . However , the nature of the backpropagation algorithm used to train CNNs ( Rumelhart et al. , 1986 ) is that the computations of a layer are performed only after the computations of the preceding layer in the forward pass of the algorithm and only after the computations of the succeeding layer in the backward pass . Further , the computations for one batch of input data are only performed after the computations of the preceding batch have updated the parameters ( i.e. , weights ) of the network . These dependences underutilize the accelerators , as shown by the space-time diagram in Figure 2 ; only one accelerator can be active at any given point in time . The underutilization of accelerators can be alleviated by pipelining the computations of the backpropagation algorithm over the accelerators ( Huang et al. , 2018 ; Harlap et al. , 2018 ; Chen et al. , 2012 ) . That is , by overlapping the computations of different input data batches using the multiple accelerators . However , pipelining causes an accelerator to potentially use weights that are yet to be updated by an accelerator further down in the pipeline . The use of such stale weights can negatively affect the statistical efficiency of the network , preventing the convergence of training or producing a model with lower inference accuracy . A0 A1 A2 A3 P0 P1 P2 P3 Idle Time P0 P1 P2 P3 Forward Backward Figure 2 : Schedule of Computations Common wisdom is that the use of stale weights must either be avoided , e.g. , with the use of microbatches ( Huang et al. , 2018 ) , be constrained to ensure the consistency of the weights within an accelerator using stashing ( Harlap et al. , 2018 ) , or by limiting the use of pipelining to very small networks ( Mostafa et al. , 2017 ) . However , these approaches either underutilize accelerators ( Huang et al. , 2018 ) or inflate memory usage to stash multiple copies of weights ( Harlap et al. , 2018 ) . In this paper we question this common wisdom and explore pipelining that allows for the full utilization of accelerators while using stale weights . This results in a pipelining scheme that , compared to existing schemes , is simpler to implement , fully utilizes the accelerators and has lower memory overhead . We evaluate this pipelining scheme using 4 CNNs : LeNet-5 ( trained on MNIST ) , AlexNet , VGG and ResNet ( all trained on CIFAR-10 ) . We analyze the impact of weight staleness and show that if pipelining is limited to early layers in the network , training does converge and the quality of the resulting models is comparable to that of models obtained with non-pipelined training . For the 4 networks , the drop in accuracy is 0.4 % , 4 % , 0.83 % and 1.45 % , respectively . However , inference accuracies drop significantly when the pipelining is deeper in the network . While this is not a limitation since the bulk of computations that can benefit from pipelining are in the early convolutional layers , we address this through a hybrid scheme that combines pipelined and non-pipelined training to maintain inference accuracy while still delivering performance improvement . Evaluation shows that our pipelined training delivers a speedup of up to 1.8X on a 2-GPU system . The remainder of this paper is organized as follows . Section 2 briefly describes the backpropagation for training of CNNs . Section 3 details our pipelining scheme . Section 4 describes how non-pipelined and pipelined backpropagation are combined . Section 5 highlights some of the implementation details . Experimental evaluation is presented in Section 6 . Related work is reviewed in Section 7 . Finally , Section 8 gives concluding remarks and directions for future work . 2 THE BACKPROPAGATION ALGORITHM . The backpropagation algorithm ( Rumelhart et al. , 1986 ) consists of two passes : a forward pass that calculates the output error and a backward pass that calculates the error gradients and updates the weights of the network . The two passes are performed for input data one mini-batch at a time . In the forward pass , a mini-batch is fed into the network , propagating from the first to the last layer . At each layer l , the activations of the layer , denoted by x ( l ) , are computed using the weights of the layer , denoted by W ( l ) . When the output of the network ( layer L ) x ( L ) is produced , it is used with the true data label to obtain a training error e for the mini-batch . In the backward pass , the error e is propagated from the last to the first layer . The error gradients with respect to pre-activations of layer l , denoted by δ ( l ) , are calculated . Further , the error gradients with respect to weights of layer l , ∂e ∂W ( l ) , are computed using the activations from layer l − 1 ( i.e. , x ( l−1 ) ) and δ ( l ) . Subsequently , δ ( l ) is used to calculate the δ ( l−1 ) . When ∂e ∂W ( l ) is computed for every layer , the weights are updated using the error gradients . In the forward pass , the activations of the layer l , x ( l ) , can not be computed until the activations of the previous layers , i.e. , x ( l−1 ) , are computed . In backward pass , ∂e ∂W ( l ) can only be computed once x ( l−1 ) and δ ( l ) have been computed . Moreover , δ ( l ) depends on δ ( l+1 ) . Finally , for a given mini-batch the backward pass can not be started until the forward pass is completed and the error e has been determined . The above dependences ensure that the weights of the layers are updated using the activations and error gradients calculated from the same batch of training data in one iteration of the backpropagation algorithm . Only when the weights are updated is the next batch of training data fed into the network . These dependences limit parallelism when a network is partitioned across multiple accelerators and allow only one accelerator to be active at any point . This results in under-utilization of the accelerators . It is this limitation that pipelining addresses . 3 PIPELINED BACKPROPAGATION . We illustrate our pipelined backpropagation implementation with the L layer network shown in Figure 3 , using conceptual pipeline registers . Two registers are inserted between layers l and l + 1 ; one register for the forward pass and a second for the backward pass . The forward register stores the activations of layer l ( x ( l ) ) . The backward register stores the gradients δ ( l+1 ) of layer l+1 . This defines a 4-stage pipelined backpropagation . The forward pass for layers 1 to l forms forward stage FS1 . The forward pass for layers l + 1 to L form forward stage FS2 . Similarly , the backwards pass for layers l + 1 to L and 1 to l form backward stages BKS1 and BKS2 respectively . The forward and backward stages are executed in a pipelined fashion on 3 accelerators : one for FS1 , one for both FS2 and BKS1 , and one for BKS21 . In cycle 0 , mini-batch 0 is fed to FS1 . The computations of the forward pass are done as in the traditional non-pipelined implementation . In cycle 1 , layer l activations x ( l ) are fed to FS2 and mini-batch 1 is fed to FS1 . In cycle 2 , the error for mini-batch 0 computed in FS2 is directly fed to BKS1 , the activations of layer l x ( l ) are forwarded to FS2 and mini-batch 2 is fed to FS1 . This pipelined execution is illustrated by the space-time diagram in Figure 4 for 5 mini-batches . The figure depicts the mini-batch processed by each accelerator cycles 0 to 6 . At steady state , all the accelerators are active in each cycle of execution . The above pipelining scheme utilizes weights in FS1 that are yet to be updated by the errors calculated by FS2 and BKS1 . At steady state , the activations of a mini-batch in FS1 are calculated using weights that are 2 execution cycles old , or 2 cycles stale . This is reflected in Figure 4 by indicating the weights used by each forward stage and the weights updated by each backward stage . The weights of a forward stage are subscripted by how stale they are ( -ve subscripts ) . Similarly , the weights updated by a backward stage are subscripted by how delayed they are ( +ve subscripts ) . Further , since the updates of the weights by BKS2 requires activations calculated for the same mini-batch in FS1 for all layers in the stage , it is necessary to save these activations until the error gradients with respect to the weights are calculated by BKS2 . Only when the weights are updated using the gradients can these activations be discarded . In the general case , we use K pairs of pipeline registers ( each pair consisting of a forward register and a backward register ) inserted between the layers of the network . We describe the placement of the register pairs by the Pipeline Placement Vector , PPV = ( p1 , p2 , ... , pK ) , where pi represents the layer number after which a pipeline register pair is inserted . Such a placement creates ( K + 1 ) forward stages , labeled FSi , i = 1 , 2 , ... , K + 1 and ( K + 1 ) backward stages , labeled BKSi , i = 1 , 2 , ... , K + 1 . Forward stage FSi and backward stage BKSK−i+2 correspond to the same set of layers . Specifically , stage FSi contains layers pi + 1 to pi+1 , inclusive . We assign each forward stage and each backward stage to an accelerator , with the exception of the FSK+1 and backward stage BKS1 , which are assigned to the same accelerator to reduce weight staleness by an execution cycle . In total 2K + 1 accelerators are used . We quantify weight staleness as follows . A forward stage FSi and backward stage BKSK−i+2 use the same weights that are 2 ( K− i+1 ) cycles old . A forward stage FSi must store the activations of 1We combine FS1 and BKS1 on the same accelerator to reduce weight staleness . all layer in the stage for all 2 ( K− i+1 ) cycles which are used for the corresponding backward stage BKSK−i+2 . Thus , we define the Degree of Staleness as 2 ( K − i + 1 ) , and these saved activations as intermediate activations . For each pair of stages FSi and BKSK−i+2 , let there be Ni weights in their corresponding layers . The layers before the last pipeline register pairs always use stale weights . Thus , we define Percentage of Stale Weight as ( ∑K i=1Ni ) / ( ∑K+1 i=1 Ni ) . On the one hand , the above pipelined execution allows a potential speedup of 2K + 1 over the nonpipelined implementation , keeping all the accelerators active at steady state . On the other hand , the use of stale weights may prevent training convergence or may result in a model that has an inferior inference accuracy . Further , it requires an increase in storage for activations . Our goal is to assess the benefit of this pipelined execution and the impact of its down sides .

Summary: This paper proposes a new pipelined training approach to speedup the training for neural networks. The approach separates forward and backpropagation processes into multiple stages, cache the activation and gradients between stages, processes stages simultaneously, and then uses the stored activations to compute gradients for updating the weights. The approach leads to stale weights and gradients. The authors studied the relation between weight staleness and show that the quality degradation mainly correlates with the percentage of the weights being stale in the pipeline. The quality degradation can also be remedied by turning off the pipelining at the later training steps while overall training speed is still faster than without pipelined training.