一般来说,为了向量数据库运行更加高效,我们需要按照以下 OpenSearch 集群的设计指导原则,选择 OpenSearch 的资源配置:

如果偏向向量数据库搜索的工作负载那么应该使用 10-30GB 的 shard 大小,如果偏向日志的工作负载,应该采用 30-50GB 的节点大小;

请尝试将分片数量设为数据节点数量的偶数倍,有助于分片在数据节点上均匀分布;

每个节点的分片数,与 JVM 堆内存成正比,每 GB 内存的分片不超过 25 个;

每个 5 个 vCPU 能对应一个分片,比如 8 个 vCPU 则最多支持 6 个 shard;

如果有启用 k-NN 字段,参考如下表格进行内存的推算。

依据目前的信息,仅仅知道所要索引的向量数据库原始文档数。由于文档切分等中间处理过程,无法估算向量数据库具体的内存用量和存储量,所以需要用小批量的实验数据进行测试推演。

在小批量实验中向量数据库尝试索引了 300 篇文档,通过切分共产生出了约 203k 条记录,4.5GB 存储量。那么按比例换算如果需要向量数据库索引 1 万篇文档,那么会产生约 700w 记录,150GB 存储。如下图所示:

从知识问答 Chatbot 场景出发,属于搜索的工作负载,Shard 的大小应在 10-30GB 范围内以增加搜索性能。shard 数量一般遵循节点数的倍数的原则,假设是 2 节点集群那么可以是 [2, 4, 8, 16 …]。以 150GB 存储总量进行计算,shard 可以为 8, 10, 12, 14, 16,当 shard 数为 8 时,每个 shard 存储量为 18.75GB,符合要求。

向量数据库检索方面为了同时保证 recall 和 latency,采用了 HNSW 算法。另外参考 中的 benchmark 结论,HNSW 算法中 m 值可以设定为 16。那么向量数据库内存规划方面,依照上表公式进行内存占用的推算:

一般每个节点的堆外内存占比 50%,根据knn.memory.circuit_breaker.limit=70% 的最佳实践设定,那么向量数据库 35% 的节点内存被 KNN 占用,那么推算出整个节点的内存应为 22.9GB / 35% = 65GB。

vCPU 的规划方面,假设 shard 数为 8,乘以 1.5vCPU/Shard 的系数,vCPU 个数至少需要为 12 以上。结合如下 C 系和 R 系的实例配置信息和价格信息,综合考虑内存和 vCPU 的要求,选择 C 系的 2 节点 c6g.4xlarge 或者 R 系的 2 节点 r6g.2xlarge。

数据完整性,保证所有的知识都能被查询到,不会因为摄入异常导致数据缺失。

构建速度,向量数据库知识召回部分可能存在反复的效果调整,需要反复多次摄入,速度对全链路开发调优的效率很重要。

向量数据库查询性能,保证场景中的实时会话体验。

整个摄入过程,从过程上基本可以划分为三个阶段:文本切分、文本向量化以及摄入 Amazon OpenSearch。其中文本切分和文本向量化的处理是临时性的工作负载,原则上可以提升 glue job 的并发数和 Amazon SageMaker Endpoint 背后的节点数来线性提高对这块工作负载的处理速度,但 OpenSearch 属于一个预分配的资源(注:今年即将发布的 OpenSearch Severless k-NN 向量数据库会改变这一点)。后两个部分,即向量化和 OpenSearch 摄入可能会是整个流程的瓶颈,完整流程测试不容易进行拆解分析性能瓶颈,所以本试验会分别对这两部分进行测试。

实验 1 – Embedding Model 吞吐测试

1. 使用 paraphrase-multilingual-deploy.ipynb 进行部署,部署 10 台 g4dn.xlarge 机型;

https://github.com/aws-samples/private-llm-qa-bot/blob/main/notebooks/embedding/paraphrase-multilingual-deploy.ipynb

2. 注释掉下游写入 OpenSearch 造成的影响,暂时注释掉相关代码;

https://github.com/aws-samples/private-llm-qa-bot/blob/main/code/aos_write_job.py

3. 利用 batch_upload_docs.py 启动多 glue job 进行并发运行。

https://github.com/aws-samples/private-llm-qa-bot/blob/main/code/batch_upload_docs.py

这部分处理流程中,通过调整 glue job 的并行度与 client-side batch size 可以调整向量化这一步骤的吞吐能力。当 GPU 利用率不足时,提高 client-side batch size 能够提高 GPU 的利用率。经过简单测试发现,确实能证明这个假设,具体数据可以参考如下实验结果:

实验 2 – Amazon OpenSearch 摄入测试

1. 随机生成向量,替换掉 Embedding 模型调用,参考如下代码:

2. 构建 OpenSearch 集群以及索引,并优化设置;

a. 构建对应的索引

向量字段涉及到的两个参数 ef_construction 和 m。ef_construction 指定构建 k-NN 图的时候的动态列表大小,值越大其向量数据的图更精确,但索引的速度也响应更慢。m指定 k-NN 中每个向量双向链表的数量,越大检索越准确,但相应内存占用会显著增大。参考博客中的benchmark结论(https://aws.amazon.com/cn/blogs/big-data/choose-the-k-nn-algorithm-for-your-billion-scale-use-case-with-opensearch/),对于当前的数据规模,参数 ef_construction:128 和 m:16 已经足以保证召回率,另外在构建索引时可以关注以下几点:

i. 添加一个 publish_date 字段方便后续根据时间来删除/更新知识;

ii. 添加 idx 整型字段用于记录对应片段在全文中的顺序,在召回时可以基于 range_search 召回相邻上下文片段;

iii. 只做过滤不做关键字召回的字段设置成 keyword 类型,有利于索引速度。具体可以参考如下代码:

b.设置 knn 相关参数 可参考《亚马逊云科技向量数据库与生成式AI的完美融合:落地实践详解(一)》的内容。

c.开启多 glue job 进行并发摄入,可以参考如下代码:

3. 部分实验结果明细

每轮实验中,调整的参数已经用加粗字体标注出来,供参考用于指导后续的数据注入中的参数调整。

实验 3 – 全流程摄入测试

a. 部分实验记录明细

b. 初步实验结论

参考以上的实验记录可知,1 万篇文档拆分成 700 万条向量后,通过调整客户端并发,推理端点的节点数和推理 Batch Size 可以在 1 小时左右完成摄入,且完整性没有问题。能够满足大规模向量数据库知识构建的要求,如果文档量继续增长,可以继续扩展 OpenSearch 节点和 SageMaker Endpoint 节点。