使用 RAPIDS RAFT 進(jìn)行機(jī)器學(xué)習(xí)和數(shù)據(jù)分析的可重用計算模式

在許多數(shù)據(jù)分析和機(jī)器學(xué)習(xí)算法中，計算瓶頸往往來自控制端到端性能的一小部分步驟。這些步驟的可重用解決方案通常需要低級別的基元，這些基元非常簡單且耗時。

NVIDIA 制造 RAPIDS RAFT 是為了解決這些瓶頸，并在為多維數(shù)據(jù)構(gòu)建算法時最大限度地重用，例如機(jī)器學(xué)習(xí)和數(shù)據(jù)分析中經(jīng)常遇到的問題。

RAPIDS 是 GPU 上的一套用于數(shù)據(jù)科學(xué)和機(jī)器學(xué)習(xí)的加速庫：

pandas 類數(shù)據(jù)結(jié)構(gòu)的 cuDF

cuGraph 用于圖形數(shù)據(jù)

cuML 用于機(jī)器學(xué)習(xí)

高度優(yōu)化的 RAFT 計算模式構(gòu)成了一個豐富的模塊化嵌入式加速器目錄，為您提供了強(qiáng)大的元素來組成新的算法或加速現(xiàn)有的庫。

這僅僅是一個開始：隨著新 GPU 架構(gòu)的發(fā)布， RAFT 組件將繼續(xù)優(yōu)化，確保您始終從硬件中獲得最佳性能。

RAFT 使您能夠花時間設(shè)計和開發(fā)應(yīng)用程序，而不必?fù)?dān)心您是否能充分利用 GPU 硬件。

在這篇文章中，我討論了 RAFT 在開發(fā)人員工具箱中的位置，使用它的環(huán)境，更重要的是，在需要時使用 RAFT 的權(quán)力。

消除常見瓶頸

NVIDIA 構(gòu)建 RAFT ，為開發(fā)者提供基本元素作為構(gòu)建塊。

最近的鄰居就是一個很好的例子。它很常見，很有用，而且計算量很大。鄰域方法涵蓋了 machine learning 的大部分算法，如聚類和降維。支持線性代數(shù)、稀疏矩陣運算、采樣、優(yōu)化和統(tǒng)計矩等應(yīng)用數(shù)學(xué)的核心工具。

事實上， RAFT 構(gòu)成了 RAPIDS cuML 中幾乎所有的算法，包括但不限于流行的 HDBSCAN 、 TSNE 、 UMAP 以及所有其他用于聚類和可視化的算法。

表 1 。 RAFT 包含幾個不同類別的構(gòu)建塊，包括數(shù)據(jù)格式、稀疏和密集運算、最近鄰等空間運算、基本聚類、迭代求解器和統(tǒng)計。

如果你想在絕對最快的作品基礎(chǔ)上創(chuàng)作新作品，那么你現(xiàn)在只會受到創(chuàng)造力的限制。除了所有的基本方法外，向量搜索領(lǐng)域的最佳方法是在搜索大型語言模型（ LLM ）和推薦系統(tǒng)方面取得令人興奮的進(jìn)展。

圖 1 。 NVIDIA H100 GPU 上的新 RAFT IVF-PQ 算法與 CPU 上的 HNSW 算法的比較

圖 1 重點介紹了一種新的 RAFT 最先進(jìn)的近似最近鄰算法，它是一類被稱為 inverted file indexes （ IVF ）的算法中 product quantization （ PQ ）的變體。 IVF 將訓(xùn)練數(shù)據(jù)劃分為一組集群，并將查詢減少到僅最近集群的子集。這是 PQ 算法的 IVF 版本，稱為 IVF-PQ 。

RAFT IVF-PQ 實現(xiàn)已針對生產(chǎn)向量相似性搜索系統(tǒng)所需的小批量查詢進(jìn)行了調(diào)整。這類針對小批量調(diào)整的算法中目前最先進(jìn)的技術(shù)被稱為分層可導(dǎo)航小世界圖（ HNSW ）算法。它尤其以在 CPU 上速度快而聞名。新的 RAFT IVF-PQ 算法的訓(xùn)練速度比 HNSW 快 95 倍，在執(zhí)行小批量查詢時快 3 倍。

打造始終如一的體驗

RAFT 在其核心是一個只有頭部的 C ++模板庫，它需要一組最小的依賴項。它主要依賴于 CUDA 工具包附帶的庫。這使您能夠靈活地根據(jù)需要專門化所需數(shù)據(jù)類型的模板（圖 2 ）。

圖 2:RAFT 技術(shù)堆棧

圖 2 顯示了堆棧從一個僅包含頭的模板庫開始。在純頭庫的基礎(chǔ)上構(gòu)建了一個共享庫，其中包含預(yù)編譯的模板專用化和主機(jī)可訪問的運行時 API 。 PyLibRAFT 為 Python 用戶提供了強(qiáng)大的 RAFT 構(gòu)建塊。

一個可選的共享庫可以通過預(yù)編譯常見類型的模板專用化來減少編譯時間。共享庫中還提供了一個運行時 API ，可以鏈接到您的構(gòu)建中。它使 RAFT API 能夠以類似于從 C 源文件調(diào)用 cuBLAS 、 cuSOLVER 和 cuSPARSE API 的方式從常規(guī) C ++源文件調(diào)用。

直接使用共享庫中的運行時 API 的是一個名為pylibraft的 Python API ，它包含 C ++ API 的輕量級包裝器。

管理設(shè)備資源

RAFT 使用raft::resources對象來管理不同的資源，如 CUDA 流、流池和各種 CUDA 庫的句柄，如 cuBLAS 和 cuSOLVER 。raft::device_resources實例是配置和管理 GPU 特定資源以調(diào)用 RAFT API 的最簡單方法。

#include 
raft::device_resources resource_handle;

下面是 Python 中的相同示例：

from pylibraft.common.handle import DeviceResources
resource_handle = DeviceResources()

使用 CUDA 構(gòu)建復(fù)雜的端到端算法傳統(tǒng)上需要低水平的專業(yè)知識和關(guān)于每個 GPU 體系結(jié)構(gòu)的能力的高級知識，以持續(xù)保持硬件繁忙。

像 Cub 、 Thrust 和 CUTLASS 這樣的庫使編寫 CUDA 應(yīng)用程序變得更加容易。它們將較低級別的 API 抽象為較高級別的原語，這些原語可用于開發(fā)各種算法。

RAFT 提供了一個類似的抽象層，但它特別關(guān)注 ML 和數(shù)據(jù)分析的稍微高級的原語。如果你已經(jīng)熟悉深度神經(jīng)網(wǎng)絡(luò)的嵌入式加速器 CUDA 庫 cuDNN ，你可以說 RAFT 與 RAPIDS cuML 等庫的關(guān)系與 cuDNN 與 TensorFlow 等庫相似。

處理多維數(shù)據(jù)

雖然開始出現(xiàn)與 GPU 交互的標(biāo)準(zhǔn)，但 RAFT C ++ API 基于 C ++ 23 STL 標(biāo)準(zhǔn)中的 mdspan （多維非擁有視圖）。

mdspan是一種非所有權(quán)視圖結(jié)構(gòu)，具有高度的表達(dá)能力、靈活性和自文檔化。它有一個靈活的 API 來表示多維數(shù)據(jù)，在精神上類似于 NumPy 的ndarray，但在 C ++中。mdspan提供了干凈和一致的 API 體驗，因為它包裹了任何現(xiàn)有的指針，無論是在主（主機(jī)）內(nèi)存還是設(shè)備內(nèi)存中。

#include 
#include 

std::vector my_floats(10);

… populate vector …

auto my_mdspan = raft::make_host_vector(my_floats.data(), 10);

下面是 Python 中的相同示例：

#include 

float *my_floats;

… allocate and populate device memory …

auto my_mdspan = raft::make_device_vector(my_floats, 10);

為了簡化具有多維表示的內(nèi)存分配，mdarray標(biāo)準(zhǔn)提供了一個 RAII 兼容的內(nèi)存，與 RAFT 也采用的mdspan相對應(yīng)。通過使用mdarray 來分配和包含內(nèi)存，可以使示例變得更容易。

#include 
#include 

raft::device_resources handle;

int n_rows = 10;
int n_cols = 10;

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

下面是 Python 中的相同示例：

#include 

int n_rows = 10;
int n_cols = 10;

auto scalar = raft::make_host_scalar(1.0);
auto vector = raft::make_host_vector(n_cols);
auto matrix = raft::make_host_matrix(n_rows, n_cols);

在mdarray實例中分配并包含內(nèi)存后，創(chuàng)建mdspan 視圖以使用.view方法調(diào)用 RAFT API 。

// Scalar mdspan on device
auto scalar_view = scalar.view();

// Vector mdspan on device
auto vector_view = vector.view();

// Matrix mdspan on device
auto matrix_view = matrix.view();

使用 C ++和 Python API

例如，使用新的 RAFT 近似最近鄰 API 在 C ++和 Python 中構(gòu)建和查詢索引。以下代碼示例使用 RAFT IVF-PQ 算法構(gòu)建索引。

#include 

raft::device_resources handle;
raft::neighbors::ivf_pq::index_params idx_params;
auto index = raft::neighbors::ivf_pq::build(handle, idx_params, dataset);

您可以使用以下代碼示例來查詢新建的索引。

#include 

raft::neighbors::ivf_pq::search_params search_params;

uint32_t n_query_rows = 1000;
uint32_t k = 5;
…

auto out_inds = raft::make_device_matrix(handle, n_query_rows, k);
auto out_dists = raft::make_device_matrix(handle, n_query_rows, k);

raft::neighbors::ivf_pq::search(handle, search_params, index, queries,                                     
            out_inds.view(), out_dists.view());

下面是 Python 中相同的索引構(gòu)建示例：

import cupy as cp
from pylibraft.neighbors import ivf_pq
n_samples = 50000
n_query_rows = 1000
index_params = ivf_pq.IndexParams(
     n_lists=1024,
     metric="sqeuclidean",
     pq_dim=10)
index = ivf_pq.build(index_params, dataset)

下面是 Python 中相同的索引搜索示例：

search_params = ivf_pq.SearchParams(n_probes=20)
k = 10
…
distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(), index,
                                     queries, k)

PyLibRAFT 可以與任何支持 __cuda_array_interface__ （ CAI ）的庫進(jìn)行互操作，如 PyTorch 、 CuPy 和 Numba 。默認(rèn)情況下， PyLibRAFT 函數(shù)可以接受任何符合 CAI 的數(shù)組，即使在需要時也可以將輸出寫入到位。您可以通過配置執(zhí)行輸出轉(zhuǎn)換的函數(shù)，將此行為自定義到您選擇的庫中，以實現(xiàn)更無縫的集成。

以下代碼示例將 PyLibRAFT 配置為返回 CuPy ndarray或 PyTorch tensors。

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 函數(shù)的所有輸出轉(zhuǎn)換為主（主機(jī)）存儲器中的 NumPy 數(shù)組。

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

獲取 RAFT

以下步驟提供了有關(guān)安裝 RAFT 的不同方式的信息。 Conda 和 PIP 是最簡單的路線，尤其是在與 PyLibRAFT 交互時。由于 RAFT 是一個核心的 C ++庫，我們還簡要展示了它可以多么容易地集成到 Cmake 中。

康達(dá)牌手表

對于 C ++和 Python 使用，安裝 RAFT 的一種簡單方法是使用 conda 。將 C ++標(biāo)頭安裝到您的 conda 環(huán)境中。

conda install -c rapidsai -c conda-forge -c nvidia libraft-headers

將預(yù)編譯的二進(jìn)制文件也安裝到您的環(huán)境中。

conda install -c rapidsai -c conda-forge -c nvidia libraft-distance libraft-nn

當(dāng)然， Python 庫也有 conda 版本：

conda install -c rapidsai -c conda-forge -c nvidia pylibraft

皮普

RAPIDS 還提供了 pip 包，可以在純 Python 環(huán)境中輕松安裝。

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

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

C 制造商

RAFT 也可以使用 CMake Package Manager （ CPM ）和 rapids-cmake 輕松集成到您的 cmake 項目中，它們還可以下載所需版本并配置raft::raft cmake 目標(biāo)。

如果 RAFT 已經(jīng)安裝，例如之前的 conda 命令，請使用 cmake 的find_package使raft::raft目標(biāo)可用于您的項目。

find_package(raft)

raft::raft目標(biāo)攜帶了許多主要 RAFT 功能所需的基本依賴項集，如 RAPIDS 內(nèi)存管理器（ RMM ）和 CUDA 工具包庫。但是，如果您使用來自以下任何命名空間的 API ，則可能需要額外的依賴項：

raft::distance

raft::neighbors

raft::comms

RAFT 使用可選的 cmake 組件，這使得新的 cmakes 目標(biāo)能夠?qū)⑦@些依賴關(guān)系傳遞到您的項目中。

find_package(raft COMPONENTS distance distributed)

此 cmake 功能可啟用目標(biāo) raft :： distance 和 raft :: distributed ，此外還有 raft ：： raft 。將它們添加到必要的構(gòu)建目標(biāo)中，以便傳遞包含路徑和依賴關(guān)系的鏈接庫信息。例如，可以為距離組件鏈接 CUTRASS ，或者為分布式組件鏈接 NVIDIA Collective Communications Library （ NCCL ）和 Unififed Communications X （ UCX ）。

主要收獲

RAFT 是一個用于機(jī)器學(xué)習(xí)和數(shù)據(jù)分析的高度可重用計算模式庫。集中化重要計算使您能夠在進(jìn)行優(yōu)化時自動獲得性能改進(jìn)的好處，例如在發(fā)布新的 GPU 體系結(jié)構(gòu)和功能時。

RAFT 還包含高度可重復(fù)使用的基礎(chǔ)設(shè)施，用于構(gòu)建具有干凈靈活接口的生產(chǎn)質(zhì)量加速庫，遵循 C ++ STL 等重要標(biāo)準(zhǔn)，以保證使用壽命。

RAFT 通過覆蓋從 C ++到 Python 的整個堆棧來簡化集成，進(jìn)一步實現(xiàn)了與其他流行的 GPU 加速庫的互操作性，如 PyTorch 、 CuPy 和 Numba ，以及 RAPIDS 生態(tài)系統(tǒng)中的庫。