RAFT: 用于向量搜索和其他的加速函数和工具

[!IMPORTANT] RAFT中的向量搜索和聚类算法正在迁移到一个名为cuVS的新库，该库专门用于向量搜索。在这次迁移过程中，我们将继续支持RAFT中的向量搜索算法，但不会在RAPIDS 24.06（6月）发布后对其进行更新。我们计划在RAPIDS 24.08（8月）发布前完成迁移。

内容

1. 有用资源
2. 什么是RAFT？
3. 用例
4. RAFT适合我吗？
5. 入门
6. 安装RAFT
7. 代码库结构和内容
8. 贡献
9. 参考

有用资源

• RAFT参考文档: API文档。
• RAFT入门: 开始使用RAFT。
• 构建和安装RAFT: 安装和构建RAFT的说明。
• 示例笔记本: 示例jupyter笔记本
• RAPIDS社区: 获取帮助、贡献和协作。
• GitHub仓库: 下载RAFT源代码。
• 问题跟踪器: 报告问题或请求功能。

什么是RAFT？

RAFT 包含了机器学习和信息检索广泛使用的算法和原语。这些算法经过 CUDA 加速，形成构建高性能应用程序的基础模块。

通过采用基于原语的方法进行算法开发，RAFT

• 缩短了算法构建时间
• 通过最大化项目间的复用来降低维护负担，并且
• 集中核心可重用计算，使未来的优化能够惠及所有使用它们的算法。

虽然不是详尽的，以下的一般类别有助于总结 RAFT 中加速的功能

RAFT 是一个仅包含 C++ 头文件的模板库，包含可选的共享库，

1. 可以加速常见模板类型的编译时间，并提供
2. 主机可访问的 "运行时" API，无需 CUDA 编译器即可使用

除了是一个 C++ 库，RAFT 还提供了 2 个 Python 库

• pylibraft - RAFT 的主机可访问 "运行时" API 的轻量级 Python 包装器。
• raft-dask - 用于在 GPU 上使用 Dask 构建分布式算法的多节点多 GPU 通信基础设施。

用例

向量相似度搜索

RAFT 包含在 GPU 上近似最近邻搜索 (ANNS) 算法的最新实现，例如

• 暴力搜索。在不使用索引的情况下执行暴力最近邻搜索。
• IVF-Flat 和 IVF-PQ。使用倒排文件索引结构将内容映射到其位置。IVF-PQ 还使用产品量化来减少向量的内存使用。这些方法最初由 FAISS 库推广。
• CAGRA (Cuda Anns GRAph-based)。使用快速 ANNS 图构建和搜索实现，针对 GPU 进行优化。CAGRA 在大批量查询、单个查询和图构建时间方面优于最先进的 CPU 方法（即 HNSW）。

使用 RAFT ANNS 算法加速向量搜索的项目包括： Milvus、Redis 和 Faiss。

请参阅示例 Jupyter 笔记本 以开始使用 Python 中的 RAFT 进行向量搜索。

信息检索

RAFT 包含了一个可重用原语目录，用于组合需要快速邻域计算的算法，例如

1. 计算向量之间的距离和计算核语法矩阵
2. 执行球半径查询以构建 epsilon 邻域
3. 聚类点以划分空间以进行更小和更快的搜索
4. 从密集向量构建邻域 "连通性" 图

机器学习

RAFT 的原语被多个 RAPIDS 库使用，包括 cuML、cuGraph 和 cuOpt，用于构建涵盖广泛不同应用的端到端机器学习算法，包括

• 数据生成
• 模型评估
• 分类和回归
• 聚类
• 流形学习
• 降维。

RAFT 还被流行的协同过滤库 implicit 用于推荐系统。

RAFT适合我吗？

RAFT 包含用于加速应用程序和工作流程的低级原语。数据源提供商和应用开发者可能会发现特定工具（如 ANN 算法）非常有用。RAFT 不建议数据科学家直接用于发现和实验。有关数据科学工具，请参阅 RAPIDS 网站。

入门指南

RAPIDS 内存管理器（RMM）

RAFT 严重依赖 RMM，它简化了在全局范围内配置不同分配策略的负担。

多维数组

RAFT 的 API 接受 mdspan 多维数组视图来表示类似于 Numpy Python 库中 ndarray 的高维数据。RAFT 还包含相应的拥有 mdarray 结构，它简化了主机和设备（GPU）内存中多维数据的分配和管理。

mdarray 在 RMM 之上形成了一个便利层，并且可以使用许多不同的辅助函数在 RAFT 中构建

#include <raft/core/device_mdarray.hpp>

int n_rows = 10;
int n_cols = 10;

auto scalar = raft::make_device_scalar<float>(handle, 1.0);
auto vector = raft::make_device_vector<float>(handle, n_cols);
auto matrix = raft::make_device_matrix<float>(handle, n_rows, n_cols);

C++ 示例

RAFT 中的大多数原语都接受一个 raft::device_resources 对象来管理创建成本高昂的资源，例如 CUDA 流、流池以及其他 CUDA 库（如 cublas 和 cusolver）的句柄。

以下示例演示了创建 RAFT 句柄，并使用 device_matrix 和 device_vector 来分配内存、生成随机聚类以及计算成对欧几里得距离。

#include <raft/core/device_resources.hpp>
#include <raft/core/device_mdarray.hpp>
#include <raft/random/make_blobs.cuh>
#include <raft/distance/distance.cuh>

raft::device_resources handle;

int n_samples = 5000;
int n_features = 50;

auto input = raft::make_device_matrix<float, int>(handle, n_samples, n_features);
auto labels = raft::make_device_vector<int, int>(handle, n_samples);
auto output = raft::make_device_matrix<float, int>(handle, n_samples, n_samples);

raft::random::make_blobs(handle, input.view(), labels.view());

auto metric = raft::distance::DistanceType::L2SqrtExpanded;
raft::distance::pairwise_distance(handle, input.view(), input.view(), output.view(), metric);

还可以创建 raft::device_mdspan 视图，使用原始指针和形状信息调用相同的 API。

#include <raft/core/device_resources.hpp>
#include <raft/core/device_mdspan.hpp>
#include <raft/random/make_blobs.cuh>
#include <raft/distance/distance.cuh>

raft::device_resources handle;

int n_samples = 5000;
int n_features = 50;

float *input;
int *labels;
float *output;

...
// Allocate input, labels, and output pointers
...

auto input_view = raft::make_device_matrix_view(input, n_samples, n_features);
auto labels_view = raft::make_device_vector_view(labels, n_samples);
auto output_view = raft::make_device_matrix_view(output, n_samples, n_samples);

raft::random::make_blobs(handle, input_view, labels_view);

auto metric = raft::distance::DistanceType::L2SqrtExpanded;
raft::distance::pairwise_distance(handle, input_view, input_view, output_view, metric);

Python 示例

pylibraft 包包含 RAFT 算法和原语的 Python API。由于具有最小的依赖项并且接受任何支持 __cuda_array_interface__ 的对象，如 CuPy 的 ndarray，pylibraft 可以很好地集成到其他库中。此包中公开的 RAFT 算法数量会随着每个版本的发布而持续增长。

以下示例演示了计算 CuPy 数组之间的成对欧几里得距离。请注意，CuPy 不是 pylibraft 的必需依赖项。

import cupy as cp

from pylibraft.distance import pairwise_distance

n_samples = 5000
n_features = 50

in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)

output = pairwise_distance(in1, in2, metric="euclidean")

上述示例中的 output 数组为 raft.common.device_ndarray 类型，它支持 cuda_array_interface，使其与其他支持该接口的库（如 CuPy、Numba、PyTorch 和 RAPIDS cuDF）互操作。CuPy 支持 DLPack，这还使得从 raft.common.device_ndarray 到 JAX 和 Tensorflow 的零拷贝转换成为可能。

以下是将输出 pylibraft.device_ndarray 转换为 CuPy 数组的示例

cupy_array = cp.asarray(output)

以及将其转换为 PyTorch 张量

import torch

torch_tensor = torch.as_tensor(output, device='cuda')

或者将其转换为 RAPIDS cuDF 数据框

cudf_dataframe = cudf.DataFrame(output)

当相应的库已安装并在您的环境中可用时，此转换也可以通过设置全局配置选项自动由所有 RAFT 计算 API 完成

import pylibraft.config
pylibraft.config.set_output_as("cupy")  # All compute APIs will return cupy arrays
pylibraft.config.set_output_as("torch") # All compute APIs will return torch tensors

您还可以指定一个接受 pylibraft.common.device_ndarray 并执行自定义转换的 callable。以下示例将所有输出转换为 numpy 数组

pylibraft.config.set_output_as(lambda device_ndarray: return device_ndarray.copy_to_host())

pylibraft 还支持写入预分配的输出数组，因此任何支持 __cuda_array_interface__ 的数组都可以就地写入

import cupy as cp

from pylibraft.distance import pairwise_distance

n_samples = 5000
n_features = 50

in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
output = cp.empty((n_samples, n_samples), dtype=cp.float32)

pairwise_distance(in1, in2, out=output, metric="euclidean")

安装

RAFT的C++和Python库都可以通过Conda安装，Python库还可以通过Pip安装。

通过Conda安装C++和Python

安装RAFT最简单的方式是通过conda，并提供了一些软件包。

• libraft-headers C++头文件
• libraft（可选）包含预编译模板实例化和运行时API的C++共享库。
• pylibraft（可选）Python库
• raft-dask（可选）用于在Dask集群中部署使用RAFT raft::comms抽象层的多节点多GPU算法的Python库。
• raft-ann-bench（可选）一个基准测试工具，可以轻松生成与其它最先进实现相比的RAFT的向量搜索算法的基准测试。
• raft-ann-bench-cpu（可选）与上述类似的可重复基准测试工具，但不需要在机器上安装CUDA。可用于具有竞争性CPU的环境中进行测试。

根据您的CUDA版本，使用以下命令通过conda安装所有RAFT软件包（将rapidsai替换为rapidsai-nightly以安装更新但不够稳定的夜间软件包）。mamba比conda命令更受欢迎。

# for CUDA 11.8
mamba install -c rapidsai -c conda-forge -c nvidia raft-dask pylibraft cuda-version=11.8

# for CUDA 12.5
mamba install -c rapidsai -c conda-forge -c nvidia raft-dask pylibraft cuda-version=12.5

注意，上述命令还将安装libraft-headers和libraft。

您也可以使用上述mamba命令单独安装conda软件包。例如，如果您想将RAFT的头文件和预编译共享库安装到您的项目中

# for CUDA 12.5
mamba install -c rapidsai -c conda-forge -c nvidia libraft libraft-headers cuda-version=12.5

如果您正在安装C++ API，请参阅使用libraft以获取有关使用预编译共享库的更多信息。您还可以参考示例C++模板项目，它提供了一个可插入您的项目并针对上述安装的RAFT开发工件进行构建的CMake配置。

通过Pip安装Python

pylibraft和raft-dask都有实验性软件包，可以通过pip安装

pip install pylibraft-cu11 --extra-index-url=https://pypi.nvidia.com
pip install raft-dask-cu11 --extra-index-url=https://pypi.nvidia.com

这些软件包静态构建了RAFT的预编译实例化，因此C++头文件和预编译共享库在您的代码中不会立即可用。

构建说明中包含有关从源代码构建RAFT并将其包含在下游项目中的更多详细信息。构建说明中的从源代码构建RAFT C++和Python部分包含上述CPM代码片段的更全面版本。

您可以在cpp/template目录中找到一个RAFT项目模板示例，它演示了如何使用RAFT构建新应用程序或将其纳入现有的CMake项目。

贡献

如果您有兴趣为RAFT项目做出贡献，请阅读我们的贡献指南。有关开发指南、工作流程和原则的详细信息，请参阅开发者指南。

参考

在一般引用RAFT时，请考虑引用此Github项目。

@misc{rapidsai,
  title={Rapidsai/raft: RAFT contains fundamental widely-used algorithms and primitives for data science, Graph and machine learning.},
  url={https://github.com/rapidsai/raft},
  journal={GitHub},
  publisher={Nvidia RAPIDS},
  author={Rapidsai},
  year={2022}
}

在引用稀疏成对距离API时，请考虑使用以下bibtex

@article{nolet2021semiring,
  title={Semiring primitives for sparse neighborhood methods on the gpu},
  author={Nolet, Corey J and Gala, Divye and Raff, Edward and Eaton, Joe and Rees, Brad and Zedlewski, John and Oates, Tim},
  journal={arXiv preprint arXiv:2104.06357},
  year={2021}
}

在引用单链聚类API时，请考虑以下bibtex

@misc{nolet2023cuslink,
      title={cuSLINK: Single-linkage Agglomerative Clustering on the GPU},
      author={Corey J. Nolet and Divye Gala and Alex Fender and Mahesh Doijade and Joe Eaton and Edward Raff and John Zedlewski and Brad Rees and Tim Oates},
      year={2023},
      eprint={2306.16354},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

在引用CAGRA时，请考虑以下bibtex

@misc{ootomo2023cagra,
      title={CAGRA: Highly Parallel Graph Construction and Approximate Nearest Neighbor Search for GPUs},
      author={Hiroyuki Ootomo and Akira Naruse and Corey Nolet and Ray Wang and Tamas Feher and Yong Wang},
      year={2024},
      series = {ICDE '24}
}

在引用k选择例程时，请考虑以下bibtex

@proceedings{10.1145/3581784,
    title = {Parallel Top-K Algorithms on GPU: A Comprehensive Study and New Methods},
    author={Jingrong Zhang, Akira Naruse, Xipeng Li, and Yong Wang},
    year = {2023},
    isbn = {9798400701092},
    publisher = {Association for Computing Machinery},
    address = {New York, NY, USA},
    location = {Denver, CO, USA},
    series = {SC '23}
}

在引用最近邻下降API时，请考虑以下bibtex

@inproceedings{10.1145/3459637.3482344,
    author = {Wang, Hui and Zhao, Wan-Lei and Zeng, Xiangxiang and Yang, Jianye},
    title = {Fast K-NN Graph Construction by GPU Based NN-Descent},
    year = {2021},
    isbn = {9781450384469},
    publisher = {Association for Computing Machinery},
    address = {New York, NY, USA},
    url = {https://doi.org/10.1145/3459637.3482344},
    doi = {10.1145/3459637.3482344},
    abstract = {NN-Descent is a classic k-NN graph construction approach. It is still widely employed in machine learning, computer vision, and information retrieval tasks due to its efficiency and genericness. However, the current design only works well on CPU. In this paper, NN-Descent has been redesigned to adapt to the GPU architecture. A new graph update strategy called selective update is proposed. It reduces the data exchange between GPU cores and GPU global memory significantly, which is the processing bottleneck under GPU computation architecture. This redesign leads to full exploitation of the parallelism of the GPU hardware. In the meantime, the genericness, as well as the simplicity of NN-Descent, are well-preserved. Moreover, a procedure that allows to k-NN graph to be merged efficiently on GPU is proposed. It makes the construction of high-quality k-NN graphs for out-of-GPU-memory datasets tractable. Our approach is 100-250\texttimes{} faster than the single-thread NN-Descent and is 2.5-5\texttimes{} faster than the existing GPU-based approaches as we tested on million as well as billion scale datasets.},
    booktitle = {Proceedings of the 30th ACM International Conference on Information \& Knowledge Management},
    pages = {1929–1938},
    numpages = {10},
    keywords = {high-dimensional, nn-descent, gpu, k-nearest neighbor graph},
    location = {Virtual Event, Queensland, Australia},
    series = {CIKM '21}
}