SHIRO: Near-Optimal Communication Strategies for Distributed Sparse Matrix Multiplication

TL;DR

SHIRO introduces a sparsity-aware and hierarchical communication strategy for distributed sparse matrix multiplication, significantly reducing communication overhead and improving scalability in high-performance computing and deep learning.

Contribution

It presents a novel framework that exploits sparsity patterns and network hierarchy to optimize communication in distributed SpMM, achieving near-optimal performance.

Findings

Achieves up to 221.5x speedup over state-of-the-art baselines.

Demonstrates scalability up to 128 GPUs.

Reduces communication overhead through sparsity-aware and hierarchical strategies.

Abstract

Distributed Sparse Matrix-Matrix Multiplication (SpMM) is a fundamental operation in high-performance computing and deep learning applications. The major performance bottleneck in distributed SpMM lies in substantial communication overhead, which limits both performance and scalability. In this paper, we identify two key sources of communication inefficiency in distributed SpMM: redundant data transfer due to sparsity unawareness, and suboptimal utilization of hierarchical network topology. To address these, we propose (1) a fine-grained, sparsity-aware communication strategy that reduces communication overhead by exploiting the sparsity pattern of the sparse matrix, and (2) a hierarchical communication strategy that maps the sparsity-aware strategy onto two-tier GPU network architectures, minimizing redundant data movement across slower inter-node links. We implement these optimizations in a comprehensive distributed SpMM framework, \method{}. Extensive evaluations on real-world datasets show that \method{} demonstrates strong scalability up to 128 GPUs, achieving geometric mean speedups of 221.5$\times$, 56.0$\times$, 23.4$\times$, and 8.8$\times$ in SpMM over four state-of-the-art baselines (CAGNET, SPA, BCL, and CoLa, respectively) at this scale.