Tensor Cores上的高效量化稀疏矩阵运算
摘要
指数级增长的模型规模推动了深度学习的持续成功，但也带来了高昂的计算与内存成本。从算法角度看，模型稀疏化与量化已被研究以缓解该问题。从架构角度看，硬件厂商提供了Tensor cores进行加速。然而，由于在Tensor cores上严格的数据布局要求以及缺乏对低精度整数高效操作的支持，从稀疏、低精度矩阵运算中获得实际加速非常困难。我们提出Magicube，一个在Tensor cores上面向低精度整数的高性能稀疏矩阵库。Magicube支持SpMM和SDDMM，即深度学习中两种主要的混合精度稀疏运算。在NVIDIA A100 GPU上的实验结果表明，Magicube相较厂商优化的稀疏核库平均实现1.44倍（最高2.37倍）加速，并在端到端稀疏Transformer推理中以可比精度相较最先进方法实现1.43倍加速。
引用
@article{arxiv.2209.06979,
title = {Efficient Quantized Sparse Matrix Operations on Tensor Cores},
author = {Shigang Li and Kazuki Osawa and Torsten Hoefler},
journal= {arXiv preprint arXiv:2209.06979},
year = {2023}
}
备注
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 $\textbf{K/BS}$$_k$ steps to obtain the final results; we fix Table 3.)