NVIDIA Accelerates Large-Scale UMAP with Multi-GPU Support
UMAP processing on 870 GB of data now takes just 8 minutes.
1What is UMAP?
Uniform Manifold Approximation and Projection (UMAP) is a dimensionality reduction technique used for data visualization and feature extraction. It has a wide range of applications from exploratory data analysis to single-cell analysis. Performance optimization is critical as UMAP often needs to be run repeatedly on large datasets.
2New Feature
NVIDIA cuML and cuVS 25.06 introduced multi-GPU support for the all-neighbors kNN graph construction in UMAP. This feature significantly reduces training times on datasets with tens to hundreds of millions of vectors. Tests on MIRACL and Wiki datasets showed 74x speedup compared to CPUs using eight NVIDIA H100 GPUs.
3Technical Approach
The multi-GPU implementation partitions datasets into balanced clusters. Local kNN graphs are computed independently per cluster and merged into a global graph. This method enables distributed computation without all-to-all communication while preserving embedding quality at massive scale.
4Memory Management
The out-of-core approach eliminates the need for the entire dataset to fit in GPU memory. Data is divided into balanced clusters with overlapping vectors between nearby clusters. This preserves nearest-neighbor relationships across cluster boundaries.
5Hyperparameters
The knn_n_clusters parameter ensures balanced distribution of vectors across GPUs. knn_overlap_factor increases preservation of true nearest neighbors across cluster boundaries. Together, these parameters provide controllable trade-offs between space, time, and quality.
6Performance Gains
With multi-GPU support, UMAP can now process hundred-gigabyte workloads in minutes instead of hours or days. A 870 GB dataset can be processed in just 8 minutes. This represents a paradigm shift for large-scale data analysis workflows.
7Application Areas
This advancement is particularly transformative for iterative workflows like single-cell analysis, topic modeling, and exploratory data analysis. Researchers can now tune parameters and see results much faster on large datasets, accelerating scientific discovery.
8Distributed Computing Advantage
Each cluster can be processed independently, enabling natural scaling to multiple GPUs. GPUs independently gather data for assigned clusters from CPU memory. Computing local kNN graphs separately eliminates communication bottlenecks common in distributed approaches.
9Quality Assurance
The method maintains high trustworthiness scores even on massive datasets. Tests confirm no degradation in embedding quality compared to single-GPU implementation. This makes the solution reliable for scientific research applications.
10Ease of Use
cuML multi-GPU UMAP follows the same steps as single-GPU UMAP. Users can migrate existing code with minimal changes. The cuVS all-neighbors API is also exposed for standalone kNN graph construction needs.
11Industry Impact
This technology marks a significant milestone in big data analytics and AI. NVIDIA's solution removes scalability barriers for researchers and data scientists, paving the way for new discoveries. It stands as another success story for GPU-accelerated distributed computing.