Partitioning sparse deep neural networks for scalable training and inference
This work addresses scalability challenges in deep learning for researchers and practitioners dealing with large, sparse models, though it is incremental as it builds on existing sparsification and parallelization methods.
The paper tackles the computational inefficiency of training and inference in sparse deep neural networks by introducing a distributed-memory parallel solution based on sparse matrix-vector multiplications and a novel hypergraph model for partitioning weight matrices, resulting in significant performance improvements and scalability as demonstrated in experiments.
The state-of-the-art deep neural networks (DNNs) have significant computational and data management requirements. The size of both training data and models continue to increase. Sparsification and pruning methods are shown to be effective in removing a large fraction of connections in DNNs. The resulting sparse networks present unique challenges to further improve the computational efficiency of training and inference in deep learning. Both the feedforward (inference) and backpropagation steps in stochastic gradient descent (SGD) algorithm for training sparse DNNs involve consecutive sparse matrix-vector multiplications (SpMVs). We first introduce a distributed-memory parallel SpMV-based solution for the SGD algorithm to improve its scalability. The parallelization approach is based on row-wise partitioning of weight matrices that represent neuron connections between consecutive layers. We then propose a novel hypergraph model for partitioning weight matrices to reduce the total communication volume and ensure computational load-balance among processors. Experiments performed on sparse DNNs demonstrate that the proposed solution is highly efficient and scalable. By utilizing the proposed matrix partitioning scheme, the performance of our solution is further improved significantly.