NVIDIA Enables Multi-GPU Training for Massive-Scale UMAP via cuML and cuVS

NVIDIA Enables Multi-GPU Training for Massive-Scale UMAP via cuML and cuVS NVIDIA has introduced multi-GPU acceleration for the computationally intensive all-neighbors kNN graph construction step in cuML and cuVS 25.06, enabling UMAP training on datasets with hundreds of millions of vectors in minutes while preserving embedding quality. cuML and cuVS 25.06 remove the single-GPU bottleneck in UMAP training by distributing all-neighbors kNN graph construction across multiple GPUs. The distributed approach partitions data into balanced clusters with overlapping vectors, avoiding expensive all-to-all GPU communication during local kNN graph computations. Two key hyperparameters, knn_n_clusters and knn_overlap_factor, allow users to tune memory consumption, runtime, and neighborhood recall. Benchmarking on an NVIDIA DGX system with eight H100 GPUs showed end-to-end processing of the 106M-vector MIRACL dataset in 8 minutes, achieving an estimated 74x speedup over CPU projections. Scaling All-Neighbors Graph Construction Across GPUs UMAP relies on building an all-neighbors kNN graph across the entire dataset, a step that becomes computationally prohibitive on tens or hundreds of millions of vectors. While previous cuML releases supported out-of-core scaling and multi-GPU inference via transform(), UMAP training was historically limited to a single GPU. The new multi-GPU feature partitions data into balanced clusters and overlaps vectors across nearby boundaries to preserve nearest-neighbor relationships. Each GPU independently fetches its assigned cluster data from CPU memory, computes a local kNN graph, and merges it into the global graph, completely avoiding costly all-to-all communication between GPUs. Configuration and Hyperparameter Trade-offs The distributed algorithm introduces two primary hyperparameters exposed in both the cuVS all_neighbors API and cuML UMAP: knn_n_clusters and knn_overlap_factor. The cluster count controls data partitioning to ensure cluster sizes fit into individual GPU memory, while the overlap factor assigns vectors to multiple neighboring clusters to preserve edge fidelity and graph recall. Increasing the overlap factor improves embedding quality and trustworthiness but demands more memory and compute. To offset memory increases from higher overlap factors, users can proportionally scale up the number of clusters. Memory Estimation and Practical Usage GPU memory consumption scales inversely with knn_n_clusters and linearly with knn_overlap_factor, split primarily between vector data and the neighbor index and distance structures. Additional temporary overhead occurs during construction, requiring headroom above calculated minimums. In cuML, users invoke multi-GPU execution through the device_ids parameter (set to 'all' or a list of specific IDs) alongside the build_kwds dictionary, or by passing a precomputed cuVS all-neighbors graph directly via precomputed_knn. Massive-Scale Benchmarks and Scaling Performance Evaluations conducted on an NVIDIA DGX system equipped with eight H100 GPUs and 2 TiB of system RAM demonstrated strong scaling on massive datasets such as Wiki and the 106-million-vector, 2048-dimensional MIRACL dataset. Whereas the reference CPU implementation failed on full datasets due to memory exhaustion, cuML multi-GPU UMAP processed the full MIRACL dataset in 8 minutes, representing an estimated 74x speedup compared to extrapolated CPU runtimes while maintaining comparable trustworthiness and visualization quality.

NVIDIA cuML与cuVS支持多GPU加速UMAP算法，数亿级向量降维缩减至数分钟 NVIDIA在cuML与cuVS 25.06中引入了多GPU分布式全邻居kNN图构建功能，消除了此前UMAP训练阶段仅能单卡运行的限制，使得数百GB、上亿级别向量的降维任务可在数分钟内完成且保持高精度。 cuML与cuVS通过在多GPU间分布构建本地kNN图并合并为全局图，避免了开销昂贵的多机全对全通信。 系统提供knn_n_clusters与knn_overlap_factor两个超参数，支持在计算时间、显存占用与嵌入质量之间进行权衡。 在包含1.06亿向量的MIRACL数据集上，8块H100 GPU使端到端运行时间缩短至8分钟，相较CPU投影性能提速达74倍。 多GPU分布式全邻居图构建原理 UMAP最耗时的步骤是在高维向量集合上构建全邻居kNN图。新方案基于核外（out-of-core）策略，将数据集划分为平衡的簇，并在临近簇之间重叠向量以保留边界邻居关系。由于各簇的本地kNN图计算彼此独立且无需访问全量数据，系统可将簇分配到不同GPU上并发处理并合并至全局图，消除了分布式计算中昂贵的全对全通信。 核心超参数配置与显存考量 用户可通过knn_n_clusters（聚类簇数）和knn_overlap_factor（重叠因子）调节性能。增加簇数可减少单GPU处理的点数并降低显存压力；增加重叠因子则能提高最近邻召回率及嵌入质量，但会成比例增加显存和计算时间。实践中推荐以knn_overlap_factor=2为起点，若提高重叠因子，应同步增加簇数以保持单卡显存占用稳定并留出临时开销余量。 大规模基准测试与加速表现 在配备8块NVIDIA H100 GPU的DGX系统上进行的Wiki和MIRACL数据集测试显示，CPU参考实现在全规模数据下因内存溢出无法完成。而cuML多GPU实现在1.06亿条2048维向量的MIRACL数据集上仅用时8分钟完成端到端处理，比CPU推算耗时快多达74倍，且随GPU数量增加展现出强扩展性，同时保持了相当的全局结构与可信度指标。

Uniform Manifold Approximation and Projection (UMAP) is a dimensionality reduction technique widely used for visualization and feature extraction. Applications range across exploratory data analysis, topic modeling, and single-cell analysis. Many of these workflows are iterative and exploratory, requiring UMAP to be run repeatedly as users analyze their data or tune parameters. As datasets grow, the cost of each UMAP run increases substantially, making interactive exploration and iterative analysis increasingly difficult. A critical step in the UMAP algorithm is the construction of an all-neighbors kNN graph over the dataset, which finds the k-nearest neighbors for every vector in the dataset. All-neighbors graph construction becomes increasingly expensive as datasets scale to tens or hundreds of millions of vectors. The previous post, Even Faster and More Scalable UMAP on the GPU with NVIDIA cuML , demonstrates how to scale UMAP using an out-of-core approach, making it possible to fit UMAP on datasets that were previously too large for a single GPU. However, the training stage remained limited to a single GPU and only the transform() step was able to use multiple GPUs. This post introduces a feature released in NVIDIA cuML and NVIDIA cuVS 25.06 that removes this limitation by distributing the expensive all-neighbors graph construction step across multiple GPUs. The feature significantly improves the scale and reduces the overall training runtime. We also explain how to use cuML multi-GPU UMAP to accelerate workflows on datasets with tens to hundreds of millions of vectors. In practice, this approach delivers substantial end-to-end speedups on massive-scale datasets, as this post will show. This enables UMAP to run on workloads of several hundred gigabytes in minutes instead of hours or even days. How does NVIDIA cuML UMAP scale across multiple GPUs? The key idea for enabling the out-of-core approach for scaling UMAP is constructing the all-neighbors kNN graph without requiring the entire dataset to fit in GPU memory at once, as introduced in the previous post. The approach accomplishes this by partitioning the dataset into balanced clusters and overlapping the vectors across nearby clusters to preserve nearest-neighbor relationships across the cluster boundaries. Local kNN graphs are computed independently for each cluster, and these local graphs are merged into a single global all-neighbors graph. This makes it possible to run UMAP at scales that were previously too large to fit in GPU memory, while still preserving embedding quality. This technique already decomposes the problem into independent units of work, so it extends naturally to a multi-GPU setting. Each cluster can be processed in isolation, meaning the computation of the local kNN graphs does not require access to the full dataset or coordination with other clusters. As a result, the clusters can be distributed across GPUs, where each GPU independently gathers the data for its assigned clusters from CPU memory. Each GPU then computes local all-neighbors kNN graphs and merges each one with the global kNN graph. Computing these local kNN graphs independently avoids the need for expensive all-to-all communication that would typically limit scale in distributed all-neighbors kNN workloads. This method can produce major end-to-end performance gains on very large datasets. How to configure UMAP for multiple GPUs This section provides copy-and-paste examples of the new feature. First, it’s important to understand two hyperparameters that provide the space, time, and quality tradeoff. cuML multi-GPU UMAP follows the same steps as the single-GPU UMAP implementation, with the following two hyperparameters: knn_n_clusters : The number of clusters the data will be partitioned into knn_overlap_factor : The total number of closest clusters each data point will be assigned to. The knn_n_clusters are approximately balanced, meaning the vectors are distributed approximately evenly across available GPUs for processing. Increasing this value reduces the number of points assigned to each cluster, lowering the amount of data that needs to fit in the memory of each GPU. For more details on the balanced k-means implementation, see our paper, Massive-Scale Out-Of-Core UMAP on the GPU . As mentioned, knn_overlap_factor increases the overlap of points across clusters, preserving more of the true nearest neighbors across cluster boundaries. Increasing this parameter generally improves the quality of the all-neighbors kNN graph, which improves the quality of the final UMAP embeddings. The improved quality comes with the tradeoff of increasing computation time and memory usage, as more vectors need to be processed per cluster. knn_overlap_factor and knn_n_clusters together provide a controllable trade-off between space, time, and quality. The distributed all-neighbors graph construction directly exposes these parameters as the cuVS all-neighbors API. While cuML UMAP uses them internally, it is also available directly for applications that require standalone all-neighbors graph construction. Python example of the cuVS all-neighbors API using multiple GPUs (Example 1): from cuvs.neighbors import all_neighbors from cuvs.common import MultiGpuResources params = all_neighbors.AllNeighborsParams( algo="nn_descent", n_clusters=32, overlap_factor=2 ) # Using all GPUs on the system res = MultiGpuResources() indices, distances = all_neighbors.build( data, k, params, distances=cupy.empty((n_rows, k)) resources=res ) The n_clusters and overlap_factor correspond to the knn_n_clusters and knn_overlap_factor arguments exposed by cuML UMAP, which forwards them to this cuVS all-neighbors API during graph construction. Practical configuration considerations A good starting point for a quality embedding is knn_overlap_factor=2 . In practice: Smaller increments of knn_overlap_factor (2->3->4) work well for moderate scales. Larger datasets that have a large knn_n_clusters (> 100) may benefit from larger increments of knn_overlap_factor (2->4->6). Increase knn_overlap_factor carefully, because while higher values of knn_overlap_factor will improve kNN recall, they will also significantly increase the computation time for each cluster. To manage the memory trade-off while improving embedding quality, raise knn_overlap_factor while also increasing knn_n_clusters to keep the memory usage fairly constant. For example, doubling the overlap factor will double the expected number of points per cluster, so doubling the number of clusters will maintain the points per cluster. It is suggested to use enough knn_n_clusters such that each local graph and vectors in that cluster comfortably fit in GPU memory. GPU memory requirements The memory required on each GPU depends on the dataset and the local kNN graph being constructed. For a dataset with \(N\) vectors and dimension \(D\), the dominant memory components are: \(M_{\mathit{data}} = N \cdot \frac{\mathit{knn\_overlap\_factor}}{\mathit{knn\_n\_clusters}} \cdot D \cdot \mathit{sizeof}(\mathit{data})\) \(M_{\mathit{graph}} = N \cdot k \cdot \frac{\mathit{knn\_overlap\_factor}}{\mathit{knn\_n\_clusters}} \cdot \bigl(\mathit{sizeof}(\mathit{index}) + \mathit{sizeof}(\mathit{distance})\bigr)\) In practice, additional memory is also required for overhead like temporary workspace allocations during graph construction. Because of this, the actual peak memory usage can be higher than estimated here. We therefore recommend choosing knn_overlap_factor and knn_n_clusters comfortably below the available GPU memory. For example, when using GPUs with 80 GB memory and a dataset with \(N\)=100M, and \(D\)=1024 float32 (sizeof(data)=4) vectors (409 GB), you might choose knn_overlap_factor=2 and knn_n_clusters=24 . This makes the data size of each cluster \(M_{\mathit{data}}\) approximately 34 GB. Using index type int64 (sizeof(index) = 8) and distance type float32 (sizeof(distance)=4),  and k=15, \(M_{\mathit{graph}}\) is about 1.5 GB, adding up to 35.5 GB. This leaves sufficient headroom for modest cluster imbalance (because the clusters are approximately balanced) from k-means splits, and other overheads during the kNN graph construction. Memory usage scales proportionally with knn_overlap_factor because it decides the number of duplications per vector. In contrast, memory usage scales inversely with knn_n_clusters . How to use multi-GPU UMAP in NVIDIA cuML Using multi-GPU UMAP in NVIDIA cuML requires only a few additional configuration options beyond previous cuML UMAP usage. In the following Python example of using multi-GPU UMAP (Example 2), the top example will construct the UMAP embedding with all available GPUs, while the bottom example will use only GPUs with IDs 0, 4, and 5: from cuml.manifold import UMAP # Using all GPUs on the system umap = UMAP( build_kwds={ "knn_n_clusters": 32, "knn_overlap_factor": 2, }, device_ids="all", ) embedding = umap.fit_transform(data) # Using a subset of GPUs umap = UMAP( build_kwds={ "knn_n_clusters": 32, "knn_overlap_factor": 2, }, device_ids=[0, 4, 5], ) embedding = umap.fit_transform(data) Here, the device_ids argument controls which GPUs participate in the workload, while knn_n_clusters and knn_overlap_factor control how the all-neighbors kNN graph is constructed. In general, increasing the number of GPUs reduces runtime by distributing the all-neighbors construction across devices. The next sections show the visualization results for a massive-scale dataset, and benchmarks that demonstrate the strong scaling effects as GPUs are added. Finally, Example 3 below shows running cuML UMAP on the precomputed all-neighbors graph generated by cuVS in Example 1: from cuml.manifold.umap import UMAP as cuUMAP # Using indices, distances from Example 1 computed by cuVS all-neighbors gpu_umap = cuUMAP( precomputed_knn=(indices, distances), ) gpu_embedding = gpu_umap.fit_transform(data) While Example 1 shows that the all-neighbors graph can be computed outside of UMAP, Example 3 demonstrates how easy it is to provide a precomputed all-neighbors graph to UMAP through the precomputed_knn argument . The CPU UMAP also accepts this argument. Visualizing massive-scale embeddings Figure 1 compares embeddings generated on the 106M x 2048 MIRACL dataset using the CPU reference implementation with a precomputed GPU all-neighbors graph and cuML native GPU UMAP implementation. We precomputed the all-neighbors graph because it is computationally infeasible to compute on the CPU at this scale. Note that UMAP is invariant to scale, translations, and rotations. This means these embeddings are functionally identical, even though the images look slightly different. Figure 1. Comparison of embeddings of CPU UMAP using a GPU-precomputed all-neighbors graph, and cuML UMAP on the MIRACL dataset The resulting embeddings show comparable global structure, demonstrating that the GPU all-neighbors graph preserves the neighborhood relationships needed to produce high-quality visualizations even at a massive scale. What is the performance of multi-GPU UMAP on datasets that exceed single-GPU memory? We evaluated the performance impact of using multi-GPU UMAP on datasets that exceed single-GPU memory ( Wiki and MIRACL ). All benchmarks were executed on an NVIDIA DGX system with eight NVIDIA H100 GPUs and an Intel Xeon 8480CL 224-core CPU with 2TiB of RAM. The trustworthiness score is a popular measure of embedding quality, returning a number between 0 and 1 (higher is better). It measures the degree to which the local neighborhood structure is preserved in the low-dimensional UMAP embedded space as compared to the original high-dimensional space. End-to-end speedup at massive scale Figure 2a (top) shows the runtime scaling behavior of the widely used CPU reference implementation on progressively larger subsamples of the Wiki and MIRACL datasets. As the dataset size increases, the runtime grows rapidly due to the memory and computation cost. At full scale, the CPU implementation failed to complete because of excessive memory consumption, even on a system with 2 TiB of RAM. To estimate runtime at these scales, we projected the full-scale CPU runtime by extrapolating the measured scaling trend from smaller subsamples. For full details of this methodology, see our paper, Massive-Scale Out-Of-Core UMAP on the GPU . Figure 2b (bottom) compares the projected CPU runtime against the actual runtime of cuML multi-GPU UMAP running on eight NVIDIA H100 GPUs. On the 106M vector MIRACL dataset, cuML achieves up to 74x end-to-end speedup over the projected CPU runtime. Figure 2. Measured and projected CPU UMAP runtimes for the Wiki and MIRACL datasets for varying subsample sizes (top) and an end-to-end runtime comparison of projected CPU UMAP runtimes and cuML GPU UMAP (bottom) Most importantly, this makes UMAP practical at previously intractable scales, completing end-to-end processing of the 106M vector dataset in just 8 minutes. Multi-GPU scaling Figure 3 shows the scaling behavior of cuML multi-GPU UMAP on the Wiki and MIRACL datasets as the number of GPUs increases from one to eight H100 GPUs. The performance improvement is achieved while preserving comparable embedding quality across GPU configurations. Figure 3. Multi-GPU scaling of the new scalable all-neighbors algorithm in UMAP: performance continues to improve as more GPUs are added, without sacrificing trustworthiness For additional details about the multi-GPU all-neighbors UMAP implementation, including more rigorous evaluation, see our paper, Massive-Scale Out-Of-Core UMAP on the GPU . Get started with UMAP using multiple GPUs NVIDIA cuVS now provides multi-GPU all-neighbors graph construction, unlocking UMAP training in NVIDIA cuML at unprecedented scale and performance. These capabilities make large-scale visualization and embedding workflows such as topic modeling and single-cell analysis more practical, enabling even faster iteration and exploration. To get started with multi-GPU UMAP in cuML and cuVS, see the RAPIDS Installation Guide . For more information about the all-neighbors APIs in NVIDIA cuVS, see the cuVS API Guide . Acknowledgments We’d like to thank Manas Singh and Mike Grauer for their valuable contributions to drafting and reviewing this post. We’re also deeply grateful to Leland McInnes for creating and sharing the UMAP algorithm with the world, and for his continued support in our advancements of its GPU acceleration.