GPU Index
本指南概述了在 Milvus 中构建支持 GPU 的 index 的步骤,这可以显著提高高吞吐量和高召回率场景下的搜索性能。有关 Milvus 支持的 GPU index 类型的详细信息,请参阅 GPU Index。
配置 Milvus 设置进行 GPU 内存控制
Milvus 使用全局图形内存池来分配 GPU 内存。
它在 Milvus 配置文件 中支持两个参数 initMemSize
和 maxMemSize
。池大小最初设置为 initMemSize
,在超过此限制后将自动扩展到 maxMemSize
。
默认的 initMemSize
是 Milvus 启动时可用 GPU 内存的 1/2,默认的 maxMemSize
等于所有可用的 GPU 内存。
直到 Milvus 2.4.1(包含版本 2.4.1),Milvus 使用统一的 GPU 内存池。对于 2.4.1 之前的版本(包含版本 2.4.1),建议将这两个值都设置为 0。
gpu:
initMemSize: 0 #设置初始内存池大小。
maxMemSize: 0 #maxMemSize 设置最大内存使用限制。当内存使用超过 initMemSize 时,Milvus 将尝试扩展内存池。
从 Milvus 2.4.1 开始,GPU 内存池仅用于搜索期间的临时 GPU 数据。因此,建议将其设置为 2048 和 4096。
gpu:
initMemSize: 2048 #设置初始内存池大小。
maxMemSize: 4096 #maxMemSize 设置最大内存使用限制。当内存使用超过 initMemSize 时,Milvus 将尝试扩展内存池。
构建 Index
以下示例演示如何构建不同类型的 GPU index。
准备 Index 参数
设置 GPU index 参数时,定义 index_type、metric_type 和 params:
-
index_type (string):用于加速向量搜索的 index 类型。有效选项包括 GPU_CAGRA、GPU_IVF_FLAT、GPU_IVF_PQ 和 GPU_BRUTE_FORCE。
-
metric_type (string):用于测量向量相似性的度量类型。有效选项为 IP 和 L2。
-
params(dict):特定于 index 的构建参数。此参数的有效选项取决于 index 类型。
以下是不同 index 类型的示例配置:
-
GPU_CAGRA index
index_params = {
"metric_type": "L2",
"index_type": "GPU_CAGRA",
"params": {
'intermediate_graph_degree': 64,
'graph_degree': 32
}
}params 的可能选项包括:
-
intermediate_graph_degree (int):通过确定修剪前图的度数来影响召回率和构建时间。推荐值为 32 或 64。
-
graph_degree (int):通过设置修剪后图的度数来影响搜索性能和召回率。通常,它是 intermediate_graph_degree 的一半。这两个度数之间的差异越大,构建时间越长。其值必须小于 intermediate_graph_degree 的值。
-
build_algo (string):选择修剪前的图生成算法。可能的选项:
-
IVF_PQ:提供更高质量但构建时间较慢。
-
NN_DESCENT:提供更快的构建但可能召回率较低。
-
-
cache_dataset_on_device (string,"true" | "false"):决定是否将原始数据集缓存在 GPU 内存中。将此设置为 "true" 通过精化搜索结果来增强召回率,而设置为 "false" 则节省 GPU 内存。
-
-
GPU_IVF_FLAT 或 GPU_IVF_PQ index
index_params = {
"metric_type": "L2",
"index_type": "GPU_IVF_FLAT", # Or GPU_IVF_PQ
"params": {
"nlist": 1024
}
} -
GPU_BRUTE_FORCE index
index_params = {
'index_type': 'GPU_BRUTE_FORCE',
'metric_type': 'L2',
'params': {}
}不需要额外的 params 配置。
构建 Index
在 index_params 中配置 index 参数后,调用 create_index()
方法构建 index。
# Get an existing collection
collection = Collection("YOUR_COLLECTION_NAME")
collection.create_index(
field_name="vector", # Name of the vector field on which an index is built
index_params=index_params
)
搜索
构建 GPU index 后,下一步是在进行搜索之前准备搜索参数。
准备搜索参数
以下是不同 index 类型的示例配置:
-
GPU_BRUTE_FORCE index
search_params = {
"metric_type": "L2",
"params": {}
}不需要额外的 params 配置。
-
GPU_CAGRA index
search_params = {
"metric_type": "L2",
"params": {
"itopk_size": 128,
"search_width": 4,
"min_iterations": 0,
"max_iterations": 0,
"team_size": 0
}
}关键搜索参数包括:
-
itopk_size:确定搜索期间保持的中间结果大小。较大的值可能会提高召回率,但会降低搜索性能。它应该至少等于最终的 top-k (limit) 值,通常是 2 的幂(例如,16、32、64、128)。
-
search_width:指定搜索期间进入 CAGRA 图的入口点数量。增加此值可以增强召回率,但可能影响搜索性能。
-
min_iterations / max_iterations:这些参数控制搜索迭代过程。默认情况下,它们设置为 0,CAGRA 根据 itopk_size 和 search_width 自动确定迭代次数。手动调整这些值可以帮助平衡性能和准确性。
-
team_size:指定用于在 GPU 上计算度量距离的 CUDA 线程数。常见值是最多 32 的 2 的幂(例如 2、4、8、16、32)。它对搜索性能的影响很小。默认值为 0,其中 Milvus 根据向量维度自动选择 team_size。
-
-
GPU_IVF_FLAT 或 GPU_IVF_PQ index
search_params = {
"metric_type": "L2",
"params": {"nprobe": 10}
}这两种 index 类型的搜索参数与 IVF_FLAT 和 IVF_PQ 中使用的类似。有关更多信息,请参阅 进行向量相似性搜索。
执行搜索
使用 search()
方法在 GPU index 上执行向量相似性搜索。
# Load data into memory
collection.load()
collection.search(
data=[[query_vector]], # Your query vector
anns_field="vector", # Name of the vector field
param=search_params,
limit=100 # Number of the results to return
)
限制
使用 GPU index 时,请注意某些约束:
-
对于 GPU_IVF_FLAT,limit 的最大值为 1024。
-
对于 GPU_IVF_PQ 和 GPU_CAGRA,limit 的最大值为 1024。
-
虽然 GPU_BRUTE_FORCE 的 limit 没有设定限制,但建议不要超过 4096 以避免潜在的性能问题。
-
目前,GPU index 不支持 COSINE 距离。如果需要 COSINE 距离,应首先对数据进行归一化,然后可以使用内积 (IP) 距离作为替代。
-
GPU index 的加载 OOM 保护不完全支持,过多的数据可能导致 QueryNode 崩溃。
FAQ
-
何时适合使用 GPU index?
GPU index 在要求高吞吐量或高召回率的情况下特别有益。例如,在处理大批量时,GPU index 的吞吐量可以超过 CPU index 多达 100 倍。在较小批量的场景中,GPU index 在性能方面仍然显著优于 CPU index。此外,如果需要快速数据插入,加入 GPU 可以大幅加速 index 构建过程。
-
CAGRA、GPU_IVF_PQ、GPU_IVF_FLAT 和 GPU_BRUTE_FORCE 等 GPU index 在哪些场景中最适合?
CAGRA index 非常适合要求增强性能的场景,尽管代价是消耗更多内存。对于优先考虑内存节约的环境,GPU_IVF_PQ index 可以帮助最小化存储需求,尽管这会带来更高的精度损失。GPU_IVF_FLAT index 作为平衡选项,在性能和内存使用之间提供折衷。最后,GPU_BRUTE_FORCE index 专为穷举搜索操作而设计,通过执行遍历搜索来保证 1 的召回率。