Enabling Efficient All-Dimension Top-k Sparsification for High-Performance Distributed DNN Training Systems

Abstract

Distributed deep neural network (DNN) training across multiple workers is widely adopted to accelerate the training of large models and datasets. However, their performance is often limited by the synchronization of inter-worker data (i.e., gradient tensors). Top-k sparsification serves as the de facto compression approach to alleviate the communication bottleneck by retaining only the k largest-magnitude gradient elements before transmission to other nodes. However, traditional top-k sparsification has two key limitations. First, gradients at each DNN layer are represented as multi-dimensional tensors, with the selected k elements typically concentrated in a subset of dimensions. This leads to many dimensions being missed during training, and hence poor convergence. Second, traditional top-k sparsification relies on the global sorting of all gradient elements per layer, leading to underutilized GPU core parallelism and low training throughput. To this end, we propose ADTopk, an all-dimension top-k sparsification scheme that selects the k largest elements across all dimensions of the per-layer gradient tensor, ensuring that every dimension contributes at least one element and eliminating dimension missing. ADTopk also enables independent local sorting within each dimension, allowing parallel execution across dimensions to enhance GPU utilization. Furthermore, we enhance ADTopk with system-level optimizations: (i) interleaved sparsification to accelerate convergence, (ii) partial sparsification to reduce sparsification overhead, and (iii) hybrid collective communication to improve sparse communication efficiency. Experiments on local and cloud clusters show that ADTopk improves convergence accuracy by 2.75%-22.31% and achieves 20.8%-268.0% higher training throughput than state-of-the-art sparsification approaches.