Communication-Avoiding Linear Algebraic Kernel K-Means on GPUs

TL;DR

The paper tackles scaling exact Kernel K-means to million-scale datasets by mapping the clustering workflow onto distributed, communication-optimized linear algebra primitives. It introduces four distributed formulations (1D, Hybrid 1D, 2D, and 1.5D) and demonstrates that a careful 1.5D partitioning strategy minimizes communication by keeping the sparse matrix $\mathbf{V}$ in 1D while using a 2D distribution for the kernel $\mathbf{K}$ produced by SUMMA. The main contributions are the first open-source fully GPU-accelerated distributed implementation of exact Kernel K-means (Vivaldi), a set of new parallel algorithms for Kernel K-means based on distributed linear algebra primitives, and substantial performance gains (up to $3.6\times$ over the baseline and up to $>2000\times$ over a sliding-window single-GPU approach) enabling clustering of over $1.5$ million points. These results highlight the practical impact of application-guided primitive composition for end-to-end performance on large-scale clustering tasks.

Abstract

Clustering is an important tool in data analysis, with K-means being popular for its simplicity and versatility. However, it cannot handle non-linearly separable clusters. Kernel K-means addresses this limitation but requires a large kernel matrix, making it computationally and memory intensive. Prior work has accelerated Kernel K-means by formulating it using sparse linear algebra primitives and implementing it on a single GPU. However, that approach cannot run on datasets with more than approximately 80,000 samples due to limited GPU memory. In this work, we address this issue by presenting a suite of distributed-memory parallel algorithms for large-scale Kernel K-means clustering on multi-GPU systems. Our approach maps the most computationally expensive components of Kernel K-means onto communication-efficient distributed linear algebra primitives uniquely tailored for Kernel K-means, enabling highly scalable implementations that efficiently cluster million-scale datasets. Central to our work is the design of partitioning schemes that enable communication-efficient composition of the linear algebra primitives that appear in Kernel K-means. Our 1.5D algorithm consistently achieves the highest performance, enabling Kernel K-means to scale to data one to two orders of magnitude larger than previously practical. On 256 GPUs, it achieves a geometric mean weak scaling efficiency of $79.7\%$ and a geometric mean strong scaling speedup of $4.2\times$. Compared to our 1D algorithm, the 1.5D approach achieves up to a $3.6\times$ speedup on 256 GPUs and reduces clustering time from over an hour to under two seconds relative to a single-GPU sliding window implementation. Our results show that distributed algorithms designed with application-specific linear algebraic formulations can achieve substantial performance improvement.

Figures (6)

• Figure 1: 1.5D SpMM algorithm on $P=4$ processes. $\mathbf{V}\xspace$ is partitioned 1D columnwise and $\mathbf{K}\xspace$ in 2D. (1) The nonzeros of each $\mathbf{V}\xspace$ partition are replicated along the corresponding process row. (2) Each process performs a local SpMM with its $\mathbf{V}\xspace$ replicas and local $\mathbf{K}\xspace$ partition, producing partial sums of $\mathbf{E}\xspace^T$. (3) The partial sums are split along columns and reduced across process columns, resulting in a 1D columnwise partitioning of $\mathbf{E}\xspace^T$.
• Figure 2: The weak scaling evaluation on three datasets and $k\in\{16,64\}$. Results for the HIGGS dataset and $k=32$ are omitted for clarity.
• Figure 3: Weak scaling runtime breakdown for MNIST8m and HIGGS for $k=64$.
• Figure 4: Strong scaling evaluated on three datasets with $k \in \{16, 64\}$. Results for the HIGGS dataset and $k=32$ are omitted for clarity.
• Figure 5: The strong scaling runtime breakdown on MNIST8m and KDD for $k=64$.