Sparse GPU Kernels for Deep Learning
Trevor Gale Stanford University United States of America Matei Zaharia Stanford University United States of America Cliff Young Google Brain United States of America Erich Elsen DeepMind United Kingdom
Abstract
Scientific workloads have traditionally exploited high levels of sparsity to accelerate computation and reduce memory requirements. While deep neural networks can be made sparse, achieving practical speedups on GPUs is difficult because these applications have relatively moderate levels of sparsity that are not sufficient for existing sparse kernels to outperform their dense counterparts. In this work, we study sparse matrices from deep learning applications and identify favorable properties that can be exploited to accelerate computation. Based on these insights, we develop high-performance GPU kernels for two sparse matrix operations widely applicable in neural networks: sparse matrix–dense matrix multiplication and sampled dense–dense matrix multiplication. Our kernels reach 27% of single-precision peak on Nvidia V100 GPUs. Using our kernels, we demonstrate sparse Transformer and MobileNet models that achieve 1.2–2.1 $\times$ speedups and up to 12.8 $\times$ memory savings without sacrificing accuracy.
中文速览
深度神经网络的权重矩阵可以被稀疏化以节省计算和内存，但现有GPU稀疏计算库（如cuSPARSE）主要针对科学计算中99%以上极高稀疏度的矩阵设计，对深度学习中常见的中等稀疏度（约50%–95%）矩阵根本快不过稠密计算。作者系统分析了深度学习稀疏矩阵的统计规律，发现它们的行长度更长、分布更均匀、稀疏度更低，并据此从零设计了两个GPU核函数：用于前向传播的稀疏矩阵-稠密矩阵乘法（SpMM）和用于反向传播及稀疏注意力的采样稠密-稠密矩阵乘法（SDDMM），引入了一维分块、子warp切分、反偏移内存对齐和行乱序负载均衡等关键技术。在V100 GPU上，这两个核函数分别比cuSPARSE快3.58倍和2.19倍，峰值达到单精度理论算力的27%，并在稀疏Transformer和MobileNet模型上实现了1.2–2.1倍端到端推理加速和最高12.8倍内存节省，且精度不损失，为在实际硬件上高效利用神经网络稀疏性奠定了基础。
原文 arXiv:2006.10901;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2006.10901v2