BPPSA: Scaling Back-propagation by Parallel Scan Algorithm

TL;DR

This paper introduces BPPSA, a novel method that reformulates back-propagation as a scan operation, enabling scalable parallel training of deep learning models with significant speedups.

Contribution

It presents a new reformulation of back-propagation as a scan operation and a modified parallel scan algorithm to improve scalability on parallel systems.

Findings

Up to 2.75x speedup in overall training time

108x speedup in backward pass

Effective for RNN and pruned network retraining

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.