Qdrant 中的大规模搜索

耗时:2 天难度:高级

在本教程中,我们将以现实世界中的数据集 LAION-400M 为例,描述一种经济高效地上传、索引和搜索海量数据的方法。

本教程旨在演示在保持合理的搜索延迟和准确性的前提下,索引和搜索大型数据集所需的最低资源配置。

所有相关的代码片段均可在 GitHub 仓库中找到。

本教程推荐的 Qdrant 版本为 v1.13.5 及更高版本。

数据集

我们将使用的数据集是 LAION-400M,这是一个包含约 4 亿个向量的集合,这些向量源自从 Common Crawl 数据集中提取的图像。每个向量均为 512 维,并使用 CLIP 模型生成。

向量关联了许多元数据字段,例如 urlcaptionLICENSE 等。

总有效载荷(payload)大小约为 200 GB,向量大小为 400 GB。

该数据集以 409 个数据块的形式提供,每个数据块包含约 100 万个向量。我们将使用以下 Python 脚本逐个上传数据集块。

硬件要求

经过初步实验,我们确定了完成此任务的最低硬件配置:

  • 8 核 CPU
  • 64 GB RAM
  • 650 GB 磁盘空间
Hardware configuration

硬件配置

此配置足以在单用户模式下索引和探索数据集;延迟足以构建交互式图表并在仪表板中导航。

当然,在生产级配置中,您可能需要更多的 CPU 核心和内存。

确保实验具备高网络带宽非常重要,因此请确保客户端和服务器运行在同一区域。

上传与索引

我们将使用以下 Python 脚本逐个上传数据集块。

export QDRANT_URL="https://xxxx-xxxx.xxxx.cloud.qdrant.io"
export QDRANT_API_KEY="xxxx-xxxx-xxxx-xxxx"

python upload.py

该脚本会逐一下载 LAION 数据集块并将其上传至 Qdrant。中间数据不会持久化到磁盘,因此脚本在客户端不需要占用过多的磁盘空间。

让我们看一下我们所使用的集合配置:

client.create_collection(
        QDRANT_COLLECTION_NAME,
        vectors_config=models.VectorParams(
            size=512, # CLIP model output size
            distance=models.Distance.COSINE, # CLIP model uses cosine distance
            datatype=models.Datatype.FLOAT16, # We only need 16 bits for float, otherwise disk usage would be 800Gb instead of 400Gb
            on_disk=True # We don't need original vectors in RAM
        ),
        # Even though CLIP vectors don't work well with binary quantization, out of the box,
        # we can rely on query-time oversampling to get more accurate results
        quantization_config=models.BinaryQuantization(
            binary=models.BinaryQuantizationConfig(
                always_ram=True,
            )
        ),
        optimizers_config=models.OptimizersConfigDiff(
            # Bigger size of segments are desired for faster search
            # However it might be slower for indexing
            max_segment_size=5_000_000, 
        ),
        # Having larger M value is desirable for higher accuracy,
        # but in our case we care more about memory usage
        # We could still achieve reasonable accuracy even with M=6 + oversampling
        hnsw_config=models.HnswConfigDiff(
            m=6, # decrease M for lower memory usage
            on_disk=False
        ),
    )

有几个重点需要注意:

  • 我们使用 FLOAT16 数据类型存储向量,这使我们能够以 FLOAT32 一半的大小存储向量,且对该数据集的准确性没有明显损失。
  • 我们使用 BinaryQuantization(并设置 always_ram=True)来启用查询时过采样(oversampling)。这使我们能够获得准确且资源高效的搜索,尽管 512 维的 CLIP 向量并不直接适合二值量化。
  • 我们使用 m=6HnswConfig 来降低内存使用。我们将在下一节中更深入地探讨内存使用情况。

此配置的目标是确保搜索过程中的预取(prefetch)组件无需从磁盘加载数据,并且至少保证向量和向量索引的最小版本始终驻留在 RAM 中。搜索的第二阶段可以明确决定我们可以承担多少次从磁盘加载数据的开销。

在我们的实验中,上传过程的速度为每秒 5000 个点。索引过程与上传并行进行,速度约为每秒 4000 个点。

Upload and indexation process

上传与索引过程

内存使用

在上传和索引过程完成后,让我们详细查看 Qdrant 服务器的内存使用情况。

Memory usage

内存使用

总体而言,内存使用由 3 个部分组成:

  • 系统内存 - 8.34 GB - 这是为内部系统和操作系统预留的内存,与数据集大小无关。
  • 数据内存 - 39.27 GB - 这是 Qdrant 进程的常驻内存,无法被驱逐;如果超过限制,Qdrant 进程将会崩溃。
  • 缓存内存 - 14.54 GB - 这是 Qdrant 使用的磁盘缓存。它对于快速搜索是必要的,但在需要时可以被驱逐。

我们最感兴趣的是数据内存和缓存内存。让我们看看这些组件中具体存储了什么。

在我们的场景中,Qdrant 使用内存存储以下组件:

  • 存储向量
  • 存储向量索引
  • 存储有关点 ID 和版本的信息

向量大小

在我们的场景中,我们仅在 RAM 中存储量化向量,因此计算所需大小相对容易。

400_000_000 * 512d / 8 bits / 1024 (Kb) / 1024 (Mb) / 1024 (Gb) = 23.84Gb

向量索引大小

向量索引稍微复杂一些,因为它不是一个简单的矩阵。

在内部,它存储为图中的连接列表,每个连接是一个 4 字节的整数。

连接数由 HNSW 索引的 M 参数定义,在我们的例子中,高层级为 6,0 层为 2 x M

这给了我们如下估算:

400_000_000 * (6 * 2) * 4 bytes / 1024 (Kb) / 1024 (Mb) / 1024 (Gb) = 17.881Gb

实际上,由于我们在 Qdrant v1.13.0 中实现的压缩技术,索引的大小会略小一些,但它仍然是一个很好的估算值。

Qdrant 中的 HNSW 索引以 mmap 形式存储,如果需要,它可以从 RAM 中驱逐。因此,HNSW 的内存消耗属于缓存内存范畴。

ID 和版本大小

Qdrant 必须存储每个点的附加信息,例如 ID 和版本。每次请求都需要这些信息,因此为了快速访问,将其保留在 RAM 中非常重要。

让我们深入了解 Qdrant 内部,以了解这些信息需要多少内存。


// This is s simplified version of the IdTracker struct
// It omits all optimizations and small details,
// but gives a good estimation of memory usage
IdTracker {
    // Mapping of internal id to version (u64), compressed to 4 bytes
    // Required for versioning and conflict resolution between segments
    internal_to_version, // 400M x 4 = 1.5Gb

    // Mapping of external id to internal id, 4 bytes per point.
    // Required to determine original point ID after search inside the segment
    internal_to_external: Vec<u128>, // 400M x 16 = 6.4Gb

    // Mapping of external id to internal id. For numeric ids it uses 8 bytes,
    //  UUIDs are stored as 16 bytes.
    // Required to determine sequential point ID inside the segment
    external_to_internal: Vec<u64, u32>, // 400M x (8 + 4) = 4.5Gb
}

在 v1.13.5 中,我们引入了一项重大优化,将 IdTracker 的内存使用量减少了约 2 倍。因此,在我们的案例中,IdTracker 的总内存使用量约为 12.4 GB

因此,在我们的案例中,Qdrant 服务器预期的总 RAM 使用量约为 23.84 GB + 17.881 GB + 12.4 GB = 54.121 GB,这非常接近我们观察到的实际内存使用量:39.27 GB + 14.54 GB = 53.81 GB

尽管我们对估算做了一些简化,但它们足以理解 Qdrant 服务器的内存使用情况。

在数据集上传并建立索引后,我们就可以开始搜索相似向量了。

我们可以先在 Web-UI 中探索数据集,这样您不仅可以从数据表格中,还可以直观地感受到搜索性能。

Web-UI Bear image

Web-UI 熊的图片

Web-UI similar Bear image

Web-UI 类似的熊的图片

Web-UI 的默认请求不使用过采样,但观察到的结果已经足够看出图像之间的相似性。

基准测试数据(Ground truth)

然而,为了更准确地评估搜索性能,我们需要将搜索结果与基准数据进行比较。遗憾的是,LAION 数据集不包含可用的基准数据,因此我们必须自己生成它。

为此,我们需要对数据集中的每个向量执行全扫描(full-scan)搜索,并将结果存储在单独的文件中。遗憾的是,这个过程非常耗时且需要大量资源,因此我们将查询数量限制为 100。我们提供了随时可用的 基准数据文件 和生成该文件的 脚本(需要 512 GB RAM 的机器,执行时间约 20 小时)。

我们的基准数据文件包含 100 个查询,每个查询有 50 个结果。使用数据集本身的前 100 个向量来生成这些查询。

搜索查询

为了精确控制过采样的数量,我们将使用以下搜索查询:


limit = 50
rescore_limit = 1000 # oversampling factor is 20

query = vectors[query_id] # One of existing vectors

response = client.query_points(
        collection_name=QDRANT_COLLECTION_NAME,
        query=query,
        limit=limit,
        # Go to disk 
        search_params=models.SearchParams(
            quantization=models.QuantizationSearchParams(
                rescore=True,
            ),
        ),
        # Prefetch is performed using only in-RAM data,
        # so querying even large amount of data is fast
        prefetch=models.Prefetch(
            query=query,
            limit=rescore_limit,
            params=models.SearchParams(
                quantization=models.QuantizationSearchParams(
                    # Avoid rescoring in prefetch
                    # We should do it explicitly on the second stage
                    rescore=False,
                ),
            )
        )
    )

如您所见,此查询包含两个阶段:

  • 第一阶段是预取,仅使用驻留在 RAM 中的数据执行。它速度非常快,并允许我们获得大量的候选者。
  • 第二阶段是重打分(rescore),使用存储在磁盘上的完整尺寸向量执行。

通过使用两阶段搜索,我们可以精确控制从磁盘加载的数据量,并确保搜索速度和准确性之间的平衡。

您可以在 eval.py 中找到搜索过程的完整代码。

性能优化

我们发现对该数据集有用的一个重要性能调整是启用 Qdrant 中的 异步 IO (Async IO)

默认情况下,Qdrant 使用同步 IO,这对于内存数据集表现良好,但在我们需要从磁盘读取大量数据时可能会成为瓶颈。

异步 IO(通过 io_uring 实现)允许向磁盘发送并行请求并充分利用磁盘带宽。

这正是我们在使用原始向量执行大规模重打分时所需要的。

与其逐个读取向量并等待磁盘响应 1000 次,不如向磁盘发送 1000 个请求并等待它们全部完成。这使我们能够饱和利用磁盘带宽并获得更快的响应。

要启用 Qdrant 中的异步 IO,您需要设置以下环境变量:

QDRANT__STORAGE__PERFORMANCE__ASYNC_SCORER=true

或者在配置文件中设置参数:

storage:
  performance:
    async_scorer: true

在 Qdrant Managed 云服务中,可以通过集群 Configuration 选项卡中的 Advanced optimizations 部分启用异步 IO。

Async IO configuration in Cloud

云端的异步 IO 配置

运行搜索请求

一旦所有准备工作完成,我们就可以运行搜索请求并评估结果。

您可以在 eval.py 中找到搜索过程的完整代码。

该脚本将以配置好的过采样因子运行 100 次搜索请求,并将结果与基准数据进行比较。

python eval.py --rescore_limit 1000

在我们的请求中,取得了以下结果:

重打分限制 (Rescore Limit)Precision@50 (准确率)单次请求耗时
100075.2%0.7 秒
500081.0%2.2 秒

使用 m=16 的额外实验证明,在 rescore_limit=1000 的情况下,我们可以达到 85% 的准确率,但这会需要更多的内存。

Log of search evaluation

搜索评估日志

结论

在本教程中,我们演示了如何在 Qdrant 中经济高效地上传、索引和搜索大数据集。如果结合查询时的过采样,即使是 512 维的向量也可以应用二值量化。

Qdrant 允许精确控制存储的每个部分位于何处,这使得在搜索速度和内存使用之间实现良好的平衡成为可能。

潜在改进

在本实验中,我们详细调查了存储的哪些部分负责内存使用,以及如何控制它们。

一个特别有趣的部分是 VectorIndex 组件,它负责存储向量之间的连接图。

在进一步的研究中,我们将探讨使 HNSW 对磁盘更友好的可能性,以便在不损失显著性能的情况下将其卸载到磁盘。

此页面有用吗?

感谢您的反馈!🙏

很遗憾听到这个消息。😔 您可以在 GitHub 上编辑此页面,或创建一个 GitHub 问题。