Adaptive Gradient Quantization with Bit Allocation for Distributed Deep Learning

Gradient compression plays a crucial role in mitigating communication overhead in distributed deep learning. Existing gradient compression methods usually employ fix-bit quantization across all layers, neglecting the varying sensitivities of different layers to compression, resulting in suboptimal performance. In this paper, we introduce a layer-wise bit allocation mechanism for gradient quantization that minimizes overall quantization error within a specified bit budget. To address the heavy computational load of conventional greedy search approach for bit allocation, we develop two acceleration techniques to reduce computational overhead, thereby making the proposed bit allocation method feasible for real-time deep learning training. Specifically, by observing the bit allocation statistics, we propose Bit Searching Range Optimization to narrow the available bit options, while the Bit Pre-Assignment selectively bypasses certain searching processes. Experimental results across various neural network models and datasets demonstrate the effectiveness of our proposed bit allocation methods for gradient quantization. The combination of proposed acceleration techniques offers an advantageous trade-off among quantization error, model training performance and time consumption. Moreover, our proposed bit allocation methods can be seamlessly integrated with existing gradient compression approaches, improving overall performance.