BPPSA: Scaling Back-propagation by Parallel Scan Algorithm

Abstract

In an era when the performance of a single compute device plateaus, software must be designed to scale on massively parallel systems for better runtime performance. However, in the context of training deep learning models, the popular back-propagation (BP) algorithm imposes a strong sequential dependency in the process of gradient computation. Under model parallelism, BP takes $\Theta (n)$ steps to complete which hinders its scalability on parallel systems ($n$ represents the number of compute devices into which a model is partitioned). In this work, in order to improve the scalability of BP, we reformulate BP into a scan operation which is a primitive that performs an in-order aggregation on a sequence of values and returns the partial result at each step. We can then scale such reformulation of BP on parallel systems by our modified version of the Blelloch scan algorithm which theoretically takes $\Theta (\log n)$ steps. We evaluate our approach on a vanilla Recurrent Neural Network (RNN) training with synthetic datasets and a RNN with Gated Recurrent Units (GRU) training with the IRMAS dataset, and demonstrate up to $2.75\times$ speedup on the overall training time and $108\times$ speedup on the backward pass. We also demonstrate that the retraining of pruned networks can be a practical use case of our method.

arXiv commentIn Proceedings of MLSys 2020: https://mlsys.org/Conferences/2020/Schedule?showEvent=1407

References in corpus (11)

• Deep Residual Learning for Image Recognition
• Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• Densely Connected Convolutional Networks
• cuDNN: Efficient Primitives for Deep Learning
• One weird trick for parallelizing convolutional neural networks
• Training Deep Nets with Sublinear Memory Cost
• Channel Pruning for Accelerating Very Deep Neural Networks
• FMA: A Dataset For Music Analysis
• Measuring the Effects of Data Parallelism on Neural Network Training
• Optimizing Performance of Recurrent Neural Networks on GPUs

Cited by in corpus (1)

• An Optimized and Energy-Efficient Parallel Implementation of Non-Iteratively Trained Recurrent Neural Networks