TernGrad: Ternary Gradients to Reduce Communication in Distributed Deep Learning

TL;DR

The paper addresses the bottleneck of gradient communication in distributed data-parallel training by introducing TernGrad, which quantizes gradients to ternary values and uses a shared scaler along with parameter localization to reduce communication. It provides a convergence analysis under a gradient-bound assumption and offers practical improvements via layer-wise ternarizing and gradient clipping. Empirical results across MNIST, CIFAR-10, and ImageNet demonstrate near-baseline accuracy with significant communication-time reductions and substantial throughput gains, especially for networks with high communication-to-computation ratios. A performance model further quantifies scalability and speedups on GPU clusters under various network configurations.

Abstract

High network communication cost for synchronizing gradients and parameters is the well-known bottleneck of distributed training. In this work, we propose TernGrad that uses ternary gradients to accelerate distributed deep learning in data parallelism. Our approach requires only three numerical levels {-1,0,1}, which can aggressively reduce the communication time. We mathematically prove the convergence of TernGrad under the assumption of a bound on gradients. Guided by the bound, we propose layer-wise ternarizing and gradient clipping to improve its convergence. Our experiments show that applying TernGrad on AlexNet does not incur any accuracy loss and can even improve accuracy. The accuracy loss of GoogLeNet induced by TernGrad is less than 2% on average. Finally, a performance model is proposed to study the scalability of TernGrad. Experiments show significant speed gains for various deep neural networks. Our source code is available.

Figures (5)

• Figure 1: Distributed SGD with data parallelism.
• Figure 2: Histograms of (a) original floating gradients, (b) clipped gradients, (c) ternary gradients and (d) final averaged gradients. Visualization by TensorBoard. The DNN is AlexNet distributed on two workers, and vertical axis is the training iteration. As examples, top row visualizes the third convolutional layer and bottom one visualizes the first fully-connected layer.
• Figure 3: Accuracy vs. worker number for baseline and TernGrad, trained with (a) momentum SGD or (b) vanilla SGD. In all experiments, total mini-batch size is $64$ and maximum iteration is $10K$.
• Figure 4: AlexNet trained on $4$ workers with mini-batch size $512$: (a) top-1 validation accuracy, (b) training data loss and (c) sparsity of gradients in first fully-connected layer (fc6) vs. iteration.
• Figure 5: Training throughput on two different GPUs clusters: (a) 128-node GPU cluster with 1Gbps Ethernet, each node has 4 NVIDIA GTX 1080 GPUs and one PCI switch; (b) 128-node GPU cluster with 100 Gbps InfiniBand network connections, each node has 4 NVIDIA Tesla P100 GPUs connected via NVLink. Mini-batch size per GPU of AlexNet, GoogLeNet and VggNet-A is $128$, $64$ and $32$, respectively

Theorems & Definitions (3)

• Lemma 1
• Theorem 1
• proof