We propose an efficient distributed out-of-memory implementation of the Non-negative Matrix Factorization (NMF) algorithm for heterogeneous high-performance-computing (HPC) systems. The proposed implementation is based on prior work on NMFk, which can perform automatic model selection and extract latent variables and patterns from data. In this work, we extend NMFk by adding support for dense and sparse matrix operation on multi-node, multi-GPU systems. The resulting algorithm is optimized for out-of-memory (OOM) problems where the memory required to factorize a given matrix is greater than the available GPU memory. Memory complexity is reduced by batching/tiling strategies, and sparse and dense matrix operations are significantly accelerated with GPU cores (or tensor cores when available). Input/Output (I/O) latency associated with batch copies between host and device is hidden using CUDA streams to overlap data transfers and compute asynchronously, and latency associated with collective communications (both intra-node and inter-node) is reduced using optimized NVIDIA Collective Communication Library NCCL based communicators. Benchmark results show significant improvement, from 32X to 76x speedup, with the new implementation using GPUs over the CPU-based NMFk. Good weak scaling was demonstrated on up to 4096 multi-GPU cluster nodes with approximately 25,000 GPUs when decomposing a dense 340 Terabyte-size matrix and an 11 Exabyte-size sparse matrix of density 10e-6.

翻译：我们提出了一种高效的分布式内存外非负矩阵分解（NMF）算法实现，适用于异构高性能计算（HPC）系统。该实现基于前期NMFk工作，能够自动进行模型选择并提取数据中的潜在变量和模式。本研究通过支持多节点、多GPU系统上的稠密与稀疏矩阵运算，扩展了NMFk算法。优化后的算法专门针对内存外（OOM）问题——即待分解矩阵所需内存超出可用GPU内存的情况。通过分块/平铺策略降低内存复杂度，并利用GPU核心（若支持则使用张量核心）显著加速稀疏与稠密矩阵运算。采用CUDA流技术隐藏主机与设备间批量拷贝的输入/输出（I/O）延迟，实现数据传输与异步计算的重叠；同时基于优化的NVIDIA集体通信库（NCCL）构建通信器，减少节点内与节点间集体通信的延迟。基准测试结果显示，新实现基于GPU相比CPU版NMFk实现了32倍至76倍的加速比。在约25,000个GPU的4096个多GPU集群节点上，对340TB稠密矩阵和密度为10e-6的11EB稀疏矩阵进行分解时，展现出良好的弱扩展性。