Pipelined Training with Stale Weights of Deep Convolutional Neural Networks

TL;DR

This paper investigates the effects of using stale weights in pipelined training of CNNs, showing that limited pipelining maintains accuracy and improves training speed, while deeper pipelining causes accuracy drops.

Contribution

It introduces a hybrid pipelined training scheme that balances accelerator utilization and accuracy, demonstrating practical implementation and performance gains.

Findings

Limited pipelining preserves accuracy within 1.45% of non-pipelined training.

Hybrid scheme mitigates accuracy loss when pipelining deeper layers.

Achieves up to 1.8X speedup on 2 GPUs with minimal accuracy drop.

Abstract

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.