Multi-stage gradient compression: Overcoming the communication bottleneck in distributed deep learning

Abstract

Due to the huge size of deep learning model and the limited bandwidth of network, communication cost has become a salient bottleneck in distributed training. Gradient compression is an effective way to relieve the pressure of bandwidth and increase the scalability of distributed training. In this paper, we propose a novel gradient compression technique, Multi-Stage Gradient Compression (MGC) with Sparsity Automatic Adjustment and Gradient Recession. These techniques divide the whole training process into three stages which fit different compression strategy. To handle error and preserve accuracy, we accumulate the quantization error and sparse gradients locally with momentum correction. Our experiments show that MGC achieves excellent compression ratio up to 3800x without incurring accuracy loss. We compress gradient size of ResNet-50 from 97 MB to 0.03 MB, for AlexNet from 233 MB to 0.06 MB. We even get a better accuracy than baseline on GoogLeNet. Experiments also show the significant scalability of MGC.