PacTrain: Pruning and Adaptive Sparse Gradient Compression for Efficient Collective Communication in Distributed Deep Learning

TL;DR

PacTrain addresses the gradient aggregation bottleneck in distributed deep learning by fusing unstructured pruning with adaptive, all-reduce–friendly sparse gradient compression. A Mask Tracker and gradient sparsity enforcement align sparsity patterns across distributed workers, enabling lightweight communication without compromising accuracy. The approach, including pruning-aware training and ternary gradient quantization, achieves substantial end-to-end speedups—up to 8.72x—under bandwidth-constrained conditions while keeping accuracy within a small margin for pruning up to 80%. This makes scalable training of large vision and language models more practical in heterogeneous network environments without sacrificing convergence or final performance.

Abstract

Large-scale deep neural networks (DNN) exhibit excellent performance for various tasks. As DNNs and datasets grow, distributed training becomes extremely time-consuming and demands larger clusters. A main bottleneck is the resulting gradient aggregation overhead. While gradient compression and sparse collective communication techniques are commonly employed to alleviate network load, many gradient compression schemes do not achieve acceleration of the training process while also preserving accuracy. This paper introduces PacTrain, a novel framework that accelerates distributed training by combining pruning with sparse gradient compression. Active pruning of the neural network makes the model weights and gradients sparse. By ensuring the global knowledge of the gradient sparsity among all distributed training workers, we can perform lightweight compression communication without harming accuracy. We show that the PacTrain compression scheme achieves a near-optimal compression strategy while remaining compatible with the all-reduce primitive. Experimental evaluations show that PacTrain improves training throughput by 1.25 to 8.72 times compared to state-of-the-art compression-enabled systems for representative vision and language models training tasks under bandwidth-constrained conditions.

Figures (6)

• Figure 1: Pruning and Fine-tune on specific data.
• Figure 2: Acceleration with masked assignment.
• Figure 3: End-to-End TTA speedup with different WAN bandwidths (relative to native all-reduce, log scale).
• Figure 4: Evaluation topology.
• Figure 5: Time-to-accuracy Comparison of performance for the CIFAR-10 classification task using ResNet152.