faiss Offline IVF 实战：基于 big batch search 的十亿级向量离线 KNN 全流程

Offline IVF 是 faiss 仓库中位于 demos/offline_ivf 的一个端到端演示工程：它把 faiss 的 big batch search 机制、分片索引、GPU 加速的 kNN 计算与集群作业调度组合在一起，用于对超大规模（示例为 SSNPP 10 亿条 256 维向量）数据集做离线全库最近邻检索。读完本文，你将掌握从数据集分片、索引训练、分片建库到集群检索与一致性校验的完整操作流程，并理解每个步骤在 offline_ivf.py 中的具体实现逻辑。

一、环境准备

按照 demos/offline_ivf/README.md 的说明，工程依赖 faiss 的 GPU 版本，官方给出的环境创建方式如下：

conda create --name oivf python=3.10
conda activate oivf
conda install -c pytorch/label/nightly -c nvidia faiss-gpu=1.7.4
conda install tqdm
conda install pyyaml
conda install -c conda-forge submitit

注意其中的适用前提：

• 安装的是 faiss-gpu=1.7.4 的 nightly 版本，检索阶段（search 命令）会实际调用 faiss.knn_gpu 与 StandardGpuResources，因此必须有可用的 GPU；
• submitit 是集群作业调度库，--cluster_run 模式依赖它向 Slurm 集群提交作业（见 run.py）；
• tqdm、pyyaml 分别用于进度显示和读取 yaml 配置。

二、数据集分片与配置文件生成

2.1 分片脚本

原始数据库（例如 big-ann-benchmarks 格式的 .u8bin 文件）需要先切分成多个可独立读取的分片文件。仓库中对应的脚本是 create_sharded_ssnpp_files.py（README 中提及的 create_sharded_dataset.py 即此脚本的现名）：

python create_sharded_ssnpp_files.py \
    --filepath /path/to/FB_ssnpp_database.u8bin \
    --data_batch 50000000 \
    --output_dir /path/to/ssnpp_data

脚本逻辑很直接：先用 xbin_mmap 以 np.memmap 方式零拷贝映射整个 .u8bin 文件（跳过 8 字节的 n/d 头部），然后按 --data_batch（默认 5000 万条/文件，要求能整除总数）逐批 np.save 成 ssnpp_0000000000.npy、ssnpp_0000000001.npy 这样的 npy 分片。

2.2 配置生成

generate_config.py 遍历指定目录下的分片文件，自动探测每个文件是 raw（无头部的原始二进制）还是 npy 格式，并统计每个文件的向量条数，最终 yaml.dump 出一份数据集描述配置。它接受脚本顶部硬编码的 root 目录与文件名列表（示例中是 20 个 ssnpp_{i:010}.npy，维度 d=256、dtype=uint8）。

2.3 配置文件结构

仓库自带的示例配置 config_ssnpp.yaml 展示了完整的配置格式，各字段含义如下：

字段	示例值	含义
d	256	向量维度
output	/checkpoint/.../ssnpp	索引与结果的输出根目录
index.prod / index.non-prod	IVF8192,PQ128 等	index factory 字符串，prod 用于生产运行，non-prod 用于多参数扫描
nprobe.prod / nprobe.non-prod	512 / 256~8192	各实验组使用的 nprobe 值
k	50	每个查询返回的近邻数
index_shard_size	50000000	每个索引分片包含的向量条数
query_batch_size	50000000	检索阶段每批处理的查询条数
evaluation_sample	10000	评估（evaluate）时的查询采样条数
training_sample	1572864	训练索引时采样的向量条数
datasets.<名称>	ssnpp_1B	数据集描述：root 目录、size 总数、files 列表（每个文件含 dtype/format/name/size）

示例中 ssnpp_1B 数据集共 20 个 5000 万条的 uint8 npy 分片，总计 10 亿条向量，这正是 README Run book 中 --xb ssnpp_1B 引用的名称。

三、Run book：五步操作命令

3.1 train_index：训练索引模板

python run.py --command train_index --config config_ssnpp.yaml --xb ssnpp_1B

对应 offline_ivf.py 的 train_index()：

1. 用 self.xb_ds.get_first_n(training_sample, np.float32) 读取数据前 1572864 条，并经 np.unique(axis=0) 去重（对去重场景的数据集尤为重要）；
2. faiss.index_factory(d, "IVF8192,PQ128", metric) 构建索引并 train；
3. 通过 faiss.extract_index_ivf 取出内部 IVF 节点，强制设置 by_residual = True（以残差编码提升 PQ 精度）；
4. 用 faiss.write_index 保存为 *.empty.faissindex 模板文件——只含训练好的量化工具，不含任何向量，后续所有分片都基于它构建。

命令已存在模板文件时会直接断言失败，避免重复训练。

3.2 index_shard：生成索引分片

python run.py --command index_shard --config config_ssnpp.yaml --xb ssnpp_1B

对应 index_shard()：

1. 读取索引模板，把 CPU quantizer 用 faiss.index_cpu_to_all_gpus 迁移到 GPU 加速粗量化；
2. 按 index_shard_size 切分数据（10 亿 / 5000 万 = 20 个分片），逐分片 index.reset() 后流式 add_with_ids，其中 ID 是全局连续的（start + jj），保证检索结果可以直接定位回原始数据；
3. 若配置了 OPQ 等前置变换（IndexPreTransform），输入会先经过 _iterate_transformed 做变换；
4. 每写完一个分片立即 faiss.write_index 落盘为 *.shard_{i}；已存在的分片文件会跳过（FileExistsError 分支），支持断点续建。

分片数量在构造函数中计算：nshards = ceil(xb_ds.size / index_shard_size)（见 offline_ivf.py#L95-L98）。

3.3 search：集群批量检索

python run.py --command search --config config_ssnpp.yaml --xb ssnpp_1B \
    --cluster_run --partition <PARTITION-NAME>

这是整个工程的核心，对应 search()。其流程：

1. 打开分片索引：_open_sharded_index 读取模板索引后，把所有分片的 InvertedLists 收集进一个指针向量，用 faiss.HStackInvertedLists 水平拼接后 replace_invlists 挂回索引（offline_ivf.py#L768-L793）。这样无需合并物理文件，内存映射即可逻辑上得到一个完整索引；
2. 按 query_batch_size 批量处理：每批查询先经 GPU quantizer 做粗量化，得到 q_assign（每个查询的 nprobe 个候选倒排列表）；
3. 调用 big_batch_search：这是离线检索效率的关键。它以 method="knn_function" 方式传入自定义 _knn_function——内部用 faiss.knn_gpu 在指定 GPU 上以 float16 计算一个倒排列表内全部（查询子集 × 库向量）的距离，并设置查询/向量内存上限（4GB / 8GB）防止显存溢出；
4. 容错与断点：每个批次输出到 I{offset}_...npy / D_approx{offset}_...npy，文件已存在则整批跳过；Slurm 任务还会写 record_{offset}.txt 记录 job id，若同一 job 重跑且结果文件为空，会清理后重算。

big_batch_search 的实现在 contrib/big_batch_search.py：reorder_assign 把所有查询按其粗量化结果做桶排序（bucket sort），随后每个倒排列表只处理被分配到该列表的查询子集，桶级计算结果再统一汇入全局 ResultHeap。相比逐查询检索，它把每个桶的计算变成一次大型矩阵运算，天然契合 GPU GEMM；threaded=faiss.get_num_gpus() * 8、prefetch_threads 等参数让"预取下一桶 + 计算当前桶 + 写回结果"流水线并行；还支持 checkpoint/checkpoint_freq 定期落盘（search 中每 7200 秒一次），进程被杀后可从检查点恢复。

关于查询集：README 特别说明，search 默认假设数据库向量即查询向量（args.xq 缺省时自动等于 args.xb，见 offline_ivf.py#L52-L53）。如果查询集不同，需要先在配置中准备新数据集（第 1 步），再传入 --xq：

python run.py --command search --config config_ssnpp.yaml --xb ssnpp_1B --xq <QUERIES_DATASET_NAME>

3.4 evaluate 与 consistency_check

README 最后一条命令用于任意步骤后的健全性检查：

python run.py --command consistency_check --config config_ssnpp.yaml --xb ssnpp_1B

consistency_check() 的校验链相当完整（offline_ivf.py#L817-L899）：

• 索引模板存在：模板文件可读；
• 索引分片自洽：从每个分片偏移 i*shard_size + offset 处抽 10000 条向量，以 nprobe=1 搜 100 近邻，断言每条向量都能在自己的 100 个近邻中找到自己；
• 拼接后索引一致：同样的抽样在 HStackInvertedLists 拼接索引上重做一遍；
• 检索结果抽查：从 knn 输出目录加载每个批次的 I/D 结果文件，与在线 index.search 的前 10000 条对比，要求索引重合率 > 0.95、距离之和 allclose(rtol=0.01)。

此外工程还支持 README 未展开的 evaluate 命令：对 evaluation_sample 条查询分别做精确检索（faiss.contrib.exhaustive_search 的 knn_ground_truth）与近似检索，输出逐 k 的 knn_intersection_measure（近邻交集率）、精确/近似/精化距离均值以及"交集基数分布"（utils.py 的 get_intersection_cardinality_frequencies）。配置中若开启 evaluate_by_margin，还会基于 2 * D(a,b) / (mean_a + mean_b) 这一双向检索 margin 公式（utils.py 的 margin）统计双向互为近邻的"可信近邻对"，这是去重类应用的典型评估口径。

四、run.py 命令行参数全解

run.py 是所有命令的统一入口，--command 取值即 OfflineIVF 的方法名（train_index / index_shard / search / merge_index / evaluate / consistency_check 等），通过 eval(f"oivf.{args.command}()") 动态分发（run.py#L32-L39）。主要参数：

参数	默认值	说明
--command（必填）	—	要执行的步骤
--config（必填）	—	数据集/索引配置 yaml
--xb（必填）	—	数据库向量数据集名（对应 config 的 datasets 键）
--xq	等于 --xb	查询向量数据集名
--nt	96	搜索 OpenMP 线程数，启动时调用 faiss.omp_set_num_threads
--no_residuals	关闭残差	训练时不设 by_residual
--cluster_run	False	通过 submitit 提交 Slurm 作业数组
--job_name / --num_nodes / --tasks_per_node	oivf / 1 / 1	作业命名与节点配置
--gpus_per_node / --cpus_per_task	8 / 80	单节点 GPU/CPU 配额
--logs_dir	/checkpoint/.../logs	集群日志目录
--partition	learnlab	Slurm 分区名，choices 中列出了若干集群分区

集群模式下，process_options_and_run_jobs 的调度策略是（run.py#L42-L75）：evaluate 命令会为 index（factory 字符串）× nprobe 的每个组合各提交一个作业（non-prod 列表用于网格扫描）；其他命令则按 factory 字符串各提交一个作业，本地运行则只取 prod 列表的最后一个值。

五、输出文件布局

所有产物都组织在 output/<xb>/、output/<xq>/、output/<xq>_in_<xb>/ 三棵目录下：

• <xb>/{factory}.empty.faissindex：训练模板；
• <xb>/{factory}.shard_{i}：索引分片；
• <xq>_in_<xb>/knn/：检索结果 I{offset}_{factory}_np{nprobe}.npy、D_approx{offset}_...npy 及检查点文件 CP...；
• <xq>_in_<xb>/eval/：评估中间量（I_a_gt.npy、D_a_ann_...npy、margin_*.npy 等）。

文件命名把 factory、nprobe 编码进文件名，使得多组参数的扫描结果互不覆盖，可独立对比。

六、测试用例

demos/offline_ivf/tests 提供了用小规模数据（8 维、数千条向量、IVF256,PQ4 或 OPQ4,IVF256,PQ4）跑通全流程的单元测试：

• test_train_index：验证 train_index 产出 *.empty.faissindex 模板；
• test_index_shard_equal_file_sizes / test_index_shard_unequal_file_sizes：覆盖分片大小整除/不整除库大小两种切分情形，逐一确认 .shard_i 文件生成；
• test_search：串联 train → index_shard → search，确认 knn/I... 结果文件落盘；
• test_evaluate_without_margin（含 OPQ 变体）：额外走 merge_index 后跑 evaluate，检查 eval 目录下全部中间文件；
• test_consistency_check：验证在未完成前置步骤时一致性检查会正确报错。

测试数据由 tests/testing_utils.py 的 TestDataCreator 生成，可直接作为小数据集跑通本工程的参考配置。

七、适用边界小结

• 该工程面向离线批处理场景：全库/全查询集的 KNN 计算，产出 npy 结果文件，不是在线服务索引；
• search 依赖 GPU（faiss.knn_gpu），而 train_index/index_shard/consistency_check 主要在 CPU 侧完成；
• 数据集要求能被 index_shard_size / query_batch_size 整除（search 中有整除性断言），配置中的每个文件 size 之和必须等于数据集总 size（dataset.py 的断言）；
• 从源码结构看，big_batch_search 目前支持 IndexIVFFlat、IndexIVFPQ、IndexIVFScalarQuantizer 三类 IVF 索引（BlockComputer），示例配置的 IVF8192,PQ128 即属于 IVFPQ 路径。