Efficient Quantized Sparse Matrix Operations on Tensor Cores

Abstract

The exponentially growing model size drives the continued success of deep learning, but it brings prohibitive computation and memory cost. From the algorithm perspective, model sparsification and quantization have been studied to alleviate the problem. From the architecture perspective, hardware vendors provide Tensor cores for acceleration. However, it is very challenging to gain practical speedups from sparse, low-precision matrix operations on Tensor cores, because of the strict requirements for data layout and lack of support for efficiently manipulating the low-precision integers. We propose Magicube, a high-performance sparse-matrix library for low-precision integers on Tensor cores. Magicube supports SpMM and SDDMM, two major sparse operations in deep learning with mixed precision. Experimental results on an NVIDIA A100 GPU show that Magicube achieves on average 1.44x (up to 2.37x) speedup over the vendor-optimized library for sparse kernels, and 1.43x speedup over the state-of-the-art with a comparable accuracy for end-to-end sparse Transformer inference.

Published in Proceedings of 2022 International Conference for High Performance Computing, Networking, Storage and Analysis (SC'22), No.: 37, Pages 1-15, Best Paper Finalist, https://dl.acm.org/doi/10.5555/3571885.3571934 (In this arXiv verion, we fix a typo at the bottom right of Page 6: For SDDMM, each thread block needs $\text{K/BS}$${}_k$ steps to obtain the final results; we fix Table 3.)

References in corpus (11)

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Scaling Laws for Neural Language Models
• Linformer: Self-Attention with Linear Complexity
• Deep Learning Scaling is Predictable, Empirically
• Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks
• Exploring the Regularity of Sparse Structure in Convolutional Neural Networks
• Integer Quantization for Deep Learning Inference: Principles and Empirical Evaluation
• Long Range Arena: A Benchmark for Efficient Transformers
• Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines
• Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning
• Joint Pruning & Quantization for Extremely Sparse Neural Networks

Cited by in corpus (3)

• Efficient Mixed Precision Quantization in Graph Neural Networks
• MO-SAE:Multi-Objective Stacked Autoencoders Optimization for Edge Anomaly Detection
• BLEST: Blazingly Efficient BFS using Tensor Cores