分布式内存机器学习稀疏核

采样稠密乘稠密矩阵乘法（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 倍。我们将算法嵌入真实应用测试其可扩展性，包括基于交替最小二乘的协同过滤和基于注意力的图神经网络推理。