arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2203.07673cs.DCcs.LG
收藏

用于机器学习的分布式存储稀疏核

Distributed-Memory Sparse Kernels for Machine Learning

University of California, Berkeley(加州大学伯克利分校);Lawrence Berkeley National Laboratory(劳伦斯伯克利国家实验室)

Vivek Bharadwaj, Aydin Buluç, James Demmel

摘要:

采样稠密时间稠密矩阵乘法(SDDMM)和稀疏时间稠密矩阵乘法(SpMM)出现在不同的设置中,如协同过滤,文档聚类和图嵌入。通常,SDDMM输出成为后续SpMM操作的输入稀疏矩阵。现有的工作集中在共享内存并行这些原语。虽然已经有广泛的分析,通信最小化分布式1.5D算法的SpMM,没有这样的分析存在SDDMM或背靠背序列的SDDMM和SpMM,称为融合MM。我们发现,分布式存储器1.5D和2.5D算法的SpMM可以转换为算法SDDMM具有相同的通信成本和输入/输出数据布局。此外,我们给出了两个通信省略策略,以进一步降低成本的FusionMM内核:要么重用SDDMM和SpMM的输入稠密矩阵的复制顺序,或融合本地SDDMM和SpMM内核。我们基准测试FusedMM算法Cori,Cray XC 40在LBNL,使用Erdos-Renyi随机矩阵和大型现实世界的稀疏矩阵。在256个节点上,每个节点有68个内核,与依次执行分布式内存SpMM和SDDMM内核相比,使用任何一种通信省略方法的1.5D FusedMM算法可以节省至少30%的通信时间。在具有数亿条边的真实矩阵上,我们所有的算法都比PETSc中的SpMM算法至少有10倍的加速。在这些矩阵上,我们的通信省略技术的运行时间比未优化的SDDMM和SpMM序列快1.6倍。我们嵌入和测试我们的算法在现实世界中的应用,包括协同过滤通过交替最小二乘法和推理的注意力为基础的图神经网络的缩放。

英文摘要:

Sampled Dense Times Dense Matrix Multiplication (SDDMM) and Sparse Times Dense Matrix Multiplication (SpMM) appear in diverse settings, such as collaborative filtering, document clustering, and graph embedding. Frequently, the SDDMM output becomes the input sparse matrix for a subsequent SpMM operation. Existing work has focused on shared memory parallelization of these primitives. While there has been extensive analysis of communication-minimizing distributed 1.5D algorithms for SpMM, no such analysis exists for SDDMM or the back-to-back sequence of SDDMM and SpMM, termed FusedMM. We show that distributed memory 1.5D and 2.5D algorithms for SpMM can be converted to algorithms for SDDMM with identical communication costs and input / output data layouts. Further, we give two communication-eliding strategies to reduce costs further for FusedMM kernels: either reusing the replication of an input dense matrix for the SDDMM and SpMM in sequence, or fusing the local SDDMM and SpMM kernels. We benchmark FusedMM algorithms on Cori, a Cray XC40 at LBNL, using Erdos-Renyi random matrices and large real-world sparse matrices. On 256 nodes with 68 cores each, 1.5D FusedMM algorithms using either communication eliding approach can save at least 30% of time spent exclusively in communication compared to executing a distributed-memory SpMM and SDDMM kernel in sequence. On real-world matrices with hundreds of millions of edges, all of our algorithms exhibit at least a 10x speedup over the SpMM algorithm in PETSc. On these matrices, our communication-eliding techniques exhibit runtimes up to 1.6 times faster than an unoptimized sequence of SDDMM and SpMM. We embed and test the scaling of our algorithms in real-world applications, including collaborative filtering via alternating-least-squares and inference for attention-based graph neural networks.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑