Accelerating Distributed Deep Learning using Lossless Homomorphic Compression

TL;DR

The paper tackles the gradient communication bottleneck in distributed deep learning by pairing lossless homomorphic compression with in-network aggregation. It introduces a two-phase scheme that uses Count Sketch for homomorphic compression and a peeling-based recovery augmented by a Bloom Filter index to achieve exact reconstruction of summed gradients without sacrificing accuracy. The approach claims asymptotically optimal compression and linear-time compression/decompression, achieving up to $6.33\times$ aggregation throughput and $3.74\times$ per-iteration speedup across diverse models when integrated with NCCL and ATP, while maintaining convergence. This combination promises scalable distributed training by leveraging both distributed compression and in-network aggregation, with an open-source implementation for practical adoption.

Abstract

As deep neural networks (DNNs) grow in complexity and size, the resultant increase in communication overhead during distributed training has become a significant bottleneck, challenging the scalability of distributed training systems. Existing solutions, while aiming to mitigate this bottleneck through worker-level compression and in-network aggregation, fall short due to their inability to efficiently reconcile the trade-offs between compression effectiveness and computational overhead, hindering overall performance and scalability. In this paper, we introduce a novel compression algorithm that effectively merges worker-level compression with in-network aggregation. Our solution is both homomorphic, allowing for efficient in-network aggregation without CPU/GPU processing, and lossless, ensuring no compromise on training accuracy. Theoretically optimal in compression and computational efficiency, our approach is empirically validated across diverse DNN models such as NCF, LSTM, VGG19, and BERT-base, showing up to a 6.33$\times$ improvement in aggregation throughput and a 3.74$\times$ increase in per-iteration training speed.

Figures (8)

• Figure 1: An example of our lossless homomorphic compression algorithm. In this example, three non-zero parameters - $x, y,$ and $z$ - identified by the index, are hashed and added to the Count Sketch. Ignoring the aggregation process and the parallel computation for the moment, to recover the data from the index and the Count Sketch, we first peel away $z$, since its value uniquely hashes to the bottom counter of the Count Sketch. Following this removal, $x$ and $y$ become eligible for peeling. This sequence allows for the lossless recovery of the entire gradient.
• Figure 2: Theoretical compressed data size.
• Figure 3: Average relative error, recovery rate, and iterations.
• Figure 4: Training loss and test accuracy.
• Figure 5: Aggregation throughput of NCCL-based experiments.