分布式内存机器学习稀疏核
摘要
采样稠密乘稠密矩阵乘法（SDDMM）与稀疏乘稠密矩阵乘法（SpMM）出现在多样的设置中，如协同过滤、文档聚类和图嵌入。经常地，SDDMM 的输出成为后续 SpMM 操作的输入稀疏矩阵。现有工作聚焦于这些原语的共享内存并行化。虽然已有大量关于 SpMM 通信最小化分布式 1.5D 算法的分析，但对于 SDDMM 或 SDDMM 与 SpMM 的背靠背序列（称为 FusedMM）尚无此类分析。我们展示了用于 SpMM 的分布式内存 1.5D 和 2.5D 算法可转换为用于 SDDMM 的算法，且通信成本与输入/输出数据布局相同。此外，我们给出两种通信省略策略以进一步降低 FusedMM 核的成本：要么复用输入稠密矩阵对 SDDMM 和 SpMM 的复制，要么融合局部 SDDMM 和 SpMM 核。我们在 Cori（LBNL 的 Cray XC40）上使用 Erdos-Renyi 随机矩阵和大型真实世界稀疏矩阵对 FusedMM 算法进行基准测试。在 256 个节点、每节点 68 核上，采用任一通信省略方法的 1.5D FusedMM 算法相比顺序执行分布式内存 SpMM 和 SDDMM 核可节省至少 30% 的纯通信时间。在具有数亿条边的真实世界矩阵上，我们所有算法相比 PETSc 中的 SpMM 算法均展现出至少 10 倍加速。在这些矩阵上，我们的通信省略技术运行时间比未优化的 SDDMM 与 SpMM 序列快达 1.6 倍。我们将算法嵌入真实世界应用测试其扩展性，包括基于交替最小二乘的协同过滤和基于注意力的图神经网络推理。
引用
@article{arxiv.2203.07673,
title = {Distributed-Memory Sparse Kernels for Machine Learning},
author = {Vivek Bharadwaj and Aydin Buluç and James Demmel},
journal= {arXiv preprint arXiv:2203.07673},
year = {2022}
}
备注
To appear at the 36th IEEE International Parallel & Distributed Processing Symposium (IPDPS'22). 12 pages, 9 figures