多GPU并行运行大规模UMAP：分钟级完成、精度无损

均匀流形逼近与投影（UMAP）是一种广泛应用于可视化和特征提取的降维技术，涵盖探索性数据分析、主题建模和单细胞分析等多种应用场景。这些工作流程通常具有迭代性和探索性，需要用户在分析数据或调整参数时反复运行UMAP。随着数据集规模不断扩大，每次UMAP运算的成本大幅上升，使得交互式探索和迭代分析愈发困难。

UMAP算法的关键步骤是在数据集上构建全邻域kNN图，即为数据集中的每个向量找到k个最近邻。当数据集规模扩展至数千万乃至数亿个向量时，全邻域图的构建代价会急剧攀升。

此前发布的文章《借助NVIDIA cuML在GPU上实现更快、更可扩展的UMAP》介绍了如何通过核外计算方式扩展UMAP，使其能够处理单个GPU内存无法容纳的超大数据集。但训练阶段仍受限于单GPU，仅transform()步骤能够使用多个GPU。

本文介绍的是NVIDIA cuML和NVIDIA cuVS 25.06版本中发布的新特性——通过将耗费资源的全邻域图构建步骤分发到多个GPU上来突破这一限制，从而显著提升可扩展性并降低整体训练时间。

本文还将介绍如何利用cuML多GPU UMAP加速数千万乃至数亿个向量规模的数据集处理工作流程。在实践中，这一方法在超大规模数据集上能带来显著的端到端加速效果，可将数百GB数据的UMAP处理时间从数小时乃至数天压缩至数分钟。

实现核外扩展的核心思路

实现UMAP核外扩展的核心思路，是在不要求整个数据集同时加载进GPU内存的情况下完成全邻域kNN图的构建，这一点已在前文中有所介绍。具体做法是将数据集划分为均衡的簇，并在相邻簇之间进行向量重叠，以保留跨簇边界的最近邻关系。

每个簇的局部kNN图独立计算，再将这些局部图合并为一张全局全邻域图。这使UMAP得以在超出GPU内存容量的数据集上运行，同时保持嵌入质量。

自然延伸至多GPU场景

这一技术本身已将问题分解为相互独立的计算单元，因此可以自然地扩展至多GPU场景。每个簇可独立处理，局部kNN图的计算无需访问完整数据集，也无需与其他簇协调。因此，各簇可分发至不同GPU，每个GPU从CPU内存中独立获取所分配簇的数据。

每个GPU独立计算局部全邻域kNN图，并将结果合并至全局kNN图。这种独立计算局部kNN图的方式，避免了分布式全邻域kNN计算中通常存在的昂贵全量通信开销，在超大规模数据集上可带来显著的端到端性能提升。

关键超参数说明

在使用该特性之前，需要了解两个提供空间、时间与质量权衡的超参数。

cuML多GPU UMAP沿用单GPU UMAP实现的相同步骤，新增以下两个超参数：

knn_n_clusters：数据将被划分的簇数量。

knn_overlap_factor：每个数据点被分配到的最近邻簇的总数量。

knn_n_clusters所划分的簇大致均衡，即向量在各GPU上近似均匀分布以供处理。增大该值会减少每个簇中分配的向量数，从而降低每个GPU内存需要容纳的数据量。更多关于均衡k-means实现的详情，请参阅论文《GPU上的超大规模核外UMAP》。

knn_overlap_factor用于增大跨簇的向量重叠，从而保留更多跨簇边界的真实最近邻。增大该参数通常可以提升全邻域kNN图的质量，进而改善最终UMAP嵌入的质量。质量提升的代价是计算时间和内存占用的增加，因为每个簇需要处理更多向量。

knn_overlap_factor与knn_n_clusters共同提供了空间、时间与质量之间可控的权衡机制。

分布式全邻域图构建功能直接通过cuVS全邻域API暴露这些参数。cuML UMAP在内部使用这些参数，但对于需要独立构建全邻域图的应用程序，也可直接调用该API。

使用多GPU的cuVS全邻域API Python示例（示例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

# 使用系统上的所有GPU

res = MultiGpuResources()

indices, distances = all_neighbors.build(

data,

k,

params,

distances=cupy.empty((n_rows, k))

resources=res

其中n_clusters和overlap_factor对应cuML UMAP暴露的knn_n_clusters和knn_overlap_factor参数，cuML UMAP在构建图时会将这些参数转发至cuVS全邻域API。

参数调优建议

高质量嵌入的起点推荐使用knn_overlap_factor=2。实践中的建议如下：

较小规模的数据集适合以小幅递增方式调整knn_overlap_factor（2→3→4）。

knn_n_clusters较大（>100）的大规模数据集可能更适合以较大幅度递增（2→4→6）。

调整knn_overlap_factor时需谨慎，虽然更高的值会改善kNN召回率，但也会显著增加每个簇的计算时间。

若要在提升嵌入质量的同时控制内存开销，可在增大knn_overlap_factor的同时相应增大knn_n_clusters，从而使内存用量基本保持不变。

例如，将重叠因子翻倍会使每个簇的预期向量数翻倍，因此同时将簇数量翻倍即可维持每个簇的向量数不变。建议选取足够的knn_n_clusters，使每个局部图及其簇内向量能够较为轻松地放入GPU内存中。

内存估算方法

每个GPU所需的内存取决于数据集及待构建的局部kNN图。对于含N个D维向量的数据集，主要内存占用如下：

数据内存：M_data = N × (knn_overlap_factor / knn_n_clusters) × D × sizeof(data)

图内存：M_graph = N × k × (knn_overlap_factor / knn_n_clusters) × (sizeof(index) + sizeof(distance))

实际使用中，图构建过程中还需要额外的工作空间等开销，因此实际峰值内存可能高于上述估算值。建议选取knn_overlap_factor和knn_n_clusters，使预估内存占用明显低于GPU可用内存。

以使用80GB内存的GPU、N=1亿、D=1024的float32向量（共409GB）数据集为例，可选knn_overlap_factor=2、knn_n_clusters=24，此时每个簇的数据大小M_data约为34GB。使用int64索引类型（8字节）、float32距离类型（4字节）且k=15时，M_graph约为1.5GB，合计约35.5GB，为簇不均衡及kNN图构建过程中的其他开销留有足够余量。

内存占用与knn_overlap_factor成正比（决定每个向量的复制份数），与knn_n_clusters成反比。

cuML多GPU UMAP使用示例

在NVIDIA cuML中使用多GPU UMAP，只需在原有cuML UMAP配置的基础上增加少量参数。

以下Python示例（示例2）中，上方代码将使用所有可用GPU构建UMAP嵌入，下方代码仅使用ID为0、4、5的GPU：

from cuml.manifold import UMAP

# 使用系统上的所有GPU

umap = UMAP(

build_kwds={

"knn_n_clusters": 32,

"knn_overlap_factor": 2,

},

device_ids="all",

embedding = umap.fit_transform(data)

# 使用部分GPU

umap = UMAP(

build_kwds={

"knn_n_clusters": 32,

"knn_overlap_factor": 2,

},

device_ids=[0, 4, 5],

embedding = umap.fit_transform(data)

其中device_ids参数控制参与计算的GPU，knn_n_clusters和knn_overlap_factor控制全邻域kNN图的构建方式。

通常，增加GPU数量可通过将全邻域构建分散到多个设备上来缩短运行时间。

示例3展示了如何在cuML UMAP中使用示例1中由cuVS生成的预计算全邻域图：

from cuml.manifold.umap import UMAP as cuUMAP

# 使用示例1中由cuVS全邻域API计算的indices和distances

gpu_umap = cuUMAP(

precomputed_knn=(indices, distances),

gpu_embedding = gpu_umap.fit_transform(data)

这展示了通过precomputed_knn参数向UMAP提供预计算全邻域图的便捷性，CPU版UMAP同样支持该参数。

嵌入质量对比

图1对比了在106M×2048维的MIRACL数据集上，使用CPU参考实现（配合预计算GPU全邻域图）与cuML原生GPU UMAP实现生成的嵌入结果。由于在此规模下CPU计算全邻域图已不可行，采用了预计算方式。

值得注意的是，UMAP对缩放、平移和旋转具有不变性，因此尽管图像外观略有差异，这些嵌入结果在功能上是等价的。生成的嵌入展示出相近的全局结构，表明GPU全邻域图在超大规模下仍能保留产生高质量可视化所需的邻域关系。

性能基准测试

我们在超过单GPU内存容量的数据集（Wiki和MIRACL）上评估了多GPU UMAP的性能影响。所有基准测试均在搭载8块NVIDIA H100 GPU的NVIDIA DGX系统上执行，CPU为Intel Xeon 8480CL 224核处理器，内存为2TiB。

可信度分数（Trustworthiness Score）是衡量嵌入质量的常用指标，取值范围为0到1（越高越好），用于衡量UMAP低维嵌入空间中局部邻域结构相较于原始高维空间的保留程度。

图2a（上）展示了广泛使用的CPU参考实现在Wiki和MIRACL数据集逐渐增大的子集上的运行时间扩展行为。随着数据集规模增大，内存和计算成本导致运行时间急剧上升。在完整规模下，即使在2TiB内存的系统上，CPU实现也因内存耗尽而无法完成。

为估算此规模下的运行时间，通过对较小子集的实测扩展趋势进行外推来预测完整规模的CPU运行时间。详细方法见论文《GPU上的超大规模核外UMAP》。

图2b（下）将预估的CPU运行时间与在8块NVIDIA H100 GPU上运行的cuML多GPU UMAP实际运行时间进行对比。在1.06亿向量的MIRACL数据集上，cuML相较于预估CPU运行时间实现了最高74倍的端到端加速。

更重要的是，这使UMAP在以往无法处理的规模上变得切实可行——对1.06亿向量数据集的端到端处理仅需8分钟即可完成。

多GPU强扩展性

图3展示了cuML多GPU UMAP在Wiki和MIRACL数据集上，GPU数量从1块增加到8块H100 GPU时的扩展行为。在性能提升的同时，不同GPU配置下的嵌入质量保持相当。

关于多GPU全邻域UMAP实现的更多细节及更严格的评估，请参阅论文《GPU上的超大规模核外UMAP》。

总结

NVIDIA cuVS现已提供多GPU全邻域图构建能力，为NVIDIA cuML中的UMAP训练解锁了前所未有的规模与性能。这些能力使主题建模、单细胞分析等大规模可视化和嵌入工作流程更加实用，实现了更快速的迭代与探索。

如需开始使用cuML和cuVS中的多GPU UMAP，请参阅RAPIDS安装指南。有关NVIDIA cuVS中全邻域API的更多信息，请参阅cuVS API指南。

致谢：感谢Manas Singh和Mike Grauer对本文撰写与审阅的宝贵贡献，以及Leland McInnes创造并开源UMAP算法并持续支持GPU加速工作所做的贡献。

Q&A

Q1：cuML多GPU UMAP的knn_overlap_factor参数该怎么设置？

A：推荐从knn_overlap_factor=2开始。中等规模数据集可小幅递增（2→3→4）；对于knn_n_clusters超过100的大规模数据集，可以较大幅度递增（2→4→6）。提高该参数能改善kNN召回率和嵌入质量，但会增加计算时间和内存占用。若要在提升质量的同时控制内存，可同步增大knn_n_clusters来平衡内存使用。

Q2：cuML多GPU UMAP在MIRACL数据集上的性能表现如何？

A：在含1.06亿个向量的MIRACL数据集上，使用8块NVIDIA H100 GPU的cuML多GPU UMAP相较于CPU参考实现实现了最高74倍的端到端加速，完整数据集的端到端处理仅需约8分钟。而CPU参考实现在同等规模下因内存耗尽（即使在2TiB内存系统上）无法完成，是通过对小规模子集的扩展趋势外推估算的。

Q3：cuML多GPU UMAP如何在代码中指定使用哪些GPU？

A：在创建UMAP对象时，通过device_ids参数指定。设置device_ids="all"可使用系统上所有可用GPU；也可传入具体的GPU ID列表，例如device_ids=[0, 4, 5]，只使用指定的GPU参与计算。同时通过build_kwds字典传入knn_n_clusters和knn_overlap_factor参数来控制全邻域kNN图的构建方式。