统一流形近似与投影 (UMAP) 是一种广泛用于可视化和特征提取的降维技术。其应用范围涵盖探索性数据分析、主题建模和单细胞分析。其中许多工作流程是迭代和探索性的，需要在用户分析数据或调整参数时重复运行 UMAP。随着数据集的增长，每个 UMAP 的运行成本都会大幅增加，这使得交互式探索和迭代分析变得越来越困难。
UMAP 算法的关键步骤是在数据集上构建全邻 kNN 图，该图可为数据集中的每个向量查找 k – 最近邻点。随着数据集扩展到数千万或数亿个向量，全邻图构建的成本变得越来越高。
上一篇文章 (借助 NVIDIA cuML 在 GPU 上实现更快、更可扩展的 UMAP) 展示了如何使用核外方法扩展 UMAP，从而能够在以前对于单个 GPU 来说过大的数据集上拟合 UMAP。但是，训练阶段仍然局限于单个 GPU，并且只有 transform() 步骤能够使用多个 GPU。
本文将介绍 NVIDIA cuML 和 NVIDIA cuVS 25.06 中发布的一项功能，该功能通过在多个 GPU 上分配昂贵的多邻图构建步骤来消除这一限制。该功能可显著改善规模并缩短整体训练运行时间。
我们还介绍了如何使用 cuML 多 GPU UMAP 来加速具有数千万到数亿个向量的数据集的工作流。在实践中，这种方法在大规模数据集上实现了显著的端到端加速，正如本文所展示的那样。这使 UMAP 能够在几分钟内运行数百 GB 的工作负载，而不是几小时甚至几天。
NVIDIA cuML UMAP 如何跨多个 GPU 进行扩展？
启用核外扩展 UMAP 方法的关键理念是构建全近邻 kNN 图，而无需将整个数据集一次性放入 GPU 显存，如上一篇文章中所述。该方法将数据集划分为平衡集群，并将向量重叠到附近集群，以保留集群边界上的最近邻关系，从而实现这一目标。
针对每个集群分别计算本地 kNN 图，并将这些本地图合并到单个全局全邻图中。这使得 UMAP 能够在以前无法容纳 GPU 显存的规模下运行，同时仍然保持嵌入质量。
这种技术已经将问题分解为独立的工作单元，因此它自然而然地扩展到多 GPU 设置。每个集群都可以单独处理，这意味着计算本地 kNN 图无需访问完整数据集或与其他集群进行协调。因此，集群可以分布在 GPU 上，每个 GPU 从 CPU 内存中为其分配的集群单独收集数据。
然后，每个 GPU 计算本地全邻 kNN 图，并将每个图与全局 kNN 图合并。独立计算这些本地 kNN 图可避免昂贵的 all-to-all 通信，而这种通信通常会限制分布式全邻 kNN 工作负载的规模。这种方法可以对超大型数据集产生重大的端到端性能提升。
如何为多个 GPU 配置 UMAP
本节提供新功能的复制粘贴示例。首先，了解两个可权衡空间、时间和质量的超参数非常重要。
cuML 多 GPU UMAP 的步骤与单 GPU UMAP 实现相同，具有以下两个超参数：
knn_n_clusters：数据将划分到的集群数量knn_overlap_factor：将为每个数据点分配的最接近集群总数。
knn_n_clusters 是近似平衡的，这意味着向量在可用的 GPU 上大约均匀分布以进行处理。增加此值会减少分配给每个集群的点数量，从而减少需要适合每个 GPU 显存的数据量。如需详细了解 balanced k-means 实现，请参阅我们的论文《基于 GPU 的 Massive-Scale Out-Of-Core UMAP》。
如前所述，knn_overlap_factor 增加了跨集群的点重叠，跨集群边界保留了更多真最近邻点。增加此参数通常会提高全邻 kNN 图的质量，从而提高最终 UMAP 嵌入的质量。由于每个集群需要处理更多向量，因此需要权衡增加计算时间和内存使用量，从而提高质量。
knn_overlap_factor 和 knn_n_clusters 共同实现了空间、时间和质量之间可控的权衡。
分布式全邻图构建将这些参数作为 cuVS 全邻 API 直接公开。虽然 cuML UMAP 在内部使用它们，但它也可直接用于需要独立的多邻图构建的应用程序。
使用多个 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
)
# Using all GPUs on the system
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 参数，这些参数会在图形构建期间将其转发到此 cuVS all-neighbors API。
实际配置注意事项
要实现高质量嵌入，一个很好的起点是 knn_overlap_factor=2。在实践中：
- 较小的
knn_overlap_factor(2->3->4) 增量适用于中等规模。 - 具有大型
knn_n_clusters( > 100) 的较大数据集可能会受益于knn_overlap_factor(2->4->6) 的较大增量。
谨慎地增加 knn_overlap_factor，因为虽然提高 knn_overlap_factor 的值会改善 kNN 的召回率，但也会显著增加每个集群的计算时间。
要在管理内存权衡的同时提高嵌入质量，请提高 knn_overlap_factor，同时提高 knn_n_clusters，以保持内存使用量相当稳定。
例如，将重叠系数翻倍将使每个集群的预期点数翻倍，因此将集群数翻倍将使每个集群的点保持不变。建议使用足够的 knn_n_clusters，以便该集群中的每个局部图形和向量都能轻松地适合 GPU 显存。
GPU 显存要求
每个 GPU 所需的显存取决于数据集和正在构建的本地 kNN 图。对于包含 \(N\) 向量和维度 \(D\) 的数据集，主要内存组件包括：
\(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)\)
实际上，在图形构建期间，还需要额外的内存来处理开销，例如临时工作空间分配。因此，实际峰值显存使用量可能高于此处的估算值。因此，我们建议选择低于可用 GPU 显存的 knn_overlap_factor 和 knn_n_clusters。
例如，当使用显存为 80 GB、数据集大小为 \(N\) = 100M 且 \(D\) = 1024 float32 (sizeof ( data) = 4) 向量 (409 GB) 的 GPU 时，您可以选择 knn_overlap_factor=2 和 knn_n_clusters=24。这使得每个集群 \(M_{\mathit{data}}\) 的数据大小约为 34 GB。使用索引类型 int64 (sizeof ( index) = 8) 和距离类型 float32 (sizeof ( distance) = 4) 以及 k = 15 时，\(M_{\mathit{graph}}\) 约为 1.5 GB，加起来高达 35.5 GB。这为 K – means 拆分以及 kNN 图构建期间的其他用度带来的适度集群不平衡 (因为集群近似平衡) 留出了足够的空间。
内存使用量随 knn_overlap_factor 成比例扩展，因为它决定了每个向量的重复次数。相比之下，使用 knn_n_clusters 时，显存使用量会呈负增长。
如何在 NVIDIA cuML 中使用多 GPU UMAP
在 NVIDIA cuML 中使用多 GPU UMAP 只需要比之前使用 cuML UMAP 更多的配置选项。
在以下使用多 GPU UMAP 的 Python 示例 (示例 2) 中，最上面的示例将使用所有可用的 GPU 构建 UMAP 嵌入，而下面的示例将仅使用 ID 为 0、4 和 5 的 GPU：
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)
此处，device_ids 参数控制参与工作负载的 GPU 数量，而 knn_n_clusters 和 knn_overlap_factor 控制如何构建全邻 kNN 图。
通常，增加 GPU 数量会在设备之间分配全邻结构，从而缩短运行时间。后续各节将展示大规模数据集的可视化结果，以及展示添加 GPU 后强扩展效果的基准测试。
最后，以下示例 3 展示了在示例 1 中由 cuVS 生成的预先计算的全邻域图上运行 cuML UMAP：
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)
示例 1 展示了可以在 UMAP 之外计算全邻图，而示例 3 则展示了通过 precomputed_knn argument 向 UMAP 提供预先计算的全邻图是多么容易。CPU UMAP 也接受此参数。
可视化大规模嵌入
图 1 比较了使用 CPU 参考实现在 106M x 2048 MIRACL 数据集上生成的嵌入与预先计算的 GPU 全邻图和 cuML 原生 GPU UMAP 实现。我们预先计算了全邻图，因为在 CPU 上以这种规模进行计算是不可行的。
请注意，UMAP 在缩放、平移和旋转方面是不变的。这意味着，即使图像看起来略有不同，这些嵌入在功能上也是相同的。
生成的嵌入显示了类似的全局结构，这表明 GPU 全邻图保留了生成高质量可视化所需的邻域关系，即使是大规模的。
多 GPU UMAP 在超过单 GPU 显存的数据集上的性能如何？
我们评估了使用多 GPU UMAP 对超过单 GPU 显存 (Wiki 和 MIRACL) 的数据集的性能影响。所有基准测试均在 NVIDIA DGX 系统上执行，该系统配备 8 个 NVIDIA H100 GPU 和 1 个 Intel Xeon 8480CL 224 核 CPU，RAM 为 2TiB。
可信度分数是衡量嵌入质量的常用指标，返回 0 到 1 之间的数字 (越高越好) 。它用于测量与原始高维空间相比，在低维 UMAP 嵌入式空间中，局部邻域结构的保留程度。
大规模端到端加速
图 2a (上) 显示了广泛使用的 CPU 参考实现在 Wiki 和 MIRACL 数据集的更大子样本上的运行时扩展行为。随着数据集大小的增加，内存和计算成本会导致运行时迅速增长。在完整范围内，由于内存消耗过多，CPU 实施无法完成，即使在 RAM 为 2 TiB 的系统上也是如此。
为了估算这些规模的运行时间，我们通过从较小的子样本推断测量的扩展趋势来预测完整的 CPU 运行时间。如需了解此方法的完整详细信息，请参阅我们的论文《 GPU 上的 Massive-Scale Out-Of-Core UMAP》。
图 2b (下) 将预测的 CPU 运行时与在 8 个 NVIDIA H100 GPU 上运行的 cuML 多 GPU UMAP 的实际运行时进行了比较。在 106M 向量 MIRACL 数据集上，cuML 实现了比预计 CPU 运行时高达 74 倍的端到端加速。
最重要的是，这使得 UMAP 能够在以前难以实现的规模上发挥作用，只需 8 分钟即可完成 1.06 亿向量数据集的端到端处理。
多 GPU 扩展
图 3 显示了当 GPU 数量从 1 个增加到 8 个 H100 GPU 时，Wiki 和 MIRACL 数据集上 cuML 多 GPU UMAP 的扩展行为。实现性能提升的同时，还能在不同 GPU 配置下保持可美的嵌入质量。
有关多 GPU 多邻 UMAP 实现的更多详细信息 (包括更严格的评估) ，请参阅我们的论文：基于 GPU 的 Massive-Scale Out-Of-Core 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 加速的持续支持。
