UMAP（Uniform Manifold Approximation and Projection）是数据科学中广泛使用的降维技术，常用于数据可视化、特征提取、主题建模与单细胞分析等场景。这类工作流往往需要反复迭代调参，随着数据集不断增大，单次 UMAP 运行的成本急剧上升，交互式探索变得越来越困难。而 UMAP 算法中最耗时的环节，是对全量数据构建"全邻居 kNN 图"——为数据集中每一个向量寻找 k 个最近邻，这一步骤在数据达到千万、上亿级向量时开销尤为惊人。
此前 NVIDIA 曾推出基于 cuML 的核外（out-of-core）UMAP 方案，让超出单卡显存的数据也能运行，但训练阶段仍局限于单张 GPU，只有 transform() 步骤能使用多卡。此次 NVIDIA 在 cuML 与 cuVS 25.06 中发布的新特性彻底移除这一限制：将最昂贵的全邻居图构建步骤分布到多张 GPU 上执行，显著提升可处理的数据规模并大幅缩短训练时间。
新方案的核心思路是把数据集划分成多个大致均衡的簇，并让邻近簇之间保留一定重叠，从而在簇边界处维持真实的近邻关系。每张 GPU 独立为分配给自己的簇计算局部 kNN 图，再与全局图合并，整个过程无需昂贵的 all-to-all 通信，因此可以天然扩展到多卡环境。用户只需通过 knn_n_clusters（划分簇数）与 knn_overlap_factor（跨簇重叠因子）两个超参数，即可在显存占用、运行时间与嵌入质量之间取得可控的平衡；device_ids 参数则用来指定参与计算的 GPU。
实测数据验证了该方案在大规模数据集上的价值。在配备 8 张 NVIDIA H100 GPU、2TiB 内存的 DGX 平台上，使用 Wiki 与 MIRACL 数据集进行基准测试：其中 MIRACL 数据集包含 1.06 亿条 2048 维向量。CPU 参考实现在全量数据上因内存耗尽而无法完成，按小样本外推估算其耗时后对比，cuML 多 GPU UMAP 实现了最高 74 倍的端到端加速，整个 1.06 亿向量数据集的处理在约 8 分钟内完成。在从 1 卡扩展到 8 卡的过程中，运行时间持续下降，而嵌入质量（以 trustworthiness 得分衡量）保持一致。
对数据科学团队而言，这项能力让 UMAP 在以往难以触及的数据规模上变得实用，例如数百 GB 量级的数据集可以从数小时甚至数天缩短到几分钟完成。cuML 与 cuVS 均为开源项目，可通过 RAPIDS 安装指南快速上手；cuVS 的 all-neighbors API 也允许独立构建全邻居图，再通过 precomputed_knn 参数交给 UMAP 使用，为大规模可视化、主题建模与单细胞分析等工作流提供更快的迭代路径。
WeChat
Profile