Gradient pruning to accelerate ring-based all-reduce in distributed deep learning

Abstract

Data-parallelism has been widely used in order to distribute the training of large Deep Neural Networks (DNN) among multiple workers (GPU/TPU). In the data-parallel context, a large proportion of the training time is spent on collective communication to synchronize gradients between model replicas on each device. It was observed that not all gradients are necessary for the model convergence and the gradient tensor could be greatly sparsified in order to reduce the communication volume. Prior works have proposed to sparsify gradients depending on their magnitude, approach commonly called top-k. However, top-k incurs computation overhead of bitonic sorting and has scalability issues on Ring All-Reduce due to the necessity to send indices. In this work, we show that random gradient pruning can achieve convergence at minimal accuracy loss. Moreover, we propose Skip-Reduce -- a novel approach to perform gradient pruning that does not rely on sending indices and instead modifies the underlying communication algorithm. Our approach has no computational overhead and is scalable on any number of devices.