From a9e9655f50da9911cb09390a89cab7b8ac61aa91 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Tue, 1 Sep 2026 09:35:50 +0000 Subject: [PATCH 01/17] Clean up RaBitQ auto tuning config Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../inc/Core/SPANN/ExtraStaticSearcher.h | 38 ++++--------------- .../build_SPANN_sift1m_rabitq3_global.ini | 1 - .../build_SPANN_sift1m_raw_static_control.ini | 1 - .../search_SPANN_sift1m_raw_static_n12.ini | 1 - .../search_SPANN_sift1m_raw_static_n20.ini | 1 - .../search_SPANN_sift1m_raw_static_n32.ini | 1 - Tools/OPQ/OPQ_gpu_train_infer.py | 6 +-- Tools/OPQ/README.md | 4 +- docs/GettingStart.md | 1 - docs/RaBitQ_Global_Quantizer.md | 3 +- 10 files changed, 14 insertions(+), 43 deletions(-) diff --git a/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h b/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h index ee01c01f7..03d323e58 100644 --- a/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h +++ b/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h @@ -10,7 +10,6 @@ #include "inc/Core/Common/TruthSet.h" #include "Compressor.h" -#include #include #include #include @@ -177,12 +176,6 @@ namespace SPTAG } virtual bool LoadIndex(Options& p_opt) override { - m_opt = &p_opt; - m_enableDeltaEncoding = p_opt.m_enableDeltaEncoding; - m_enablePostingListRearrange = p_opt.m_enablePostingListRearrange; - m_enableDataCompression = p_opt.m_enableDataCompression; - m_enableDictTraining = p_opt.m_enableDictTraining; - m_extraFullGraphFile = p_opt.m_indexDirectory + FolderSep + p_opt.m_ssdIndex; std::string curFile = m_extraFullGraphFile + "_" + std::to_string(m_layer); p_opt.m_searchPostingPageLimit = max(p_opt.m_searchPostingPageLimit, static_cast((p_opt.m_postingVectorLimit * (p_opt.m_dim * sizeof(ValueType) + sizeof(SizeType)) + PageSize - 1) / PageSize)); @@ -218,6 +211,12 @@ namespace SPTAG } while (fileexists(curFile.c_str())); m_oneContext = (m_indexFiles.size() == 1); + m_opt = &p_opt; + m_enableDeltaEncoding = p_opt.m_enableDeltaEncoding; + m_enablePostingListRearrange = p_opt.m_enablePostingListRearrange; + m_enableDataCompression = p_opt.m_enableDataCompression; + m_enableDictTraining = p_opt.m_enableDictTraining; + if (m_enablePostingListRearrange) m_parsePosting = &ExtraStaticSearcher::ParsePostingListRearrange; else m_parsePosting = &ExtraStaticSearcher::ParsePostingList; if (m_enableDeltaEncoding) m_parseEncoding = &ExtraStaticSearcher::ParseDeltaEncoding; @@ -242,20 +241,12 @@ namespace SPTAG bool) { const uint32_t postingListCount = static_cast(p_exWorkSpace->m_postingIDs.size()); - if (postingListCount > p_exWorkSpace->m_pageBuffers.size() || - postingListCount > p_exWorkSpace->m_diskRequests.size()) { - SPTAGLIB_LOG(Helper::LogLevel::LL_Error, - "Static search workspace is too small: postings=%u buffers=%zu requests=%zu.\n", - postingListCount, p_exWorkSpace->m_pageBuffers.size(), - p_exWorkSpace->m_diskRequests.size()); - return ErrorCode::Fail; - } COMMON::QueryResultSet& queryResults = *((COMMON::QueryResultSet*)&p_queryResults); + int diskRead = 0; int diskIO = 0; int listElements = 0; - int missingPostingIDs = 0; #if defined(ASYNC_READ) && !defined(BATCH_READ) int unprocessed = 0; @@ -266,7 +257,6 @@ namespace SPTAG auto curPostingID = p_exWorkSpace->m_postingIDs[pi]; auto it = m_globalVectorIDToHeadMap.find(curPostingID); if (it == m_globalVectorIDToHeadMap.end()) { - ++missingPostingIDs; auto& request = p_exWorkSpace->m_diskRequests[pi]; request.m_readSize = 0; request.m_success = false; @@ -286,13 +276,6 @@ namespace SPTAG listElements += listInfo->listEleCount; size_t totalBytes = (static_cast(listInfo->listPageCount) << PageSizeEx); - if (totalBytes > p_exWorkSpace->m_pageBuffers[pi].GetPageSize()) { - SPTAGLIB_LOG(Helper::LogLevel::LL_Error, - "Static posting %d requires %zu bytes but its workspace buffer has %zu bytes.\n", - curPostingID, totalBytes, - p_exWorkSpace->m_pageBuffers[pi].GetPageSize()); - return ErrorCode::DiskIOFail; - } #ifdef ASYNC_READ auto& request = p_exWorkSpace->m_diskRequests[pi]; @@ -348,12 +331,6 @@ namespace SPTAG #endif } - if (missingPostingIDs > 0) { - SPTAGLIB_LOG(Helper::LogLevel::LL_Warning, - "Static search skipped %d of %u posting IDs that are absent from the head map.\n", - missingPostingIDs, postingListCount); - } - #ifdef ASYNC_READ #ifdef BATCH_READ BatchReadFileAsync(m_indexFiles, (p_exWorkSpace->m_diskRequests).data(), postingListCount); @@ -1819,6 +1796,7 @@ namespace SPTAG int m_vectorInfoSize = 0; int m_iDataDimension = 0; + int m_totalListCount = 0; int m_listPerFile = 0; diff --git a/Script_AE/iniFile/build_SPANN_sift1m_rabitq3_global.ini b/Script_AE/iniFile/build_SPANN_sift1m_rabitq3_global.ini index c80706484..d029c7521 100644 --- a/Script_AE/iniFile/build_SPANN_sift1m_rabitq3_global.ini +++ b/Script_AE/iniFile/build_SPANN_sift1m_rabitq3_global.ini @@ -53,7 +53,6 @@ TmpDir=/datadisk/yfcc_fast/sptag_rabitq_parity_sift1m_rabitq3/tmp EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false -PostingQuantizer=None Rerank=0 EnableADC=true diff --git a/Script_AE/iniFile/build_SPANN_sift1m_raw_static_control.ini b/Script_AE/iniFile/build_SPANN_sift1m_raw_static_control.ini index b3e27a5a7..17753c161 100644 --- a/Script_AE/iniFile/build_SPANN_sift1m_raw_static_control.ini +++ b/Script_AE/iniFile/build_SPANN_sift1m_raw_static_control.ini @@ -53,7 +53,6 @@ TmpDir=/datadisk/yfcc_fast/sptag_rabitq_parity_sift1m_raw_static_lambda001/tmp EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false -PostingQuantizer=None Rerank=0 [SearchSSDIndex] diff --git a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n12.ini b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n12.ini index 9a6da2b5a..86e139ab3 100644 --- a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n12.ini +++ b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n12.ini @@ -29,7 +29,6 @@ NumberOfThreads=24 EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false -PostingQuantizer=None Rerank=0 [SearchSSDIndex] diff --git a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n20.ini b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n20.ini index 5d31e46b2..4ab99f786 100644 --- a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n20.ini +++ b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n20.ini @@ -29,7 +29,6 @@ NumberOfThreads=24 EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false -PostingQuantizer=None Rerank=0 [SearchSSDIndex] diff --git a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n32.ini b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n32.ini index 533f51033..b67eb0695 100644 --- a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n32.ini +++ b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n32.ini @@ -29,7 +29,6 @@ NumberOfThreads=24 EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false -PostingQuantizer=None Rerank=0 [SearchSSDIndex] diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index 638f9d859..121d5a9ca 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -386,7 +386,7 @@ def train_rabitq(args): os.rename(args.output_truth, os.path.join(output_dir, 'truth.txt' + '.' + str(args.task))) if args.quan_test > 0: - queryreader = DataReader(args.query_file, args.dim, -1, args.query_normalize, args.data_type, args.target_type) + queryreader = DataReader(args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) numQuery, query = queryreader.readbatch() qid2ground_truths = {} @@ -508,7 +508,7 @@ def train_pq(args): os.rename(args.output_truth, os.path.join(output_dir, 'truth.txt' + '.' + str(args.task))) if args.quan_test > 0: - queryreader = DataReader(args.query_file, args.dim, -1, args.query_normalize, args.data_type, args.target_type) + queryreader = DataReader(args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) numQuery, query = queryreader.readbatch() qid2ground_truths = {} @@ -644,7 +644,7 @@ def train_opq(args): os.rename(args.output_truth, os.path.join(output_dir, 'truth.txt' + '.' + str(args.task))) if args.quan_test > 0: - queryreader = DataReader(args.query_file, args.dim, -1, args.query_normalize, args.data_type, args.target_type) + queryreader = DataReader(args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) numQuery, query = queryreader.readbatch() qid2ground_truths = {} diff --git a/Tools/OPQ/README.md b/Tools/OPQ/README.md index 8fccd1f03..6ed9507f7 100644 --- a/Tools/OPQ/README.md +++ b/Tools/OPQ/README.md @@ -12,5 +12,5 @@ ## Example python3 OPQ_gpu_train_infer.py --data_file perftest_vector.bin --query_file perftest_query.bin --task 0 --data_type float32 --k 5 --dim 64 --B 1000000 --Q 1000 --D L2 --data_format DEFAULT --T 20 --train_samples 1000000 --quan_type opq --quan_dim 32 --output_quantizer quantizer.bin -python3 OPQ_gpu_train_infer.py --data_file laion_5M.bin --query_file laion_test.bin --data_type float32 --target_type float32 --k 100 --dim 768 --B 1000000 --Q 1000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 96 --quan_test 1 -python3 OPQ_gpu_train_infer.py --data_file openai_5M.bin --query_file openai_test.bin --data_type float32 --target_type float32 --k 100 --dim 1536 --B 1000000 --Q 1000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 192 --quan_test 1 --data_normalize 1 --query_normalize 1 \ No newline at end of file +python3 OPQ_gpu_train_infer.py --data_file laion_5M.bin --query_file laion_test.bin --data_type float32 --target_type float32 --k 100 --dim 768 --B 1000000 --Q 10000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 96 --quan_test 1 +python3 OPQ_gpu_train_infer.py --data_file openai_5M.bin --query_file openai_test.bin --data_type float32 --target_type float32 --k 100 --dim 1536 --B 1000000 --Q 10000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 192 --quan_test 1 --data_normalize 1 --query_normalize 1 \ No newline at end of file diff --git a/docs/GettingStart.md b/docs/GettingStart.md index c2baa2752..a743fafb6 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -316,7 +316,6 @@ TmpDir=/tmp/sift1m-spann-rabitq3 EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false -PostingQuantizer=None Rerank=0 EnableADC=true diff --git a/docs/RaBitQ_Global_Quantizer.md b/docs/RaBitQ_Global_Quantizer.md index fd09365b7..3d563cfac 100644 --- a/docs/RaBitQ_Global_Quantizer.md +++ b/docs/RaBitQ_Global_Quantizer.md @@ -34,5 +34,4 @@ multiple of 64. The current adapter supports the official L2 estimator; cosine distance is intentionally unsupported. `Script_AE/iniFile/build_SPANN_sift1m_rabitq3_global.ini` is the canonical -SIFT1M example. It uses STATIC postings containing the global RaBitQ codes; -keep `PostingQuantizer=None` because RaBitQ is already the global quantizer. +SIFT1M example. It uses STATIC postings containing the global RaBitQ codes. From 5a19d93171b14e0b6025321b4ee4ab08877695f2 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:09:05 +0000 Subject: [PATCH 02/17] Implement RaBitQ storage bit auto tuning Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Tools/OPQ/OPQ_gpu_train_infer.py | 252 +++++++++++++++++---------- Tools/OPQ/README.md | 22 ++- Tools/OPQ/test_rabitq_auto_tuning.py | 69 ++++++++ 3 files changed, 247 insertions(+), 96 deletions(-) create mode 100644 Tools/OPQ/test_rabitq_auto_tuning.py diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index 121d5a9ca..eb3750ffb 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -8,6 +8,7 @@ import heapq import argparse import copy +import json from operator import itemgetter import os import subprocess @@ -40,13 +41,25 @@ def get_config(): parser.add_argument('--output_quan_vector_file', type = str, default = "", help='quantized vectors') parser.add_argument('--output_rec_vector_file', type = str, default = "", help = "reconstruct vectors") parser.add_argument('--quan_test', type = int, default = 0, help='compare with ground truth') + parser.add_argument('--rabitq_auto_tune', action = 'store_true', help='select the minimum RaBitQ storage bits before encoding') + parser.add_argument('--rabitq_target_recall', type = float, default = 0.95, help='minimum Recall@k for RaBitQ auto tuning') + parser.add_argument('--rabitq_min_bits', type = int, default = 1, help='minimum RaBitQ storage bits to evaluate') + parser.add_argument('--rabitq_max_bits', type = int, default = 8, help='maximum RaBitQ storage bits to evaluate') + parser.add_argument('--rabitq_tuning_result', type = str, default = 'rabitq_auto_tuning.json', help='auto-tuning result file under output_dir') args = parser.parse_args() return args +def is_binary_vector_file(filename): + suffixes = ('.bin', '.fbin', '.u8bin', '.i8bin') + if filename.endswith(suffixes): + return True + stem, separator, shard = filename.rpartition('.') + return separator != '' and shard.isdigit() and stem.endswith(suffixes) + class DataReader: def __init__(self, filename, featuredim, batchsize, normalize, datatype, targettype='float32'): self.mytype = targettype - if filename.find('.bin') >= 0: + if is_binary_vector_file(filename): self.fin = open(filename, 'rb') R = unpack('i', self.fin.read(4))[0] self.featuredim = unpack('i', self.fin.read(4))[0] @@ -287,110 +300,161 @@ def evaluate(retrieve_results: List[List[int]], return MRR, Recall +def sptag_rabitq_storage_bytes(dim, bits): + if dim <= 0: + raise ValueError('RaBitQ dimension must be positive') + if bits < 1 or bits > 8: + raise ValueError(f'RaBitQ bits must be in [1, 8], got {bits}') + padded_dimension = ((dim + 63) // 64) * 64 + return padded_dimension * bits // 8 + 5 * np.dtype(np.float32).itemsize + +def rabitq_bits_from_quantized_dimension(dim, quan_dim): + total_bits = quan_dim * 8 + if dim <= 0 or quan_dim <= 0 or total_bits % dim != 0: + raise ValueError(f'quan_dim={quan_dim} does not represent an integral RaBitQ bit count for dimension {dim}') + bits = total_bits // dim + if bits < 1 or bits > 8: + raise ValueError(f'RaBitQ bits must be in [1, 8], got {bits}') + return bits + +def load_ground_truth(path, query_count, topk): + ground_truths = [] + with open(path, 'r') as truth_file: + for query_id in range(query_count): + line = truth_file.readline() + if not line: + raise ValueError(f'ground truth contains only {query_id} queries, expected {query_count}') + neighbors = line.strip().split() + if len(neighbors) < topk: + raise ValueError( + f'ground truth query {query_id} contains {len(neighbors)} neighbors, expected at least {topk}') + ground_truths.append(set(int(neighbor) for neighbor in neighbors[:topk])) + return ground_truths + +def recall_at_k(faiss_index, queries, ground_truths, topk, batch_size=64): + if len(queries) != len(ground_truths): + raise ValueError('query and ground-truth counts differ') + if len(queries) == 0: + raise ValueError('at least one query is required') + recall_sum = 0.0 + for start in tqdm.tqdm(range(0, len(queries), batch_size)): + end = min(start + batch_size, len(queries)) + _, results = faiss_index.search(np.asarray(queries[start:end]), topk) + for offset, candidates in enumerate(results): + truth = ground_truths[start + offset] + recall_sum += len(truth.intersection(int(candidate) for candidate in candidates if candidate >= 0)) / len(truth) + return recall_sum / len(queries) + +def create_rabitq_index(faiss, dim, bits, training_data): + faiss_index = faiss.index_factory(dim, f"RaBitQ{bits}", faiss.METRIC_L2) + faiss_index.train(training_data) + return faiss_index + +def add_rabitq_data(args, faiss_index): + datareader = DataReader( + args.data_file, args.dim, args.B, args.data_normalize, args.data_type, args.target_type) + total = 0 + while True: + num_data, data = datareader.readbatch() + if num_data == 0: + break + faiss_index.add(data) + total += num_data + datareader.close() + if total == 0: + raise ValueError('RaBitQ input data is empty') + return total + +def tune_rabitq_bits(args, faiss, training_data, queries, ground_truths): + if not 0.0 < args.rabitq_target_recall <= 1.0: + raise ValueError('rabitq_target_recall must be in (0, 1]') + if args.rabitq_min_bits < 1 or args.rabitq_max_bits > 8 or args.rabitq_min_bits > args.rabitq_max_bits: + raise ValueError('RaBitQ tuning range must satisfy 1 <= min_bits <= max_bits <= 8') + + trials = [] + for bits in range(args.rabitq_min_bits, args.rabitq_max_bits + 1): + print(f'Auto tuning RaBitQ{bits} for Recall@{args.k} >= {args.rabitq_target_recall:.6f}') + candidate = create_rabitq_index(faiss, args.dim, bits, training_data) + data_count = add_rabitq_data(args, candidate) + recall = recall_at_k(candidate, queries, ground_truths, args.k) + trials.append({'bits': bits, 'recall': recall}) + print(f'RaBitQ{bits} Recall@{args.k}: {recall:.6f}') + if recall >= args.rabitq_target_recall: + return bits, candidate, data_count, trials + + measured = ', '.join(f'{trial["bits"]}-bit={trial["recall"]:.6f}' for trial in trials) + raise RuntimeError( + f'No RaBitQ bit count in [{args.rabitq_min_bits}, {args.rabitq_max_bits}] ' + f'reached Recall@{args.k} >= {args.rabitq_target_recall:.6f}; {measured}') + def train_rabitq(args): import faiss output_dir = args.output_dir - if args.train_samples > args.B: args.train_samples = args.B - - datareader = DataReader(args.data_file, args.dim, args.B, args.data_normalize, args.data_type, args.target_type) - - print (f'train RabitQ using {args.train_samples} samples ...') - - numData, data = datareader.readbatch() - + if args.D not in ('L2', 'Cosine'): + raise ValueError('RaBitQ training supports only L2, or Cosine with normalized data and queries') + if args.D == 'Cosine' and (args.data_normalize == 0 or args.query_normalize == 0): + raise ValueError('Cosine RaBitQ tuning requires data_normalize=1 and query_normalize=1') + if args.train_samples <= 0: + raise ValueError('train_samples must be positive') + if args.k <= 0 or args.Q <= 0: + raise ValueError('k and Q must be positive') + if len(args.output_quan_vector_file) > 0 or len(args.output_rec_vector_file) > 0: + raise ValueError( + 'RaBitQ vectors must be generated by the native SPTAG quantizer after tuning; ' + 'do not use output_quan_vector_file or output_rec_vector_file') + training_reader = DataReader( + args.data_file, args.dim, args.train_samples, args.data_normalize, args.data_type, args.target_type) + num_training, training_data = training_reader.readbatch() + training_reader.close() + if num_training == 0: + raise ValueError('RaBitQ training data is empty') + print(f'train RaBitQ using {num_training} samples ...') faiss.omp_set_num_threads(args.T) - nbytes = int(args.quan_dim * 8 // args.dim) - print (f'nbytes:{nbytes}') - faiss_index = faiss.index_factory(len(data[0]), f"RaBitQ{nbytes}", faiss.METRIC_L2) - print('Training the index with doc embeddings') - - faiss_index.train(data[0:args.train_samples]) - - rtype = np.uint8(0) - if args.data_type == 'uint8': - rtype = np.uint8(1) - elif args.data_type == 'int16': - rtype = np.uint8(2) - elif args.data_type == 'float32': - rtype = np.uint8(3) - - ivf_index = faiss.downcast_index(faiss_index) - #centroid_embedings = faiss.vector_to_array(ivf_index.pq.centroids) - #codebooks = centroid_embedings.reshape(ivf_index.pq.M, ivf_index.pq.ksub, ivf_index.pq.dsub) - #print ('codebooks shape:') - #print (codebooks.shape) - - #codebooks = codebooks.astype(np.float32) - #with open(os.path.join(output_dir, args.output_quantizer + '.' + str(args.task)),'wb') as f: - # f.write(pack('B', 1)) - # f.write(pack('B', rtype)) - # f.write(pack('i', codebooks.shape[0])) - # f.write(pack('i', codebooks.shape[1])) - # f.write(pack('i', codebooks.shape[2])) - # f.write(codebooks.tobytes()) - - if args.quan_test == 0 and len(args.output_quan_vector_file) == 0 and len(args.output_rec_vector_file) == 0: - os.rename(args.output_truth, os.path.join(output_dir, 'truth.txt' + '.' + str(args.task))) - return - - if len(args.output_quan_vector_file) > 0: - fquan = open(os.path.join(output_dir, args.output_quan_vector_file + '.' + str(args.task) + '.tmp'), 'wb') - fquan.write(pack('i', 0)) - fquan.write(pack('i', args.quan_dim)) - - if len(args.output_rec_vector_file) > 0: - frec = open(os.path.join(output_dir, args.output_rec_vector_file + '.' + str(args.task) + '.tmp'), 'wb') - frec.write(pack('i', 0)) - frec.write(pack('i', data.shape[1])) - - writeitems = 0 - while numData > 0: - if args.quan_test > 0: faiss_index.add(data) - - if len(args.output_quan_vector_file) > 0: - codes = ivf_index.pq.compute_codes(data) - print ('codes shape:') - print (codes.shape) - fquan.write(codes.tobytes()) - - if len(args.output_rec_vector_file) > 0: - reconstructed = ivf_index.pq.decode(codes).astype(args.data_type) - frec.write(reconstructed.tobytes()) - - writeitems += numData - numData, data = datareader.readbatch() - - datareader.close() - - if len(args.output_quan_vector_file) > 0: - p = fquan.tell() - fquan.seek(0) - fquan.write(pack('i', writeitems)) - fquan.seek(p) - fquan.close() - if os.path.exists(os.path.join(output_dir, args.output_quan_vector_file + '.' + str(args.task))): - os.remove(os.path.join(output_dir, args.output_quan_vector_file + '.' + str(args.task))) - os.rename(os.path.join(output_dir, args.output_quan_vector_file + '.' + str(args.task) + '.tmp'), os.path.join(output_dir, args.output_quan_vector_file + '.' + str(args.task))) - if len(args.output_rec_vector_file) > 0: - p = frec.tell() - frec.seek(0) - frec.write(pack('i', writeitems)) - frec.seek(p) - frec.close() - if os.path.exists(os.path.join(output_dir, args.output_rec_vector_file + '.' + str(args.task))): - os.remove(os.path.join(output_dir, args.output_rec_vector_file + '.' + str(args.task))) - os.rename(os.path.join(output_dir, args.output_rec_vector_file + '.' + str(args.task) + '.tmp'), os.path.join(output_dir, args.output_rec_vector_file + '.' + str(args.task))) - - os.rename(args.output_truth, os.path.join(output_dir, 'truth.txt' + '.' + str(args.task))) + trials = [] + data_count = None + if args.rabitq_auto_tune: + if args.quan_test <= 0: + raise ValueError('rabitq_auto_tune requires quan_test > 0 and a pre-generated ground truth') + queryreader = DataReader( + args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) + num_query, queries = queryreader.readbatch() + queryreader.close() + if num_query != args.Q: + raise ValueError(f'query file contains {num_query} queries, but configured Q is {args.Q}') + ground_truths = load_ground_truth(args.output_truth, num_query, args.k) + bits, faiss_index, data_count, trials = tune_rabitq_bits( + args, faiss, training_data, queries, ground_truths) + result = { + 'selected_bits': bits, + 'native_quantizer_qd': bits, + 'storage_bytes_per_vector': sptag_rabitq_storage_bytes(args.dim, bits), + 'target_recall': args.rabitq_target_recall, + 'recall_at': args.k, + 'query_count': num_query, + 'data_count': data_count, + 'trials': trials, + } + result_path = os.path.join(output_dir, args.rabitq_tuning_result) + temporary_result_path = result_path + '.tmp' + with open(temporary_result_path, 'w') as result_file: + json.dump(result, result_file, indent=2) + result_file.write('\n') + os.replace(temporary_result_path, result_path) + print(f'Selected RaBitQ storage bits: {bits}; result: {result_path}') + else: + bits = rabitq_bits_from_quantized_dimension(args.dim, args.quan_dim) + faiss_index = create_rabitq_index(faiss, args.dim, bits, training_data) + if args.quan_test > 0: + data_count = add_rabitq_data(args, faiss_index) - if args.quan_test > 0: + if args.quan_test > 0 and not args.rabitq_auto_tune: queryreader = DataReader(args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) numQuery, query = queryreader.readbatch() qid2ground_truths = {} - f = open(os.path.join(output_dir, 'truth.txt.' + str(args.task)), 'r') + f = open(args.output_truth, 'r') for i in range(numQuery): items = f.readline()[0:-1].strip().split(' ') qid2ground_truths[i] = set([int(gt) for gt in items]) diff --git a/Tools/OPQ/README.md b/Tools/OPQ/README.md index 6ed9507f7..bf5cc3145 100644 --- a/Tools/OPQ/README.md +++ b/Tools/OPQ/README.md @@ -12,5 +12,23 @@ ## Example python3 OPQ_gpu_train_infer.py --data_file perftest_vector.bin --query_file perftest_query.bin --task 0 --data_type float32 --k 5 --dim 64 --B 1000000 --Q 1000 --D L2 --data_format DEFAULT --T 20 --train_samples 1000000 --quan_type opq --quan_dim 32 --output_quantizer quantizer.bin -python3 OPQ_gpu_train_infer.py --data_file laion_5M.bin --query_file laion_test.bin --data_type float32 --target_type float32 --k 100 --dim 768 --B 1000000 --Q 10000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 96 --quan_test 1 -python3 OPQ_gpu_train_infer.py --data_file openai_5M.bin --query_file openai_test.bin --data_type float32 --target_type float32 --k 100 --dim 1536 --B 1000000 --Q 10000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 192 --quan_test 1 --data_normalize 1 --query_normalize 1 \ No newline at end of file +python3 OPQ_gpu_train_infer.py --data_file laion_5M.bin --query_file laion_test.bin --output_truth truth.txt --output_dir rabitq_tuned --data_type float32 --target_type float32 --k 100 --dim 768 --B 1000000 --Q 10000 --D L2 --train_samples 1000000 --quan_type rabitq --quan_test 1 --rabitq_auto_tune --rabitq_target_recall 0.95 --rabitq_min_bits 1 --rabitq_max_bits 8 +python3 OPQ_gpu_train_infer.py --data_file openai_5M.bin --query_file openai_test.bin --output_truth truth.txt --output_dir rabitq_tuned --data_type float32 --target_type float32 --k 100 --dim 1536 --B 1000000 --Q 10000 --D Cosine --train_samples 1000000 --quan_type rabitq --quan_test 1 --rabitq_auto_tune --rabitq_target_recall 0.95 --data_normalize 1 --query_normalize 1 + +## RaBitQ storage-bit auto tuning + +`--rabitq_auto_tune` runs before vector encoding or index construction. It evaluates +RaBitQ bit counts in ascending order and selects the first (therefore minimum) +count whose exhaustive `Recall@k` reaches `--rabitq_target_recall`. Every candidate +uses the same first `--train_samples` vectors, the same configured `--Q` queries, +and the pre-generated `--output_truth` top-`k` neighbors. + +The command fails instead of silently choosing a bit count when the configured +query/ground-truth count is unavailable or no candidate in +`[--rabitq_min_bits, --rabitq_max_bits]` reaches the target. The selected storage +bit count, actual SPTAG byte width (including padded dimensions and five Float +factors), measured recalls, and sample counts are written atomically to +`/rabitq_auto_tuning.json`. Use `native_quantizer_qd` as `-qd` for the +native SPTAG `quantizer` command that generates the quantizer and vectors consumed +by the subsequent index build. The tuner deliberately does not emit Faiss codes: +their persisted layout is not the SPTAG global RaBitQ format. \ No newline at end of file diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py new file mode 100644 index 000000000..f8469527c --- /dev/null +++ b/Tools/OPQ/test_rabitq_auto_tuning.py @@ -0,0 +1,69 @@ +import importlib.util +from pathlib import Path +import tempfile +import unittest + +import numpy as np + + +MODULE_PATH = Path(__file__).with_name('OPQ_gpu_train_infer.py') +SPEC = importlib.util.spec_from_file_location('opq_gpu_train_infer', MODULE_PATH) +MODULE = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(MODULE) + + +class FakeIndex: + def __init__(self, results): + self.results = np.asarray(results, dtype=np.int64) + + def search(self, queries, topk): + count = len(queries) + results = self.results[:count, :topk] + self.results = self.results[count:] + return np.zeros(results.shape, dtype=np.float32), results + + +class RaBitQAutoTuningTest(unittest.TestCase): + def test_data_reader_accepts_fbin(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'vectors.fbin' + with path.open('wb') as output: + np.asarray([2, 3], dtype=np.int32).tofile(output) + np.arange(6, dtype=np.float32).tofile(output) + reader = MODULE.DataReader(str(path), 3, 2, 0, 'float32') + count, vectors = reader.readbatch() + reader.close() + self.assertEqual(2, count) + np.testing.assert_array_equal(vectors, np.arange(6, dtype=np.float32).reshape(2, 3)) + + def test_binary_reader_accepts_numbered_shards(self): + self.assertTrue(MODULE.is_binary_vector_file('vectors.bin.0')) + self.assertTrue(MODULE.is_binary_vector_file('vectors.fbin.12')) + self.assertFalse(MODULE.is_binary_vector_file('vectors.txt.0')) + + def test_sptag_storage_bytes_include_padding_and_factors(self): + self.assertEqual(68, MODULE.sptag_rabitq_storage_bytes(128, 3)) + self.assertEqual(68, MODULE.sptag_rabitq_storage_bytes(127, 3)) + self.assertEqual(92, MODULE.sptag_rabitq_storage_bytes(129, 3)) + + def test_rejects_invalid_storage_bits(self): + with self.assertRaises(ValueError): + MODULE.sptag_rabitq_storage_bytes(128, 0) + + def test_load_ground_truth_honors_query_count_and_topk(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'truth.txt' + path.write_text('1 2 3\n4 5 6\n', encoding='ascii') + self.assertEqual([{1, 2}, {4, 5}], MODULE.load_ground_truth(path, 2, 2)) + with self.assertRaises(ValueError): + MODULE.load_ground_truth(path, 3, 2) + + def test_recall_is_query_equal_weighted(self): + index = FakeIndex([[1, 9], [4, 5]]) + queries = np.zeros((2, 4), dtype=np.float32) + recall = MODULE.recall_at_k(index, queries, [{1, 2}, {3, 4}], 2, batch_size=2) + self.assertEqual(0.5, recall) + + +if __name__ == '__main__': + unittest.main() From 66d1c297a9a75f1a5da7050dea4a0f09a59255f8 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:14:56 +0000 Subject: [PATCH 03/17] Remove RaBitQ tuning leftovers Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Tools/OPQ/OPQ_gpu_train_infer.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index eb3750ffb..d239fd704 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -3,7 +3,6 @@ import tqdm import time from struct import pack, unpack, calcsize -from struct import pack, unpack, calcsize from typing import Dict, List import heapq import argparse @@ -380,7 +379,7 @@ def tune_rabitq_bits(args, faiss, training_data, queries, ground_truths): trials.append({'bits': bits, 'recall': recall}) print(f'RaBitQ{bits} Recall@{args.k}: {recall:.6f}') if recall >= args.rabitq_target_recall: - return bits, candidate, data_count, trials + return bits, data_count, trials measured = ', '.join(f'{trial["bits"]}-bit={trial["recall"]:.6f}' for trial in trials) raise RuntimeError( @@ -412,8 +411,6 @@ def train_rabitq(args): raise ValueError('RaBitQ training data is empty') print(f'train RaBitQ using {num_training} samples ...') faiss.omp_set_num_threads(args.T) - trials = [] - data_count = None if args.rabitq_auto_tune: if args.quan_test <= 0: raise ValueError('rabitq_auto_tune requires quan_test > 0 and a pre-generated ground truth') @@ -424,7 +421,7 @@ def train_rabitq(args): if num_query != args.Q: raise ValueError(f'query file contains {num_query} queries, but configured Q is {args.Q}') ground_truths = load_ground_truth(args.output_truth, num_query, args.k) - bits, faiss_index, data_count, trials = tune_rabitq_bits( + bits, data_count, trials = tune_rabitq_bits( args, faiss, training_data, queries, ground_truths) result = { 'selected_bits': bits, @@ -447,7 +444,7 @@ def train_rabitq(args): bits = rabitq_bits_from_quantized_dimension(args.dim, args.quan_dim) faiss_index = create_rabitq_index(faiss, args.dim, bits, training_data) if args.quan_test > 0: - data_count = add_rabitq_data(args, faiss_index) + add_rabitq_data(args, faiss_index) if args.quan_test > 0 and not args.rabitq_auto_tune: queryreader = DataReader(args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) From a3528f9d0a80dc138573fe98bd451b7a6df5c038 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:16:24 +0000 Subject: [PATCH 04/17] Document RaBitQ auto tuning workflow Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/GettingStart.md | 47 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 46 insertions(+), 1 deletion(-) diff --git a/docs/GettingStart.md b/docs/GettingStart.md index a743fafb6..a109b2ccb 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -247,7 +247,52 @@ use the generated model through `QuantizerFilePath` in the normal SPANN workflow. For 128-dimensional SIFT vectors, the encoded `UInt8` vectors use `Dim=68` at 3 bits (48 compact code bytes plus five Float factors). -Train and encode SIFT1M: +To select the minimum storage bit count before building the index, run the +RaBitQ tuner against a pre-generated exact top-`k` ground truth. The tuning +inputs below use SPTAG `DEFAULT` binary files (`int32 count`, `int32 dimension`, +then vector payload) and a text ground-truth file containing one +space-separated neighbor-ID list per query: + +```bash +python3 Tools/OPQ/OPQ_gpu_train_infer.py \ + --data_file sift1m/sift_base.bin \ + --query_file sift1m/sift_query.bin \ + --output_truth sift1m/sift_groundtruth_top1000.txt \ + --output_dir sift1m/rabitq_tuning \ + --data_type float32 --target_type float32 \ + --dim 128 --B 1000000 --Q 10000 --k 1000 --D L2 --T 46 \ + --train_samples 1000000 --quan_type rabitq --quan_test 1 \ + --rabitq_auto_tune --rabitq_target_recall 0.95 \ + --rabitq_min_bits 1 --rabitq_max_bits 8 +``` + +The tuner evaluates bit counts in ascending order using exactly the configured +10,000 queries and selects the first bit count meeting `Recall@1000 >= 0.95`. +It fails if fewer queries/ground-truth rows are available or no candidate +qualifies. The ground truth is not moved or modified. Results are written +atomically to `sift1m/rabitq_tuning/rabitq_auto_tuning.json`; use +`native_quantizer_qd` for `-qd` and `storage_bytes_per_vector` for the SPANN +`[Base] Dim`: + +```bash +BITS=$(python3 -c \ + 'import json; print(json.load(open("sift1m/rabitq_tuning/rabitq_auto_tuning.json"))["native_quantizer_qd"])') +STORAGE_DIM=$(python3 -c \ + 'import json; print(json.load(open("sift1m/rabitq_tuning/rabitq_auto_tuning.json"))["storage_bytes_per_vector"])') + +Release/quantizer \ + -d 128 -v Float -f XVEC \ + -i sift1m/sift_base.fvecs \ + -o "sift1m/sift_base.rabitq${BITS}.u8bin" \ + -oq "sift1m/official_rabitq${BITS}_global.bin" \ + -qt RaBitQQuantizer -qd "$BITS" -ts 1000000 +``` + +Set `VectorPath` and `QuantizerFilePath` to those generated files and set +`Dim=$STORAGE_DIM` in the build INI before starting SPANN construction. Do not +use the manually supplied `quan_dim` to override an auto-tuning result. + +For reference, the fixed 3-bit SIFT1M command is: ```bash Release/quantizer \ From 6a383851c718d3ae141e6d3b5f84a8f8f28a6e76 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:26:37 +0000 Subject: [PATCH 05/17] Add authoritative INI for RaBitQ tuning Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Script_AE/iniFile/rabitq_auto_tune_sift1m.ini | 22 +++++ Tools/OPQ/OPQ_gpu_train_infer.py | 93 ++++++++++++++++++- Tools/OPQ/README.md | 14 ++- Tools/OPQ/test_rabitq_auto_tuning.py | 69 ++++++++++++++ docs/GettingStart.md | 54 ++++++++--- 5 files changed, 233 insertions(+), 19 deletions(-) create mode 100644 Script_AE/iniFile/rabitq_auto_tune_sift1m.ini diff --git a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini new file mode 100644 index 000000000..ee07e1971 --- /dev/null +++ b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini @@ -0,0 +1,22 @@ +[RaBitQAutoTune] +isExecute=true +DataFile=sift1m/sift_base.bin +QueryFile=sift1m/sift_query.bin +TruthFile=sift1m/sift_groundtruth_top1000.txt +OutputDir=sift1m/rabitq_tuning +DataType=float32 +TargetType=float32 +Dimension=128 +DataBatchSize=1000000 +RecallAt=1000 +Distance=L2 +Threads=46 +TrainingSamples=1000000 +TargetRecall=0.95 +MinBits=1 +MaxBits=8 +DataNormalize=0 +QueryNormalize=0 + +[SearchSSDIndex] +QueryCountLimit=10000 diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index d239fd704..710803d94 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -7,13 +7,16 @@ import heapq import argparse import copy +import configparser import json from operator import itemgetter import os import subprocess +import sys -def get_config(): +def get_cli_parser(): parser = argparse.ArgumentParser(description ='implementation of nnsearch.') + parser.add_argument('--config', type = str, help='native INI containing [RaBitQAutoTune]') parser.add_argument('--data_file', default = 'traindata', type = str, help = 'binary data file') parser.add_argument('--query_file', default = 'query.tsv', type= str, help='query tsv file') parser.add_argument('--data_normalize', default = 0, type = int, help='normalize data vectors') @@ -45,8 +48,92 @@ def get_config(): parser.add_argument('--rabitq_min_bits', type = int, default = 1, help='minimum RaBitQ storage bits to evaluate') parser.add_argument('--rabitq_max_bits', type = int, default = 8, help='maximum RaBitQ storage bits to evaluate') parser.add_argument('--rabitq_tuning_result', type = str, default = 'rabitq_auto_tuning.json', help='auto-tuning result file under output_dir') - args = parser.parse_args() - return args + return parser + +def load_rabitq_auto_tune_ini(path): + config = configparser.ConfigParser(interpolation=None) + if not config.read(path): + raise ValueError(f'cannot read INI file: {path}') + if config.defaults(): + raise ValueError('[DEFAULT] parameters are forbidden in strict RaBitQ INI mode') + section_name = 'RaBitQAutoTune' + if not config.has_section(section_name): + raise ValueError(f'INI file is missing [{section_name}]') + section = config[section_name] + allowed_keys = { + 'isexecute', 'datafile', 'queryfile', 'truthfile', 'outputdir', + 'datatype', 'targettype', 'dimension', 'databatchsize', 'querycount', + 'recallat', 'distance', 'threads', 'trainingsamples', 'targetrecall', + 'minbits', 'maxbits', 'tuningresult', 'datanormalize', 'querynormalize', + 'dataformat', 'task', + } + unknown_keys = set(section.keys()) - allowed_keys + if unknown_keys: + raise ValueError( + f'unknown [{section_name}] parameter(s): {", ".join(sorted(unknown_keys))}') + required_keys = ('DataFile', 'QueryFile', 'TruthFile', 'OutputDir', 'Dimension', 'RecallAt') + missing_keys = [key for key in required_keys if not section.get(key)] + if missing_keys: + raise ValueError( + f'[{section_name}] is missing required parameter(s): {", ".join(missing_keys)}') + if not section.getboolean('isExecute', fallback=False): + raise ValueError(f'[{section_name}] isExecute must be true') + + query_count = section.getint('QueryCount', fallback=None) + if query_count is None: + if not config.has_option('SearchSSDIndex', 'QueryCountLimit'): + raise ValueError( + f'query count must be set by [{section_name}] QueryCount or ' + '[SearchSSDIndex] QueryCountLimit') + query_count = config.getint('SearchSSDIndex', 'QueryCountLimit') + + return argparse.Namespace( + config=path, + data_file=section['DataFile'], + query_file=section['QueryFile'], + data_normalize=section.getint('DataNormalize', fallback=0), + query_normalize=section.getint('QueryNormalize', fallback=0), + data_type=section.get('DataType', fallback='float32'), + target_type=section.get('TargetType', fallback='float32'), + k=section.getint('RecallAt'), + dim=section.getint('Dimension'), + B=section.getint('DataBatchSize', fallback=-1), + Q=query_count, + S=1000, + D=section.get('Distance', fallback='L2'), + output_truth=section['TruthFile'], + data_format=section.get('DataFormat', fallback='DEFAULT'), + task=section.getint('Task', fallback=0), + log_dir='', + T=section.getint('Threads', fallback=32), + train_samples=section.getint('TrainingSamples', fallback=1000000), + quan_type='rabitq', + quan_dim=-1, + output_dir=section['OutputDir'], + output_quantizer='quantizer.bin', + output_quan_vector_file='', + output_rec_vector_file='', + quan_test=1, + rabitq_auto_tune=True, + rabitq_target_recall=section.getfloat('TargetRecall', fallback=0.95), + rabitq_min_bits=section.getint('MinBits', fallback=1), + rabitq_max_bits=section.getint('MaxBits', fallback=8), + rabitq_tuning_result=section.get( + 'TuningResult', fallback='rabitq_auto_tuning.json'), + ) + +def get_config(argv=None): + if argv is None: + argv = sys.argv[1:] + config_probe = argparse.ArgumentParser(add_help=False) + config_probe.add_argument('--config') + config_args, remaining = config_probe.parse_known_args(argv) + if config_args.config is not None: + if remaining: + raise ValueError( + 'INI mode accepts only --config; command-line parameter overrides are forbidden') + return load_rabitq_auto_tune_ini(config_args.config) + return get_cli_parser().parse_args(argv) def is_binary_vector_file(filename): suffixes = ('.bin', '.fbin', '.u8bin', '.i8bin') diff --git a/Tools/OPQ/README.md b/Tools/OPQ/README.md index bf5cc3145..7796f440f 100644 --- a/Tools/OPQ/README.md +++ b/Tools/OPQ/README.md @@ -31,4 +31,16 @@ factors), measured recalls, and sample counts are written atomically to `/rabitq_auto_tuning.json`. Use `native_quantizer_qd` as `-qd` for the native SPTAG `quantizer` command that generates the quantizer and vectors consumed by the subsequent index build. The tuner deliberately does not emit Faiss codes: -their persisted layout is not the SPTAG global RaBitQ format. \ No newline at end of file +their persisted layout is not the SPTAG global RaBitQ format. + +The same parameters can be supplied exclusively through the +`[RaBitQAutoTune]` section of an INI: + +```bash +python3 Tools/OPQ/OPQ_gpu_train_infer.py \ + --config Script_AE/iniFile/rabitq_auto_tune_sift1m.ini +``` + +In INI mode no additional CLI parameters are accepted. `QueryCount` defaults to +`[SearchSSDIndex] QueryCountLimit`, keeping the configured query count in one +place. \ No newline at end of file diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py index f8469527c..d6b812fba 100644 --- a/Tools/OPQ/test_rabitq_auto_tuning.py +++ b/Tools/OPQ/test_rabitq_auto_tuning.py @@ -64,6 +64,75 @@ def test_recall_is_query_equal_weighted(self): recall = MODULE.recall_at_k(index, queries, [{1, 2}, {3, 4}], 2, batch_size=2) self.assertEqual(0.5, recall) + def test_ini_is_authoritative_and_reuses_query_count_limit(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'build.ini' + path.write_text( + '[RaBitQAutoTune]\n' + 'isExecute=true\n' + 'DataFile=base.bin\n' + 'QueryFile=query.bin\n' + 'TruthFile=truth.txt\n' + 'OutputDir=tuning\n' + 'Dimension=128\n' + 'RecallAt=1000\n' + 'TargetRecall=0.97\n' + 'MinBits=2\n' + 'MaxBits=7\n' + '\n' + '[SearchSSDIndex]\n' + 'QueryCountLimit=10000\n', + encoding='ascii') + args = MODULE.get_config(['--config', str(path)]) + self.assertTrue(args.rabitq_auto_tune) + self.assertEqual(10000, args.Q) + self.assertEqual(1000, args.k) + self.assertEqual(0.97, args.rabitq_target_recall) + self.assertEqual((2, 7), (args.rabitq_min_bits, args.rabitq_max_bits)) + + with self.assertRaises(ValueError): + MODULE.get_config(['--config', str(path), '--Q', '1']) + + def test_ini_rejects_unknown_parameters(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'build.ini' + path.write_text( + '[RaBitQAutoTune]\n' + 'isExecute=true\n' + 'DataFile=base.bin\n' + 'QueryFile=query.bin\n' + 'TruthFile=truth.txt\n' + 'OutputDir=tuning\n' + 'Dimension=128\n' + 'RecallAt=1000\n' + 'QueryCount=10000\n' + 'TypoTargetRecal=0.95\n', + encoding='ascii') + with self.assertRaises(ValueError): + MODULE.get_config(['--config', str(path)]) + + def test_ini_rejects_inherited_defaults(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'build.ini' + path.write_text( + '[DEFAULT]\n' + 'QueryCount=1\n' + '\n' + '[RaBitQAutoTune]\n' + 'isExecute=true\n' + 'DataFile=base.bin\n' + 'QueryFile=query.bin\n' + 'TruthFile=truth.txt\n' + 'OutputDir=tuning\n' + 'Dimension=128\n' + 'RecallAt=1000\n' + '\n' + '[SearchSSDIndex]\n' + 'QueryCountLimit=10000\n', + encoding='ascii') + with self.assertRaises(ValueError): + MODULE.get_config(['--config', str(path)]) + if __name__ == '__main__': unittest.main() diff --git a/docs/GettingStart.md b/docs/GettingStart.md index a109b2ccb..16c06e3d4 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -247,27 +247,51 @@ use the generated model through `QuantizerFilePath` in the normal SPANN workflow. For 128-dimensional SIFT vectors, the encoded `UInt8` vectors use `Dim=68` at 3 bits (48 compact code bytes plus five Float factors). -To select the minimum storage bit count before building the index, run the -RaBitQ tuner against a pre-generated exact top-`k` ground truth. The tuning -inputs below use SPTAG `DEFAULT` binary files (`int32 count`, `int32 dimension`, -then vector payload) and a text ground-truth file containing one -space-separated neighbor-ID list per query: +To select the minimum storage bit count before building the index, add a +`[RaBitQAutoTune]` section to the SPANN INI. The tuner uses a pre-generated +exact top-`k` ground truth. Its vector inputs use SPTAG `DEFAULT` binary files +(`int32 count`, `int32 dimension`, then vector payload), and its text +ground-truth file contains one space-separated neighbor-ID list per query: + +```ini +[RaBitQAutoTune] +isExecute=true +DataFile=sift1m/sift_base.bin +QueryFile=sift1m/sift_query.bin +TruthFile=sift1m/sift_groundtruth_top1000.txt +OutputDir=sift1m/rabitq_tuning +DataType=float32 +TargetType=float32 +Dimension=128 +DataBatchSize=1000000 +RecallAt=1000 +Distance=L2 +Threads=46 +TrainingSamples=1000000 +TargetRecall=0.95 +MinBits=1 +MaxBits=8 +DataNormalize=0 +QueryNormalize=0 + +[SearchSSDIndex] +QueryCountLimit=10000 +``` + +Run the pre-build tuning stage with only the INI path: ```bash python3 Tools/OPQ/OPQ_gpu_train_infer.py \ - --data_file sift1m/sift_base.bin \ - --query_file sift1m/sift_query.bin \ - --output_truth sift1m/sift_groundtruth_top1000.txt \ - --output_dir sift1m/rabitq_tuning \ - --data_type float32 --target_type float32 \ - --dim 128 --B 1000000 --Q 10000 --k 1000 --D L2 --T 46 \ - --train_samples 1000000 --quan_type rabitq --quan_test 1 \ - --rabitq_auto_tune --rabitq_target_recall 0.95 \ - --rabitq_min_bits 1 --rabitq_max_bits 8 + --config Script_AE/iniFile/rabitq_auto_tune_sift1m.ini ``` The tuner evaluates bit counts in ascending order using exactly the configured -10,000 queries and selects the first bit count meeting `Recall@1000 >= 0.95`. +`SearchSSDIndex.QueryCountLimit` queries and selects the first bit count meeting +`Recall@1000 >= 0.95`. When `--config` is used, additional command-line +parameters are rejected, so they cannot override the INI. A +`RaBitQAutoTune.QueryCount` may be set explicitly only when the tuning query +count intentionally differs from `SearchSSDIndex.QueryCountLimit`. Strict mode +also rejects unknown `[RaBitQAutoTune]` keys and inherited `[DEFAULT]` values. It fails if fewer queries/ground-truth rows are available or no candidate qualifies. The ground truth is not moved or modified. Results are written atomically to `sift1m/rabitq_tuning/rabitq_auto_tuning.json`; use From c39964b8d7730fb6faf5f8962923c5ecc4b536c6 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:31:25 +0000 Subject: [PATCH 06/17] Use Base as RaBitQ tuning data source Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Script_AE/iniFile/rabitq_auto_tune_sift1m.ini | 17 +++++--- Tools/OPQ/OPQ_gpu_train_infer.py | 37 +++++++++++++----- Tools/OPQ/test_rabitq_auto_tuning.py | 39 +++++++++++++------ docs/GettingStart.md | 35 ++++++++++------- 4 files changed, 87 insertions(+), 41 deletions(-) diff --git a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini index ee07e1971..1a39f0212 100644 --- a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini +++ b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini @@ -1,15 +1,20 @@ +[Base] +ValueType=Float +DistCalcMethod=L2 +Dim=128 +VectorPath=sift1m/sift_base.bin +VectorType=DEFAULT +QueryPath=sift1m/sift_query.bin +QueryType=DEFAULT +TruthPath=sift1m/sift_groundtruth_top1000.txt +TruthType=DEFAULT + [RaBitQAutoTune] isExecute=true -DataFile=sift1m/sift_base.bin -QueryFile=sift1m/sift_query.bin -TruthFile=sift1m/sift_groundtruth_top1000.txt OutputDir=sift1m/rabitq_tuning -DataType=float32 TargetType=float32 -Dimension=128 DataBatchSize=1000000 RecallAt=1000 -Distance=L2 Threads=46 TrainingSamples=1000000 TargetRecall=0.95 diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index 710803d94..5ccac1d18 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -59,11 +59,12 @@ def load_rabitq_auto_tune_ini(path): section_name = 'RaBitQAutoTune' if not config.has_section(section_name): raise ValueError(f'INI file is missing [{section_name}]') + if not config.has_section('Base'): + raise ValueError('INI file is missing [Base]') section = config[section_name] allowed_keys = { - 'isexecute', 'datafile', 'queryfile', 'truthfile', 'outputdir', - 'datatype', 'targettype', 'dimension', 'databatchsize', 'querycount', - 'recallat', 'distance', 'threads', 'trainingsamples', 'targetrecall', + 'isexecute', 'outputdir', 'targettype', 'databatchsize', 'querycount', + 'recallat', 'threads', 'trainingsamples', 'targetrecall', 'minbits', 'maxbits', 'tuningresult', 'datanormalize', 'querynormalize', 'dataformat', 'task', } @@ -71,7 +72,7 @@ def load_rabitq_auto_tune_ini(path): if unknown_keys: raise ValueError( f'unknown [{section_name}] parameter(s): {", ".join(sorted(unknown_keys))}') - required_keys = ('DataFile', 'QueryFile', 'TruthFile', 'OutputDir', 'Dimension', 'RecallAt') + required_keys = ('OutputDir', 'RecallAt') missing_keys = [key for key in required_keys if not section.get(key)] if missing_keys: raise ValueError( @@ -79,6 +80,22 @@ def load_rabitq_auto_tune_ini(path): if not section.getboolean('isExecute', fallback=False): raise ValueError(f'[{section_name}] isExecute must be true') + base = config['Base'] + base_required_keys = ('VectorPath', 'QueryPath', 'TruthPath', 'Dim', 'ValueType', 'DistCalcMethod') + missing_base_keys = [key for key in base_required_keys if not base.get(key)] + if missing_base_keys: + raise ValueError( + f'[Base] is missing required parameter(s): {", ".join(missing_base_keys)}') + value_types = { + 'float': 'float32', + 'int8': 'int8', + 'uint8': 'uint8', + 'int16': 'int16', + } + value_type = base['ValueType'].lower() + if value_type not in value_types: + raise ValueError(f'unsupported [Base] ValueType for RaBitQ tuning: {base["ValueType"]}') + query_count = section.getint('QueryCount', fallback=None) if query_count is None: if not config.has_option('SearchSSDIndex', 'QueryCountLimit'): @@ -89,19 +106,19 @@ def load_rabitq_auto_tune_ini(path): return argparse.Namespace( config=path, - data_file=section['DataFile'], - query_file=section['QueryFile'], + data_file=base['VectorPath'], + query_file=base['QueryPath'], data_normalize=section.getint('DataNormalize', fallback=0), query_normalize=section.getint('QueryNormalize', fallback=0), - data_type=section.get('DataType', fallback='float32'), + data_type=value_types[value_type], target_type=section.get('TargetType', fallback='float32'), k=section.getint('RecallAt'), - dim=section.getint('Dimension'), + dim=base.getint('Dim'), B=section.getint('DataBatchSize', fallback=-1), Q=query_count, S=1000, - D=section.get('Distance', fallback='L2'), - output_truth=section['TruthFile'], + D=base['DistCalcMethod'], + output_truth=base['TruthPath'], data_format=section.get('DataFormat', fallback='DEFAULT'), task=section.getint('Task', fallback=0), log_dir='', diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py index d6b812fba..a513fb6a4 100644 --- a/Tools/OPQ/test_rabitq_auto_tuning.py +++ b/Tools/OPQ/test_rabitq_auto_tuning.py @@ -68,13 +68,17 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / 'build.ini' path.write_text( + '[Base]\n' + 'ValueType=Float\n' + 'DistCalcMethod=L2\n' + 'Dim=128\n' + 'VectorPath=base.bin\n' + 'QueryPath=query.bin\n' + 'TruthPath=truth.txt\n' + '\n' '[RaBitQAutoTune]\n' 'isExecute=true\n' - 'DataFile=base.bin\n' - 'QueryFile=query.bin\n' - 'TruthFile=truth.txt\n' 'OutputDir=tuning\n' - 'Dimension=128\n' 'RecallAt=1000\n' 'TargetRecall=0.97\n' 'MinBits=2\n' @@ -89,6 +93,9 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): self.assertEqual(1000, args.k) self.assertEqual(0.97, args.rabitq_target_recall) self.assertEqual((2, 7), (args.rabitq_min_bits, args.rabitq_max_bits)) + self.assertEqual('base.bin', args.data_file) + self.assertEqual('query.bin', args.query_file) + self.assertEqual('truth.txt', args.output_truth) with self.assertRaises(ValueError): MODULE.get_config(['--config', str(path), '--Q', '1']) @@ -97,13 +104,17 @@ def test_ini_rejects_unknown_parameters(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / 'build.ini' path.write_text( + '[Base]\n' + 'ValueType=Float\n' + 'DistCalcMethod=L2\n' + 'Dim=128\n' + 'VectorPath=base.bin\n' + 'QueryPath=query.bin\n' + 'TruthPath=truth.txt\n' + '\n' '[RaBitQAutoTune]\n' 'isExecute=true\n' - 'DataFile=base.bin\n' - 'QueryFile=query.bin\n' - 'TruthFile=truth.txt\n' 'OutputDir=tuning\n' - 'Dimension=128\n' 'RecallAt=1000\n' 'QueryCount=10000\n' 'TypoTargetRecal=0.95\n', @@ -118,13 +129,17 @@ def test_ini_rejects_inherited_defaults(self): '[DEFAULT]\n' 'QueryCount=1\n' '\n' + '[Base]\n' + 'ValueType=Float\n' + 'DistCalcMethod=L2\n' + 'Dim=128\n' + 'VectorPath=base.bin\n' + 'QueryPath=query.bin\n' + 'TruthPath=truth.txt\n' + '\n' '[RaBitQAutoTune]\n' 'isExecute=true\n' - 'DataFile=base.bin\n' - 'QueryFile=query.bin\n' - 'TruthFile=truth.txt\n' 'OutputDir=tuning\n' - 'Dimension=128\n' 'RecallAt=1000\n' '\n' '[SearchSSDIndex]\n' diff --git a/docs/GettingStart.md b/docs/GettingStart.md index 16c06e3d4..f8c97a25e 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -248,24 +248,31 @@ workflow. For 128-dimensional SIFT vectors, the encoded `UInt8` vectors use `Dim=68` at 3 bits (48 compact code bytes plus five Float factors). To select the minimum storage bit count before building the index, add a -`[RaBitQAutoTune]` section to the SPANN INI. The tuner uses a pre-generated -exact top-`k` ground truth. Its vector inputs use SPTAG `DEFAULT` binary files -(`int32 count`, `int32 dimension`, then vector payload), and its text -ground-truth file contains one space-separated neighbor-ID list per query: +`[RaBitQAutoTune]` section to a pre-build INI. `[Base]` is the single source for +the raw vector, query, ground-truth, dimension, value type, and distance. The +tuner uses a pre-generated exact top-`k` ground truth. Vector inputs use SPTAG +`DEFAULT` binary files (`int32 count`, `int32 dimension`, then vector payload), +and the text ground-truth file contains one space-separated neighbor-ID list +per query: ```ini +[Base] +ValueType=Float +DistCalcMethod=L2 +Dim=128 +VectorPath=sift1m/sift_base.bin +VectorType=DEFAULT +QueryPath=sift1m/sift_query.bin +QueryType=DEFAULT +TruthPath=sift1m/sift_groundtruth_top1000.txt +TruthType=DEFAULT + [RaBitQAutoTune] isExecute=true -DataFile=sift1m/sift_base.bin -QueryFile=sift1m/sift_query.bin -TruthFile=sift1m/sift_groundtruth_top1000.txt OutputDir=sift1m/rabitq_tuning -DataType=float32 TargetType=float32 -Dimension=128 DataBatchSize=1000000 RecallAt=1000 -Distance=L2 Threads=46 TrainingSamples=1000000 TargetRecall=0.95 @@ -312,9 +319,11 @@ Release/quantizer \ -qt RaBitQQuantizer -qd "$BITS" -ts 1000000 ``` -Set `VectorPath` and `QuantizerFilePath` to those generated files and set -`Dim=$STORAGE_DIM` in the build INI before starting SPANN construction. Do not -use the manually supplied `quan_dim` to override an auto-tuning result. +The pre-build INI continues to describe raw Float data. Set `VectorPath` and +`QuantizerFilePath` to the generated quantized files and set +`ValueType=UInt8`, `Dim=$STORAGE_DIM` in the subsequent build INI before +starting SPANN construction. Do not use a manually supplied bit count to +override an auto-tuning result. For reference, the fixed 3-bit SIFT1M command is: From bb240752845c2bc6a0ce04c0d209e0829bbdf281 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:37:52 +0000 Subject: [PATCH 07/17] Infer RaBitQ recall cutoff from truth Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Script_AE/iniFile/rabitq_auto_tune_sift1m.ini | 1 - Tools/OPQ/OPQ_gpu_train_infer.py | 34 ++++++++++++++----- Tools/OPQ/test_rabitq_auto_tuning.py | 22 +++++++++--- docs/GettingStart.md | 8 +++-- 4 files changed, 47 insertions(+), 18 deletions(-) diff --git a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini index 1a39f0212..fd4d84839 100644 --- a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini +++ b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini @@ -14,7 +14,6 @@ isExecute=true OutputDir=sift1m/rabitq_tuning TargetType=float32 DataBatchSize=1000000 -RecallAt=1000 Threads=46 TrainingSamples=1000000 TargetRecall=0.95 diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index 5ccac1d18..16325ffd5 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -64,7 +64,7 @@ def load_rabitq_auto_tune_ini(path): section = config[section_name] allowed_keys = { 'isexecute', 'outputdir', 'targettype', 'databatchsize', 'querycount', - 'recallat', 'threads', 'trainingsamples', 'targetrecall', + 'threads', 'trainingsamples', 'targetrecall', 'minbits', 'maxbits', 'tuningresult', 'datanormalize', 'querynormalize', 'dataformat', 'task', } @@ -72,7 +72,7 @@ def load_rabitq_auto_tune_ini(path): if unknown_keys: raise ValueError( f'unknown [{section_name}] parameter(s): {", ".join(sorted(unknown_keys))}') - required_keys = ('OutputDir', 'RecallAt') + required_keys = ('OutputDir',) missing_keys = [key for key in required_keys if not section.get(key)] if missing_keys: raise ValueError( @@ -112,7 +112,7 @@ def load_rabitq_auto_tune_ini(path): query_normalize=section.getint('QueryNormalize', fallback=0), data_type=value_types[value_type], target_type=section.get('TargetType', fallback='float32'), - k=section.getint('RecallAt'), + k=None, dim=base.getint('Dim'), B=section.getint('DataBatchSize', fallback=-1), Q=query_count, @@ -420,19 +420,30 @@ def rabitq_bits_from_quantized_dimension(dim, quan_dim): raise ValueError(f'RaBitQ bits must be in [1, 8], got {bits}') return bits -def load_ground_truth(path, query_count, topk): +def load_ground_truth(path, query_count, topk=None): ground_truths = [] + infer_topk = topk is None with open(path, 'r') as truth_file: for query_id in range(query_count): line = truth_file.readline() if not line: raise ValueError(f'ground truth contains only {query_id} queries, expected {query_count}') neighbors = line.strip().split() + if topk is None: + topk = len(neighbors) + if topk == 0: + raise ValueError('ground truth query 0 contains no neighbors') if len(neighbors) < topk: raise ValueError( f'ground truth query {query_id} contains {len(neighbors)} neighbors, expected at least {topk}') - ground_truths.append(set(int(neighbor) for neighbor in neighbors[:topk])) - return ground_truths + if infer_topk and len(neighbors) != topk: + raise ValueError( + f'ground truth query {query_id} contains {len(neighbors)} neighbors, expected exactly {topk}') + truth = set(int(neighbor) for neighbor in neighbors[:topk]) + if len(truth) != topk: + raise ValueError(f'ground truth query {query_id} contains duplicate neighbor IDs') + ground_truths.append(truth) + return ground_truths, topk def recall_at_k(faiss_index, queries, ground_truths, topk, batch_size=64): if len(queries) != len(ground_truths): @@ -501,8 +512,10 @@ def train_rabitq(args): raise ValueError('Cosine RaBitQ tuning requires data_normalize=1 and query_normalize=1') if args.train_samples <= 0: raise ValueError('train_samples must be positive') - if args.k <= 0 or args.Q <= 0: - raise ValueError('k and Q must be positive') + if args.Q <= 0: + raise ValueError('Q must be positive') + if not args.rabitq_auto_tune and (args.k is None or args.k <= 0): + raise ValueError('k must be positive') if len(args.output_quan_vector_file) > 0 or len(args.output_rec_vector_file) > 0: raise ValueError( 'RaBitQ vectors must be generated by the native SPTAG quantizer after tuning; ' @@ -524,7 +537,10 @@ def train_rabitq(args): queryreader.close() if num_query != args.Q: raise ValueError(f'query file contains {num_query} queries, but configured Q is {args.Q}') - ground_truths = load_ground_truth(args.output_truth, num_query, args.k) + ground_truths, truth_topk = load_ground_truth( + args.output_truth, num_query, args.k) + if args.k is None: + args.k = truth_topk bits, data_count, trials = tune_rabitq_bits( args, faiss, training_data, queries, ground_truths) result = { diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py index a513fb6a4..09b78d923 100644 --- a/Tools/OPQ/test_rabitq_auto_tuning.py +++ b/Tools/OPQ/test_rabitq_auto_tuning.py @@ -54,10 +54,25 @@ def test_load_ground_truth_honors_query_count_and_topk(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / 'truth.txt' path.write_text('1 2 3\n4 5 6\n', encoding='ascii') - self.assertEqual([{1, 2}, {4, 5}], MODULE.load_ground_truth(path, 2, 2)) + truths, topk = MODULE.load_ground_truth(path, 2) + self.assertEqual(3, topk) + self.assertEqual([{1, 2, 3}, {4, 5, 6}], truths) + truths, topk = MODULE.load_ground_truth(path, 2, 2) + self.assertEqual(2, topk) + self.assertEqual([{1, 2}, {4, 5}], truths) with self.assertRaises(ValueError): MODULE.load_ground_truth(path, 3, 2) + def test_inferred_ground_truth_topk_requires_uniform_unique_rows(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'truth.txt' + path.write_text('1 2\n3 4 5\n', encoding='ascii') + with self.assertRaises(ValueError): + MODULE.load_ground_truth(path, 2) + path.write_text('1 2\n3 3\n', encoding='ascii') + with self.assertRaises(ValueError): + MODULE.load_ground_truth(path, 2) + def test_recall_is_query_equal_weighted(self): index = FakeIndex([[1, 9], [4, 5]]) queries = np.zeros((2, 4), dtype=np.float32) @@ -79,7 +94,6 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): '[RaBitQAutoTune]\n' 'isExecute=true\n' 'OutputDir=tuning\n' - 'RecallAt=1000\n' 'TargetRecall=0.97\n' 'MinBits=2\n' 'MaxBits=7\n' @@ -90,7 +104,7 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): args = MODULE.get_config(['--config', str(path)]) self.assertTrue(args.rabitq_auto_tune) self.assertEqual(10000, args.Q) - self.assertEqual(1000, args.k) + self.assertIsNone(args.k) self.assertEqual(0.97, args.rabitq_target_recall) self.assertEqual((2, 7), (args.rabitq_min_bits, args.rabitq_max_bits)) self.assertEqual('base.bin', args.data_file) @@ -115,7 +129,6 @@ def test_ini_rejects_unknown_parameters(self): '[RaBitQAutoTune]\n' 'isExecute=true\n' 'OutputDir=tuning\n' - 'RecallAt=1000\n' 'QueryCount=10000\n' 'TypoTargetRecal=0.95\n', encoding='ascii') @@ -140,7 +153,6 @@ def test_ini_rejects_inherited_defaults(self): '[RaBitQAutoTune]\n' 'isExecute=true\n' 'OutputDir=tuning\n' - 'RecallAt=1000\n' '\n' '[SearchSSDIndex]\n' 'QueryCountLimit=10000\n', diff --git a/docs/GettingStart.md b/docs/GettingStart.md index f8c97a25e..9f294b496 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -272,7 +272,6 @@ isExecute=true OutputDir=sift1m/rabitq_tuning TargetType=float32 DataBatchSize=1000000 -RecallAt=1000 Threads=46 TrainingSamples=1000000 TargetRecall=0.95 @@ -294,8 +293,11 @@ python3 Tools/OPQ/OPQ_gpu_train_infer.py \ The tuner evaluates bit counts in ascending order using exactly the configured `SearchSSDIndex.QueryCountLimit` queries and selects the first bit count meeting -`Recall@1000 >= 0.95`. When `--config` is used, additional command-line -parameters are rejected, so they cannot override the INI. A +the target Recall. The Recall cutoff is inferred from the exact ground-truth +width: 1,000 neighbor IDs per query means `Recall@1000`, so there is no second +top-k setting. All consumed ground-truth rows must have exactly that width with +no duplicate IDs. When `--config` is used, additional command-line parameters +are rejected, so they cannot override the INI. A `RaBitQAutoTune.QueryCount` may be set explicitly only when the tuning query count intentionally differs from `SearchSSDIndex.QueryCountLimit`. Strict mode also rejects unknown `[RaBitQAutoTune]` keys and inherited `[DEFAULT]` values. From 43d990399907000286fe944459aac2bb381a4fb6 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:59:08 +0000 Subject: [PATCH 08/17] Restore RaBitQ reranking recall semantics Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Script_AE/iniFile/rabitq_auto_tune_sift1m.ini | 1 + Tools/OPQ/OPQ_gpu_train_infer.py | 81 +++++++++++-------- Tools/OPQ/README.md | 6 +- Tools/OPQ/test_rabitq_auto_tuning.py | 47 +++++------ docs/GettingStart.md | 15 ++-- 5 files changed, 82 insertions(+), 68 deletions(-) diff --git a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini index fd4d84839..c17e40e50 100644 --- a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini +++ b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini @@ -24,3 +24,4 @@ QueryNormalize=0 [SearchSSDIndex] QueryCountLimit=10000 +ResultNum=100 diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index 16325ffd5..d624ae0ed 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -103,6 +103,8 @@ def load_rabitq_auto_tune_ini(path): f'query count must be set by [{section_name}] QueryCount or ' '[SearchSSDIndex] QueryCountLimit') query_count = config.getint('SearchSSDIndex', 'QueryCountLimit') + if not config.has_option('SearchSSDIndex', 'ResultNum'): + raise ValueError('[SearchSSDIndex] ResultNum is required for RaBitQ tuning') return argparse.Namespace( config=path, @@ -112,7 +114,7 @@ def load_rabitq_auto_tune_ini(path): query_normalize=section.getint('QueryNormalize', fallback=0), data_type=value_types[value_type], target_type=section.get('TargetType', fallback='float32'), - k=None, + k=config.getint('SearchSSDIndex', 'ResultNum'), dim=base.getint('Dim'), B=section.getint('DataBatchSize', fallback=-1), Q=query_count, @@ -420,43 +422,54 @@ def rabitq_bits_from_quantized_dimension(dim, quan_dim): raise ValueError(f'RaBitQ bits must be in [1, 8], got {bits}') return bits -def load_ground_truth(path, query_count, topk=None): - ground_truths = [] - infer_topk = topk is None +def load_ground_truth(path, query_count): + rows = [] + candidate_count = None with open(path, 'r') as truth_file: for query_id in range(query_count): line = truth_file.readline() if not line: raise ValueError(f'ground truth contains only {query_id} queries, expected {query_count}') neighbors = line.strip().split() - if topk is None: - topk = len(neighbors) - if topk == 0: + if candidate_count is None: + candidate_count = len(neighbors) + if candidate_count == 0: raise ValueError('ground truth query 0 contains no neighbors') - if len(neighbors) < topk: + if len(neighbors) != candidate_count: raise ValueError( - f'ground truth query {query_id} contains {len(neighbors)} neighbors, expected at least {topk}') - if infer_topk and len(neighbors) != topk: - raise ValueError( - f'ground truth query {query_id} contains {len(neighbors)} neighbors, expected exactly {topk}') - truth = set(int(neighbor) for neighbor in neighbors[:topk]) - if len(truth) != topk: + f'ground truth query {query_id} contains {len(neighbors)} neighbors, ' + f'expected exactly {candidate_count}') + row = [int(neighbor) for neighbor in neighbors] + if len(set(row)) != candidate_count: raise ValueError(f'ground truth query {query_id} contains duplicate neighbor IDs') - ground_truths.append(truth) - return ground_truths, topk + rows.append(row) + return np.asarray(rows, dtype=np.int64) -def recall_at_k(faiss_index, queries, ground_truths, topk, batch_size=64): - if len(queries) != len(ground_truths): +def reranking_recall_at_k(faiss, faiss_index, queries, candidates, topk): + if len(queries) != len(candidates): raise ValueError('query and ground-truth counts differ') if len(queries) == 0: raise ValueError('at least one query is required') + if topk <= 0: + raise ValueError('ResultNum must be positive') + candidate_count = candidates.shape[1] + if candidate_count <= topk: + raise ValueError( + f'ground-truth candidate depth must exceed ResultNum: {candidate_count} <= {topk}') + if np.any(candidates < 0) or np.any(candidates >= faiss_index.ntotal): + raise ValueError('ground truth contains a vector ID outside the base data') + recall_sum = 0.0 - for start in tqdm.tqdm(range(0, len(queries), batch_size)): - end = min(start + batch_size, len(queries)) - _, results = faiss_index.search(np.asarray(queries[start:end]), topk) - for offset, candidates in enumerate(results): - truth = ground_truths[start + offset] - recall_sum += len(truth.intersection(int(candidate) for candidate in candidates if candidate >= 0)) / len(truth) + for query_id in tqdm.tqdm(range(len(queries))): + candidate_ids = np.ascontiguousarray(candidates[query_id]) + parameters = faiss.SearchParameters() + parameters.sel = faiss.IDSelectorArray( + candidate_count, faiss.swig_ptr(candidate_ids)) + _, results = faiss_index.search( + np.ascontiguousarray(queries[query_id:query_id + 1]), topk, + params=parameters) + expected = set(int(candidate) for candidate in candidate_ids[:topk]) + recall_sum += len(expected.intersection(int(candidate) for candidate in results[0])) / topk return recall_sum / len(queries) def create_rabitq_index(faiss, dim, bits, training_data): @@ -479,7 +492,7 @@ def add_rabitq_data(args, faiss_index): raise ValueError('RaBitQ input data is empty') return total -def tune_rabitq_bits(args, faiss, training_data, queries, ground_truths): +def tune_rabitq_bits(args, faiss, training_data, queries, ground_truth_candidates): if not 0.0 < args.rabitq_target_recall <= 1.0: raise ValueError('rabitq_target_recall must be in (0, 1]') if args.rabitq_min_bits < 1 or args.rabitq_max_bits > 8 or args.rabitq_min_bits > args.rabitq_max_bits: @@ -490,7 +503,8 @@ def tune_rabitq_bits(args, faiss, training_data, queries, ground_truths): print(f'Auto tuning RaBitQ{bits} for Recall@{args.k} >= {args.rabitq_target_recall:.6f}') candidate = create_rabitq_index(faiss, args.dim, bits, training_data) data_count = add_rabitq_data(args, candidate) - recall = recall_at_k(candidate, queries, ground_truths, args.k) + recall = reranking_recall_at_k( + faiss, candidate, queries, ground_truth_candidates, args.k) trials.append({'bits': bits, 'recall': recall}) print(f'RaBitQ{bits} Recall@{args.k}: {recall:.6f}') if recall >= args.rabitq_target_recall: @@ -512,10 +526,8 @@ def train_rabitq(args): raise ValueError('Cosine RaBitQ tuning requires data_normalize=1 and query_normalize=1') if args.train_samples <= 0: raise ValueError('train_samples must be positive') - if args.Q <= 0: - raise ValueError('Q must be positive') - if not args.rabitq_auto_tune and (args.k is None or args.k <= 0): - raise ValueError('k must be positive') + if args.Q <= 0 or args.k <= 0: + raise ValueError('Q and k must be positive') if len(args.output_quan_vector_file) > 0 or len(args.output_rec_vector_file) > 0: raise ValueError( 'RaBitQ vectors must be generated by the native SPTAG quantizer after tuning; ' @@ -537,18 +549,17 @@ def train_rabitq(args): queryreader.close() if num_query != args.Q: raise ValueError(f'query file contains {num_query} queries, but configured Q is {args.Q}') - ground_truths, truth_topk = load_ground_truth( - args.output_truth, num_query, args.k) - if args.k is None: - args.k = truth_topk + ground_truth_candidates = load_ground_truth( + args.output_truth, num_query) bits, data_count, trials = tune_rabitq_bits( - args, faiss, training_data, queries, ground_truths) + args, faiss, training_data, queries, ground_truth_candidates) result = { 'selected_bits': bits, 'native_quantizer_qd': bits, 'storage_bytes_per_vector': sptag_rabitq_storage_bytes(args.dim, bits), 'target_recall': args.rabitq_target_recall, 'recall_at': args.k, + 'rerank_candidate_count': ground_truth_candidates.shape[1], 'query_count': num_query, 'data_count': data_count, 'trials': trials, diff --git a/Tools/OPQ/README.md b/Tools/OPQ/README.md index 7796f440f..ca973f97d 100644 --- a/Tools/OPQ/README.md +++ b/Tools/OPQ/README.md @@ -19,9 +19,11 @@ python3 OPQ_gpu_train_infer.py --data_file openai_5M.bin --query_file openai_tes `--rabitq_auto_tune` runs before vector encoding or index construction. It evaluates RaBitQ bit counts in ascending order and selects the first (therefore minimum) -count whose exhaustive `Recall@k` reaches `--rabitq_target_recall`. Every candidate +count whose reranking Recall reaches `--rabitq_target_recall`. Every candidate uses the same first `--train_samples` vectors, the same configured `--Q` queries, -and the pre-generated `--output_truth` top-`k` neighbors. +and the full pre-generated ground-truth candidate pool. In INI mode, +`[SearchSSDIndex] ResultNum` supplies K: the candidate pool is reranked by +RaBitQ distance and its first K IDs are compared with the exact first K. The command fails instead of silently choosing a bit count when the configured query/ground-truth count is unavailable or no candidate in diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py index 09b78d923..cacdd6082 100644 --- a/Tools/OPQ/test_rabitq_auto_tuning.py +++ b/Tools/OPQ/test_rabitq_auto_tuning.py @@ -12,17 +12,6 @@ SPEC.loader.exec_module(MODULE) -class FakeIndex: - def __init__(self, results): - self.results = np.asarray(results, dtype=np.int64) - - def search(self, queries, topk): - count = len(queries) - results = self.results[:count, :topk] - self.results = self.results[count:] - return np.zeros(results.shape, dtype=np.float32), results - - class RaBitQAutoTuningTest(unittest.TestCase): def test_data_reader_accepts_fbin(self): with tempfile.TemporaryDirectory() as directory: @@ -54,14 +43,10 @@ def test_load_ground_truth_honors_query_count_and_topk(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / 'truth.txt' path.write_text('1 2 3\n4 5 6\n', encoding='ascii') - truths, topk = MODULE.load_ground_truth(path, 2) - self.assertEqual(3, topk) - self.assertEqual([{1, 2, 3}, {4, 5, 6}], truths) - truths, topk = MODULE.load_ground_truth(path, 2, 2) - self.assertEqual(2, topk) - self.assertEqual([{1, 2}, {4, 5}], truths) + truths = MODULE.load_ground_truth(path, 2) + np.testing.assert_array_equal(truths, [[1, 2, 3], [4, 5, 6]]) with self.assertRaises(ValueError): - MODULE.load_ground_truth(path, 3, 2) + MODULE.load_ground_truth(path, 3) def test_inferred_ground_truth_topk_requires_uniform_unique_rows(self): with tempfile.TemporaryDirectory() as directory: @@ -73,11 +58,19 @@ def test_inferred_ground_truth_topk_requires_uniform_unique_rows(self): with self.assertRaises(ValueError): MODULE.load_ground_truth(path, 2) - def test_recall_is_query_equal_weighted(self): - index = FakeIndex([[1, 9], [4, 5]]) - queries = np.zeros((2, 4), dtype=np.float32) - recall = MODULE.recall_at_k(index, queries, [{1, 2}, {3, 4}], 2, batch_size=2) - self.assertEqual(0.5, recall) + def test_reranking_recall_uses_result_num_with_deeper_candidates(self): + import faiss + base = np.asarray([[0.0], [1.0], [2.0], [3.0]], dtype=np.float32) + queries = np.asarray([[0.1]], dtype=np.float32) + index = faiss.IndexFlatL2(1) + index.add(base) + candidates = np.asarray([[0, 1, 2, 3]], dtype=np.int64) + recall = MODULE.reranking_recall_at_k( + faiss, index, queries, candidates, 2) + self.assertEqual(1.0, recall) + with self.assertRaises(ValueError): + MODULE.reranking_recall_at_k( + faiss, index, queries, candidates[:, :2], 2) def test_ini_is_authoritative_and_reuses_query_count_limit(self): with tempfile.TemporaryDirectory() as directory: @@ -99,12 +92,13 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): 'MaxBits=7\n' '\n' '[SearchSSDIndex]\n' - 'QueryCountLimit=10000\n', + 'QueryCountLimit=10000\n' + 'ResultNum=100\n', encoding='ascii') args = MODULE.get_config(['--config', str(path)]) self.assertTrue(args.rabitq_auto_tune) self.assertEqual(10000, args.Q) - self.assertIsNone(args.k) + self.assertEqual(100, args.k) self.assertEqual(0.97, args.rabitq_target_recall) self.assertEqual((2, 7), (args.rabitq_min_bits, args.rabitq_max_bits)) self.assertEqual('base.bin', args.data_file) @@ -155,7 +149,8 @@ def test_ini_rejects_inherited_defaults(self): 'OutputDir=tuning\n' '\n' '[SearchSSDIndex]\n' - 'QueryCountLimit=10000\n', + 'QueryCountLimit=10000\n' + 'ResultNum=100\n', encoding='ascii') with self.assertRaises(ValueError): MODULE.get_config(['--config', str(path)]) diff --git a/docs/GettingStart.md b/docs/GettingStart.md index 9f294b496..cb7ffb6de 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -282,6 +282,7 @@ QueryNormalize=0 [SearchSSDIndex] QueryCountLimit=10000 +ResultNum=100 ``` Run the pre-build tuning stage with only the INI path: @@ -293,11 +294,15 @@ python3 Tools/OPQ/OPQ_gpu_train_infer.py \ The tuner evaluates bit counts in ascending order using exactly the configured `SearchSSDIndex.QueryCountLimit` queries and selects the first bit count meeting -the target Recall. The Recall cutoff is inferred from the exact ground-truth -width: 1,000 neighbor IDs per query means `Recall@1000`, so there is no second -top-k setting. All consumed ground-truth rows must have exactly that width with -no duplicate IDs. When `--config` is used, additional command-line parameters -are rejected, so they cannot override the INI. A +the target Recall. `SearchSSDIndex.ResultNum` is the expected result K. The +ground-truth width is a separate, deeper reranking candidate pool: with +`ResultNum=100` and 1,000 exact IDs per query, each RaBitQ candidate reranks +those 1,000 IDs and is evaluated as `Recall@100` against the exact first 100. +The candidate depth must be strictly greater than `ResultNum`; this prevents a +top100 ground truth from being used to evaluate K=100. All consumed +ground-truth rows must have exactly the same width with no duplicate IDs. When +`--config` is used, additional command-line parameters are rejected, so they +cannot override the INI. A `RaBitQAutoTune.QueryCount` may be set explicitly only when the tuning query count intentionally differs from `SearchSSDIndex.QueryCountLimit`. Strict mode also rejects unknown `[RaBitQAutoTune]` keys and inherited `[DEFAULT]` values. From 434b9127d9d3f4e7c3af3ee01372bf994d5579b4 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:03:41 +0000 Subject: [PATCH 09/17] Simplify RaBitQ tuning configuration Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Script_AE/iniFile/rabitq_auto_tune_sift1m.ini | 4 ---- Tools/OPQ/OPQ_gpu_train_infer.py | 21 +++++++++------- Tools/OPQ/test_rabitq_auto_tuning.py | 24 +++++++++++++++++++ docs/GettingStart.md | 6 ++--- 4 files changed, 39 insertions(+), 16 deletions(-) diff --git a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini index c17e40e50..208c6ff9a 100644 --- a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini +++ b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini @@ -12,10 +12,6 @@ TruthType=DEFAULT [RaBitQAutoTune] isExecute=true OutputDir=sift1m/rabitq_tuning -TargetType=float32 -DataBatchSize=1000000 -Threads=46 -TrainingSamples=1000000 TargetRecall=0.95 MinBits=1 MaxBits=8 diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index d624ae0ed..d7a0fc79d 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -14,6 +14,8 @@ import subprocess import sys +RABITQ_BATCH_SIZE = 1000000 + def get_cli_parser(): parser = argparse.ArgumentParser(description ='implementation of nnsearch.') parser.add_argument('--config', type = str, help='native INI containing [RaBitQAutoTune]') @@ -63,8 +65,7 @@ def load_rabitq_auto_tune_ini(path): raise ValueError('INI file is missing [Base]') section = config[section_name] allowed_keys = { - 'isexecute', 'outputdir', 'targettype', 'databatchsize', 'querycount', - 'threads', 'trainingsamples', 'targetrecall', + 'isexecute', 'outputdir', 'querycount', 'targetrecall', 'minbits', 'maxbits', 'tuningresult', 'datanormalize', 'querynormalize', 'dataformat', 'task', } @@ -113,10 +114,10 @@ def load_rabitq_auto_tune_ini(path): data_normalize=section.getint('DataNormalize', fallback=0), query_normalize=section.getint('QueryNormalize', fallback=0), data_type=value_types[value_type], - target_type=section.get('TargetType', fallback='float32'), + target_type='float32', k=config.getint('SearchSSDIndex', 'ResultNum'), dim=base.getint('Dim'), - B=section.getint('DataBatchSize', fallback=-1), + B=RABITQ_BATCH_SIZE, Q=query_count, S=1000, D=base['DistCalcMethod'], @@ -124,8 +125,8 @@ def load_rabitq_auto_tune_ini(path): data_format=section.get('DataFormat', fallback='DEFAULT'), task=section.getint('Task', fallback=0), log_dir='', - T=section.getint('Threads', fallback=32), - train_samples=section.getint('TrainingSamples', fallback=1000000), + T=None, + train_samples=RABITQ_BATCH_SIZE, quan_type='rabitq', quan_dim=-1, output_dir=section['OutputDir'], @@ -179,7 +180,10 @@ def __init__(self, filename, featuredim, batchsize, normalize, datatype, targett self.isbinary = False self.type = self.mytype - if batchsize <= 0: batchsize = R + if batchsize <= 0: + batchsize = R + else: + batchsize = min(batchsize, R) self.query = np.zeros([batchsize, self.featuredim], dtype=self.mytype) self.normalize = normalize @@ -539,7 +543,8 @@ def train_rabitq(args): if num_training == 0: raise ValueError('RaBitQ training data is empty') print(f'train RaBitQ using {num_training} samples ...') - faiss.omp_set_num_threads(args.T) + if args.T is not None: + faiss.omp_set_num_threads(args.T) if args.rabitq_auto_tune: if args.quan_test <= 0: raise ValueError('rabitq_auto_tune requires quan_test > 0 and a pre-generated ground truth') diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py index cacdd6082..6152ffecd 100644 --- a/Tools/OPQ/test_rabitq_auto_tuning.py +++ b/Tools/OPQ/test_rabitq_auto_tuning.py @@ -129,6 +129,30 @@ def test_ini_rejects_unknown_parameters(self): with self.assertRaises(ValueError): MODULE.get_config(['--config', str(path)]) + def test_ini_rejects_removed_execution_parameters(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'build.ini' + path.write_text( + '[Base]\n' + 'ValueType=Float\n' + 'DistCalcMethod=L2\n' + 'Dim=128\n' + 'VectorPath=base.bin\n' + 'QueryPath=query.bin\n' + 'TruthPath=truth.txt\n' + '\n' + '[RaBitQAutoTune]\n' + 'isExecute=true\n' + 'OutputDir=tuning\n' + 'Threads=46\n' + '\n' + '[SearchSSDIndex]\n' + 'QueryCountLimit=10000\n' + 'ResultNum=100\n', + encoding='ascii') + with self.assertRaises(ValueError): + MODULE.get_config(['--config', str(path)]) + def test_ini_rejects_inherited_defaults(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / 'build.ini' diff --git a/docs/GettingStart.md b/docs/GettingStart.md index cb7ffb6de..887174715 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -270,10 +270,6 @@ TruthType=DEFAULT [RaBitQAutoTune] isExecute=true OutputDir=sift1m/rabitq_tuning -TargetType=float32 -DataBatchSize=1000000 -Threads=46 -TrainingSamples=1000000 TargetRecall=0.95 MinBits=1 MaxBits=8 @@ -306,6 +302,8 @@ cannot override the INI. A `RaBitQAutoTune.QueryCount` may be set explicitly only when the tuning query count intentionally differs from `SearchSSDIndex.QueryCountLimit`. Strict mode also rejects unknown `[RaBitQAutoTune]` keys and inherited `[DEFAULT]` values. +Faiss input is always Float32; batching and centroid training use internal +one-million-vector batches, and Faiss/OpenMP controls its own worker count. It fails if fewer queries/ground-truth rows are available or no candidate qualifies. The ground truth is not moved or modified. Results are written atomically to `sift1m/rabitq_tuning/rabitq_auto_tuning.json`; use From 66524cb969e6f3bc83046a240f149d880b85118e Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:06:52 +0000 Subject: [PATCH 10/17] Reuse unified RaBitQ thread configuration Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Script_AE/iniFile/rabitq_auto_tune_sift1m.ini | 3 +++ Tools/OPQ/OPQ_gpu_train_infer.py | 9 ++++++--- Tools/OPQ/README.md | 2 +- Tools/OPQ/test_rabitq_auto_tuning.py | 4 ++++ docs/GettingStart.md | 8 ++++++-- 5 files changed, 20 insertions(+), 6 deletions(-) diff --git a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini index 208c6ff9a..d3026a7c4 100644 --- a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini +++ b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini @@ -18,6 +18,9 @@ MaxBits=8 DataNormalize=0 QueryNormalize=0 +[BuildSSDIndex] +NumberOfThreads=46 + [SearchSSDIndex] QueryCountLimit=10000 ResultNum=100 diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index d7a0fc79d..44c8ff51f 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -106,6 +106,8 @@ def load_rabitq_auto_tune_ini(path): query_count = config.getint('SearchSSDIndex', 'QueryCountLimit') if not config.has_option('SearchSSDIndex', 'ResultNum'): raise ValueError('[SearchSSDIndex] ResultNum is required for RaBitQ tuning') + if not config.has_option('BuildSSDIndex', 'NumberOfThreads'): + raise ValueError('[BuildSSDIndex] NumberOfThreads is required for RaBitQ tuning') return argparse.Namespace( config=path, @@ -125,7 +127,7 @@ def load_rabitq_auto_tune_ini(path): data_format=section.get('DataFormat', fallback='DEFAULT'), task=section.getint('Task', fallback=0), log_dir='', - T=None, + T=config.getint('BuildSSDIndex', 'NumberOfThreads'), train_samples=RABITQ_BATCH_SIZE, quan_type='rabitq', quan_dim=-1, @@ -543,8 +545,9 @@ def train_rabitq(args): if num_training == 0: raise ValueError('RaBitQ training data is empty') print(f'train RaBitQ using {num_training} samples ...') - if args.T is not None: - faiss.omp_set_num_threads(args.T) + if args.T <= 0: + raise ValueError('NumberOfThreads must be positive') + faiss.omp_set_num_threads(args.T) if args.rabitq_auto_tune: if args.quan_test <= 0: raise ValueError('rabitq_auto_tune requires quan_test > 0 and a pre-generated ground truth') diff --git a/Tools/OPQ/README.md b/Tools/OPQ/README.md index ca973f97d..72de74642 100644 --- a/Tools/OPQ/README.md +++ b/Tools/OPQ/README.md @@ -45,4 +45,4 @@ python3 Tools/OPQ/OPQ_gpu_train_infer.py \ In INI mode no additional CLI parameters are accepted. `QueryCount` defaults to `[SearchSSDIndex] QueryCountLimit`, keeping the configured query count in one -place. \ No newline at end of file +place. Faiss/OpenMP reuses `[BuildSSDIndex] NumberOfThreads`. \ No newline at end of file diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py index 6152ffecd..929ba2d8d 100644 --- a/Tools/OPQ/test_rabitq_auto_tuning.py +++ b/Tools/OPQ/test_rabitq_auto_tuning.py @@ -91,6 +91,9 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): 'MinBits=2\n' 'MaxBits=7\n' '\n' + '[BuildSSDIndex]\n' + 'NumberOfThreads=46\n' + '\n' '[SearchSSDIndex]\n' 'QueryCountLimit=10000\n' 'ResultNum=100\n', @@ -99,6 +102,7 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): self.assertTrue(args.rabitq_auto_tune) self.assertEqual(10000, args.Q) self.assertEqual(100, args.k) + self.assertEqual(46, args.T) self.assertEqual(0.97, args.rabitq_target_recall) self.assertEqual((2, 7), (args.rabitq_min_bits, args.rabitq_max_bits)) self.assertEqual('base.bin', args.data_file) diff --git a/docs/GettingStart.md b/docs/GettingStart.md index 887174715..13109caec 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -276,6 +276,9 @@ MaxBits=8 DataNormalize=0 QueryNormalize=0 +[BuildSSDIndex] +NumberOfThreads=46 + [SearchSSDIndex] QueryCountLimit=10000 ResultNum=100 @@ -302,8 +305,9 @@ cannot override the INI. A `RaBitQAutoTune.QueryCount` may be set explicitly only when the tuning query count intentionally differs from `SearchSSDIndex.QueryCountLimit`. Strict mode also rejects unknown `[RaBitQAutoTune]` keys and inherited `[DEFAULT]` values. -Faiss input is always Float32; batching and centroid training use internal -one-million-vector batches, and Faiss/OpenMP controls its own worker count. +Faiss input is always Float32, and batching and centroid training use internal +one-million-vector batches. Faiss/OpenMP uses the existing +`BuildSSDIndex.NumberOfThreads` value; there is no separate tuner thread setting. It fails if fewer queries/ground-truth rows are available or no candidate qualifies. The ground truth is not moved or modified. Results are written atomically to `sift1m/rabitq_tuning/rabitq_auto_tuning.json`; use From 574a0465afb0018e137b1dd990d88bb7a784f7d8 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:09:50 +0000 Subject: [PATCH 11/17] Use full Base data for RaBitQ centroid Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Tools/OPQ/OPQ_gpu_train_infer.py | 50 ++++++++++++++++++---------- Tools/OPQ/README.md | 6 ++-- Tools/OPQ/test_rabitq_auto_tuning.py | 20 +++++++++++ docs/GettingStart.md | 6 ++-- 4 files changed, 59 insertions(+), 23 deletions(-) diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index 44c8ff51f..486567c22 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -96,6 +96,8 @@ def load_rabitq_auto_tune_ini(path): value_type = base['ValueType'].lower() if value_type not in value_types: raise ValueError(f'unsupported [Base] ValueType for RaBitQ tuning: {base["ValueType"]}') + if value_type != 'float': + raise ValueError('RaBitQ tuning requires [Base] ValueType=Float') query_count = section.getint('QueryCount', fallback=None) if query_count is None: @@ -116,7 +118,7 @@ def load_rabitq_auto_tune_ini(path): data_normalize=section.getint('DataNormalize', fallback=0), query_normalize=section.getint('QueryNormalize', fallback=0), data_type=value_types[value_type], - target_type='float32', + target_type=value_types[value_type], k=config.getint('SearchSSDIndex', 'ResultNum'), dim=base.getint('Dim'), B=RABITQ_BATCH_SIZE, @@ -128,7 +130,7 @@ def load_rabitq_auto_tune_ini(path): task=section.getint('Task', fallback=0), log_dir='', T=config.getint('BuildSSDIndex', 'NumberOfThreads'), - train_samples=RABITQ_BATCH_SIZE, + train_samples=None, quan_type='rabitq', quan_dim=-1, output_dir=section['OutputDir'], @@ -478,11 +480,28 @@ def reranking_recall_at_k(faiss, faiss_index, queries, candidates, topk): recall_sum += len(expected.intersection(int(candidate) for candidate in results[0])) / topk return recall_sum / len(queries) -def create_rabitq_index(faiss, dim, bits, training_data): +def create_rabitq_index(faiss, dim, bits, centroid): faiss_index = faiss.index_factory(dim, f"RaBitQ{bits}", faiss.METRIC_L2) - faiss_index.train(training_data) + faiss_index.train(np.ascontiguousarray(centroid.reshape(1, dim), dtype=np.float32)) return faiss_index +def compute_streaming_centroid(args): + datareader = DataReader( + args.data_file, args.dim, RABITQ_BATCH_SIZE, + args.data_normalize, args.data_type, args.target_type) + accumulator = np.zeros(args.dim, dtype=np.float64) + total = 0 + while True: + num_data, data = datareader.readbatch() + if num_data == 0: + break + accumulator += np.sum(data, axis=0, dtype=np.float64) + total += num_data + datareader.close() + if total == 0: + raise ValueError('RaBitQ input data is empty') + return (accumulator / total).astype(np.float32), total + def add_rabitq_data(args, faiss_index): datareader = DataReader( args.data_file, args.dim, args.B, args.data_normalize, args.data_type, args.target_type) @@ -498,7 +517,7 @@ def add_rabitq_data(args, faiss_index): raise ValueError('RaBitQ input data is empty') return total -def tune_rabitq_bits(args, faiss, training_data, queries, ground_truth_candidates): +def tune_rabitq_bits(args, faiss, centroid, centroid_vector_count, queries, ground_truth_candidates): if not 0.0 < args.rabitq_target_recall <= 1.0: raise ValueError('rabitq_target_recall must be in (0, 1]') if args.rabitq_min_bits < 1 or args.rabitq_max_bits > 8 or args.rabitq_min_bits > args.rabitq_max_bits: @@ -507,8 +526,11 @@ def tune_rabitq_bits(args, faiss, training_data, queries, ground_truth_candidate trials = [] for bits in range(args.rabitq_min_bits, args.rabitq_max_bits + 1): print(f'Auto tuning RaBitQ{bits} for Recall@{args.k} >= {args.rabitq_target_recall:.6f}') - candidate = create_rabitq_index(faiss, args.dim, bits, training_data) + candidate = create_rabitq_index(faiss, args.dim, bits, centroid) data_count = add_rabitq_data(args, candidate) + if data_count != centroid_vector_count: + raise RuntimeError( + f'base vector count changed during tuning: {centroid_vector_count} != {data_count}') recall = reranking_recall_at_k( faiss, candidate, queries, ground_truth_candidates, args.k) trials.append({'bits': bits, 'recall': recall}) @@ -530,21 +552,14 @@ def train_rabitq(args): raise ValueError('RaBitQ training supports only L2, or Cosine with normalized data and queries') if args.D == 'Cosine' and (args.data_normalize == 0 or args.query_normalize == 0): raise ValueError('Cosine RaBitQ tuning requires data_normalize=1 and query_normalize=1') - if args.train_samples <= 0: - raise ValueError('train_samples must be positive') if args.Q <= 0 or args.k <= 0: raise ValueError('Q and k must be positive') if len(args.output_quan_vector_file) > 0 or len(args.output_rec_vector_file) > 0: raise ValueError( 'RaBitQ vectors must be generated by the native SPTAG quantizer after tuning; ' 'do not use output_quan_vector_file or output_rec_vector_file') - training_reader = DataReader( - args.data_file, args.dim, args.train_samples, args.data_normalize, args.data_type, args.target_type) - num_training, training_data = training_reader.readbatch() - training_reader.close() - if num_training == 0: - raise ValueError('RaBitQ training data is empty') - print(f'train RaBitQ using {num_training} samples ...') + centroid, centroid_vector_count = compute_streaming_centroid(args) + print(f'train RaBitQ using the centroid of all {centroid_vector_count} base vectors ...') if args.T <= 0: raise ValueError('NumberOfThreads must be positive') faiss.omp_set_num_threads(args.T) @@ -560,7 +575,7 @@ def train_rabitq(args): ground_truth_candidates = load_ground_truth( args.output_truth, num_query) bits, data_count, trials = tune_rabitq_bits( - args, faiss, training_data, queries, ground_truth_candidates) + args, faiss, centroid, centroid_vector_count, queries, ground_truth_candidates) result = { 'selected_bits': bits, 'native_quantizer_qd': bits, @@ -570,6 +585,7 @@ def train_rabitq(args): 'rerank_candidate_count': ground_truth_candidates.shape[1], 'query_count': num_query, 'data_count': data_count, + 'centroid_vector_count': centroid_vector_count, 'trials': trials, } result_path = os.path.join(output_dir, args.rabitq_tuning_result) @@ -581,7 +597,7 @@ def train_rabitq(args): print(f'Selected RaBitQ storage bits: {bits}; result: {result_path}') else: bits = rabitq_bits_from_quantized_dimension(args.dim, args.quan_dim) - faiss_index = create_rabitq_index(faiss, args.dim, bits, training_data) + faiss_index = create_rabitq_index(faiss, args.dim, bits, centroid) if args.quan_test > 0: add_rabitq_data(args, faiss_index) diff --git a/Tools/OPQ/README.md b/Tools/OPQ/README.md index 72de74642..ebe7f16db 100644 --- a/Tools/OPQ/README.md +++ b/Tools/OPQ/README.md @@ -12,16 +12,14 @@ ## Example python3 OPQ_gpu_train_infer.py --data_file perftest_vector.bin --query_file perftest_query.bin --task 0 --data_type float32 --k 5 --dim 64 --B 1000000 --Q 1000 --D L2 --data_format DEFAULT --T 20 --train_samples 1000000 --quan_type opq --quan_dim 32 --output_quantizer quantizer.bin -python3 OPQ_gpu_train_infer.py --data_file laion_5M.bin --query_file laion_test.bin --output_truth truth.txt --output_dir rabitq_tuned --data_type float32 --target_type float32 --k 100 --dim 768 --B 1000000 --Q 10000 --D L2 --train_samples 1000000 --quan_type rabitq --quan_test 1 --rabitq_auto_tune --rabitq_target_recall 0.95 --rabitq_min_bits 1 --rabitq_max_bits 8 -python3 OPQ_gpu_train_infer.py --data_file openai_5M.bin --query_file openai_test.bin --output_truth truth.txt --output_dir rabitq_tuned --data_type float32 --target_type float32 --k 100 --dim 1536 --B 1000000 --Q 10000 --D Cosine --train_samples 1000000 --quan_type rabitq --quan_test 1 --rabitq_auto_tune --rabitq_target_recall 0.95 --data_normalize 1 --query_normalize 1 ## RaBitQ storage-bit auto tuning `--rabitq_auto_tune` runs before vector encoding or index construction. It evaluates RaBitQ bit counts in ascending order and selects the first (therefore minimum) count whose reranking Recall reaches `--rabitq_target_recall`. Every candidate -uses the same first `--train_samples` vectors, the same configured `--Q` queries, -and the full pre-generated ground-truth candidate pool. In INI mode, +uses the centroid streamed over the complete base file, the same configured +`--Q` queries, and the full pre-generated ground-truth candidate pool. In INI mode, `[SearchSSDIndex] ResultNum` supplies K: the candidate pool is reranked by RaBitQ distance and its first K IDs are compared with the exact first K. diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py index 929ba2d8d..b83c7eb5d 100644 --- a/Tools/OPQ/test_rabitq_auto_tuning.py +++ b/Tools/OPQ/test_rabitq_auto_tuning.py @@ -103,6 +103,8 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): self.assertEqual(10000, args.Q) self.assertEqual(100, args.k) self.assertEqual(46, args.T) + self.assertEqual('float32', args.target_type) + self.assertIsNone(args.train_samples) self.assertEqual(0.97, args.rabitq_target_recall) self.assertEqual((2, 7), (args.rabitq_min_bits, args.rabitq_max_bits)) self.assertEqual('base.bin', args.data_file) @@ -112,6 +114,24 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): with self.assertRaises(ValueError): MODULE.get_config(['--config', str(path), '--Q', '1']) + def test_streaming_centroid_uses_all_base_vectors(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'vectors.bin' + vectors = np.arange(20, dtype=np.float32).reshape(5, 4) + with path.open('wb') as output: + np.asarray([5, 4], dtype=np.int32).tofile(output) + vectors.tofile(output) + args = type('Args', (), { + 'data_file': str(path), + 'dim': 4, + 'data_normalize': 0, + 'data_type': 'float32', + 'target_type': 'float32', + })() + centroid, count = MODULE.compute_streaming_centroid(args) + self.assertEqual(5, count) + np.testing.assert_array_equal(centroid, vectors.mean(axis=0)) + def test_ini_rejects_unknown_parameters(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / 'build.ini' diff --git a/docs/GettingStart.md b/docs/GettingStart.md index 13109caec..9841c41c9 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -305,8 +305,10 @@ cannot override the INI. A `RaBitQAutoTune.QueryCount` may be set explicitly only when the tuning query count intentionally differs from `SearchSSDIndex.QueryCountLimit`. Strict mode also rejects unknown `[RaBitQAutoTune]` keys and inherited `[DEFAULT]` values. -Faiss input is always Float32, and batching and centroid training use internal -one-million-vector batches. Faiss/OpenMP uses the existing +The tuner reads the type from `[Base].ValueType`; global RaBitQ currently +requires `Float`. It computes the centroid over the complete base file in +bounded one-million-vector batches, so there is no training-sample parameter. +Faiss/OpenMP uses the existing `BuildSSDIndex.NumberOfThreads` value; there is no separate tuner thread setting. It fails if fewer queries/ground-truth rows are available or no candidate qualifies. The ground truth is not moved or modified. Results are written From ac98813000a5951716b90add1085676fef5deac0 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:11:33 +0000 Subject: [PATCH 12/17] Remove RaBitQ normalization overrides Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Script_AE/iniFile/rabitq_auto_tune_sift1m.ini | 2 -- Tools/OPQ/OPQ_gpu_train_infer.py | 15 ++++----- Tools/OPQ/README.md | 4 ++- Tools/OPQ/test_rabitq_auto_tuning.py | 33 +++++++++++++++++++ docs/GettingStart.md | 7 ++-- 5 files changed, 46 insertions(+), 15 deletions(-) diff --git a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini index d3026a7c4..3c6253527 100644 --- a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini +++ b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini @@ -15,8 +15,6 @@ OutputDir=sift1m/rabitq_tuning TargetRecall=0.95 MinBits=1 MaxBits=8 -DataNormalize=0 -QueryNormalize=0 [BuildSSDIndex] NumberOfThreads=46 diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index 486567c22..7c0fd4e2a 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -66,8 +66,7 @@ def load_rabitq_auto_tune_ini(path): section = config[section_name] allowed_keys = { 'isexecute', 'outputdir', 'querycount', 'targetrecall', - 'minbits', 'maxbits', 'tuningresult', 'datanormalize', 'querynormalize', - 'dataformat', 'task', + 'minbits', 'maxbits', 'tuningresult', 'dataformat', 'task', } unknown_keys = set(section.keys()) - allowed_keys if unknown_keys: @@ -98,6 +97,8 @@ def load_rabitq_auto_tune_ini(path): raise ValueError(f'unsupported [Base] ValueType for RaBitQ tuning: {base["ValueType"]}') if value_type != 'float': raise ValueError('RaBitQ tuning requires [Base] ValueType=Float') + if base['DistCalcMethod'].lower() != 'l2': + raise ValueError('RaBitQ tuning requires [Base] DistCalcMethod=L2') query_count = section.getint('QueryCount', fallback=None) if query_count is None: @@ -115,8 +116,8 @@ def load_rabitq_auto_tune_ini(path): config=path, data_file=base['VectorPath'], query_file=base['QueryPath'], - data_normalize=section.getint('DataNormalize', fallback=0), - query_normalize=section.getint('QueryNormalize', fallback=0), + data_normalize=0, + query_normalize=0, data_type=value_types[value_type], target_type=value_types[value_type], k=config.getint('SearchSSDIndex', 'ResultNum'), @@ -548,10 +549,8 @@ def train_rabitq(args): output_dir = args.output_dir - if args.D not in ('L2', 'Cosine'): - raise ValueError('RaBitQ training supports only L2, or Cosine with normalized data and queries') - if args.D == 'Cosine' and (args.data_normalize == 0 or args.query_normalize == 0): - raise ValueError('Cosine RaBitQ tuning requires data_normalize=1 and query_normalize=1') + if args.D != 'L2': + raise ValueError('SPTAG global RaBitQ tuning requires [Base] DistCalcMethod=L2') if args.Q <= 0 or args.k <= 0: raise ValueError('Q and k must be positive') if len(args.output_quan_vector_file) > 0 or len(args.output_rec_vector_file) > 0: diff --git a/Tools/OPQ/README.md b/Tools/OPQ/README.md index ebe7f16db..ae4a3ffb6 100644 --- a/Tools/OPQ/README.md +++ b/Tools/OPQ/README.md @@ -43,4 +43,6 @@ python3 Tools/OPQ/OPQ_gpu_train_infer.py \ In INI mode no additional CLI parameters are accepted. `QueryCount` defaults to `[SearchSSDIndex] QueryCountLimit`, keeping the configured query count in one -place. Faiss/OpenMP reuses `[BuildSSDIndex] NumberOfThreads`. \ No newline at end of file +place. Faiss/OpenMP reuses `[BuildSSDIndex] NumberOfThreads`. Global RaBitQ +tuning requires `[Base] ValueType=Float` and `DistCalcMethod=L2`; there are no +separate data/query normalization settings. \ No newline at end of file diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py index b83c7eb5d..8a93b9edf 100644 --- a/Tools/OPQ/test_rabitq_auto_tuning.py +++ b/Tools/OPQ/test_rabitq_auto_tuning.py @@ -177,6 +177,39 @@ def test_ini_rejects_removed_execution_parameters(self): with self.assertRaises(ValueError): MODULE.get_config(['--config', str(path)]) + def test_ini_rejects_normalization_and_non_l2_metric(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'build.ini' + config = ( + '[Base]\n' + 'ValueType=Float\n' + 'DistCalcMethod=L2\n' + 'Dim=128\n' + 'VectorPath=base.bin\n' + 'QueryPath=query.bin\n' + 'TruthPath=truth.txt\n' + '\n' + '[RaBitQAutoTune]\n' + 'isExecute=true\n' + 'OutputDir=tuning\n' + 'DataNormalize=1\n' + '\n' + '[BuildSSDIndex]\n' + 'NumberOfThreads=46\n' + '\n' + '[SearchSSDIndex]\n' + 'QueryCountLimit=10000\n' + 'ResultNum=100\n') + path.write_text(config, encoding='ascii') + with self.assertRaises(ValueError): + MODULE.get_config(['--config', str(path)]) + path.write_text( + config.replace('DistCalcMethod=L2', 'DistCalcMethod=Cosine') + .replace('DataNormalize=1\n', ''), + encoding='ascii') + with self.assertRaises(ValueError): + MODULE.get_config(['--config', str(path)]) + def test_ini_rejects_inherited_defaults(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / 'build.ini' diff --git a/docs/GettingStart.md b/docs/GettingStart.md index 9841c41c9..abd84456d 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -273,8 +273,6 @@ OutputDir=sift1m/rabitq_tuning TargetRecall=0.95 MinBits=1 MaxBits=8 -DataNormalize=0 -QueryNormalize=0 [BuildSSDIndex] NumberOfThreads=46 @@ -306,8 +304,9 @@ cannot override the INI. A count intentionally differs from `SearchSSDIndex.QueryCountLimit`. Strict mode also rejects unknown `[RaBitQAutoTune]` keys and inherited `[DEFAULT]` values. The tuner reads the type from `[Base].ValueType`; global RaBitQ currently -requires `Float`. It computes the centroid over the complete base file in -bounded one-million-vector batches, so there is no training-sample parameter. +requires `Float` and `[Base].DistCalcMethod=L2`. It does not independently +normalize data or queries. It computes the centroid over the complete base file +in bounded one-million-vector batches, so there is no training-sample parameter. Faiss/OpenMP uses the existing `BuildSSDIndex.NumberOfThreads` value; there is no separate tuner thread setting. It fails if fewer queries/ground-truth rows are available or no candidate From 9c2ce65010a87adb15b3987dbba9d4002060923e Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:16:15 +0000 Subject: [PATCH 13/17] Fix RaBitQ tuning range to supported bits Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- Script_AE/iniFile/rabitq_auto_tune_sift1m.ini | 2 -- Tools/OPQ/OPQ_gpu_train_infer.py | 14 +++++--------- Tools/OPQ/README.md | 4 ++-- Tools/OPQ/test_rabitq_auto_tuning.py | 8 +++----- docs/GettingStart.md | 5 ++--- 5 files changed, 12 insertions(+), 21 deletions(-) diff --git a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini index 3c6253527..c3438d655 100644 --- a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini +++ b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini @@ -13,8 +13,6 @@ TruthType=DEFAULT isExecute=true OutputDir=sift1m/rabitq_tuning TargetRecall=0.95 -MinBits=1 -MaxBits=8 [BuildSSDIndex] NumberOfThreads=46 diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index 7c0fd4e2a..1fc74ac16 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -15,6 +15,8 @@ import sys RABITQ_BATCH_SIZE = 1000000 +RABITQ_MIN_BITS = 1 +RABITQ_MAX_BITS = 8 def get_cli_parser(): parser = argparse.ArgumentParser(description ='implementation of nnsearch.') @@ -47,8 +49,6 @@ def get_cli_parser(): parser.add_argument('--quan_test', type = int, default = 0, help='compare with ground truth') parser.add_argument('--rabitq_auto_tune', action = 'store_true', help='select the minimum RaBitQ storage bits before encoding') parser.add_argument('--rabitq_target_recall', type = float, default = 0.95, help='minimum Recall@k for RaBitQ auto tuning') - parser.add_argument('--rabitq_min_bits', type = int, default = 1, help='minimum RaBitQ storage bits to evaluate') - parser.add_argument('--rabitq_max_bits', type = int, default = 8, help='maximum RaBitQ storage bits to evaluate') parser.add_argument('--rabitq_tuning_result', type = str, default = 'rabitq_auto_tuning.json', help='auto-tuning result file under output_dir') return parser @@ -66,7 +66,7 @@ def load_rabitq_auto_tune_ini(path): section = config[section_name] allowed_keys = { 'isexecute', 'outputdir', 'querycount', 'targetrecall', - 'minbits', 'maxbits', 'tuningresult', 'dataformat', 'task', + 'tuningresult', 'dataformat', 'task', } unknown_keys = set(section.keys()) - allowed_keys if unknown_keys: @@ -141,8 +141,6 @@ def load_rabitq_auto_tune_ini(path): quan_test=1, rabitq_auto_tune=True, rabitq_target_recall=section.getfloat('TargetRecall', fallback=0.95), - rabitq_min_bits=section.getint('MinBits', fallback=1), - rabitq_max_bits=section.getint('MaxBits', fallback=8), rabitq_tuning_result=section.get( 'TuningResult', fallback='rabitq_auto_tuning.json'), ) @@ -521,11 +519,9 @@ def add_rabitq_data(args, faiss_index): def tune_rabitq_bits(args, faiss, centroid, centroid_vector_count, queries, ground_truth_candidates): if not 0.0 < args.rabitq_target_recall <= 1.0: raise ValueError('rabitq_target_recall must be in (0, 1]') - if args.rabitq_min_bits < 1 or args.rabitq_max_bits > 8 or args.rabitq_min_bits > args.rabitq_max_bits: - raise ValueError('RaBitQ tuning range must satisfy 1 <= min_bits <= max_bits <= 8') trials = [] - for bits in range(args.rabitq_min_bits, args.rabitq_max_bits + 1): + for bits in range(RABITQ_MIN_BITS, RABITQ_MAX_BITS + 1): print(f'Auto tuning RaBitQ{bits} for Recall@{args.k} >= {args.rabitq_target_recall:.6f}') candidate = create_rabitq_index(faiss, args.dim, bits, centroid) data_count = add_rabitq_data(args, candidate) @@ -541,7 +537,7 @@ def tune_rabitq_bits(args, faiss, centroid, centroid_vector_count, queries, grou measured = ', '.join(f'{trial["bits"]}-bit={trial["recall"]:.6f}' for trial in trials) raise RuntimeError( - f'No RaBitQ bit count in [{args.rabitq_min_bits}, {args.rabitq_max_bits}] ' + f'No RaBitQ bit count in [{RABITQ_MIN_BITS}, {RABITQ_MAX_BITS}] ' f'reached Recall@{args.k} >= {args.rabitq_target_recall:.6f}; {measured}') def train_rabitq(args): diff --git a/Tools/OPQ/README.md b/Tools/OPQ/README.md index ae4a3ffb6..3bfc52418 100644 --- a/Tools/OPQ/README.md +++ b/Tools/OPQ/README.md @@ -24,8 +24,8 @@ uses the centroid streamed over the complete base file, the same configured RaBitQ distance and its first K IDs are compared with the exact first K. The command fails instead of silently choosing a bit count when the configured -query/ground-truth count is unavailable or no candidate in -`[--rabitq_min_bits, --rabitq_max_bits]` reaches the target. The selected storage +query/ground-truth count is unavailable or no candidate in RaBitQ's fixed +1-through-8-bit range reaches the target. The selected storage bit count, actual SPTAG byte width (including padded dimensions and five Float factors), measured recalls, and sample counts are written atomically to `/rabitq_auto_tuning.json`. Use `native_quantizer_qd` as `-qd` for the diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py index 8a93b9edf..a1362d339 100644 --- a/Tools/OPQ/test_rabitq_auto_tuning.py +++ b/Tools/OPQ/test_rabitq_auto_tuning.py @@ -88,8 +88,6 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): 'isExecute=true\n' 'OutputDir=tuning\n' 'TargetRecall=0.97\n' - 'MinBits=2\n' - 'MaxBits=7\n' '\n' '[BuildSSDIndex]\n' 'NumberOfThreads=46\n' @@ -106,7 +104,6 @@ def test_ini_is_authoritative_and_reuses_query_count_limit(self): self.assertEqual('float32', args.target_type) self.assertIsNone(args.train_samples) self.assertEqual(0.97, args.rabitq_target_recall) - self.assertEqual((2, 7), (args.rabitq_min_bits, args.rabitq_max_bits)) self.assertEqual('base.bin', args.data_file) self.assertEqual('query.bin', args.query_file) self.assertEqual('truth.txt', args.output_truth) @@ -153,7 +150,7 @@ def test_ini_rejects_unknown_parameters(self): with self.assertRaises(ValueError): MODULE.get_config(['--config', str(path)]) - def test_ini_rejects_removed_execution_parameters(self): + def test_ini_rejects_configurable_bit_range(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / 'build.ini' path.write_text( @@ -168,7 +165,8 @@ def test_ini_rejects_removed_execution_parameters(self): '[RaBitQAutoTune]\n' 'isExecute=true\n' 'OutputDir=tuning\n' - 'Threads=46\n' + 'MinBits=2\n' + 'MaxBits=7\n' '\n' '[SearchSSDIndex]\n' 'QueryCountLimit=10000\n' diff --git a/docs/GettingStart.md b/docs/GettingStart.md index abd84456d..7b482d0c1 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -271,8 +271,6 @@ TruthType=DEFAULT isExecute=true OutputDir=sift1m/rabitq_tuning TargetRecall=0.95 -MinBits=1 -MaxBits=8 [BuildSSDIndex] NumberOfThreads=46 @@ -291,7 +289,8 @@ python3 Tools/OPQ/OPQ_gpu_train_infer.py \ The tuner evaluates bit counts in ascending order using exactly the configured `SearchSSDIndex.QueryCountLimit` queries and selects the first bit count meeting -the target Recall. `SearchSSDIndex.ResultNum` is the expected result K. The +the target Recall. The supported 1 through 8 bit range is fixed by RaBitQ and is +always evaluated from the minimum upward. `SearchSSDIndex.ResultNum` is the expected result K. The ground-truth width is a separate, deeper reranking candidate pool: with `ResultNum=100` and 1,000 exact IDs per query, each RaBitQ candidate reranks those 1,000 IDs and is evaluated as `Recall@100` against the exact first 100. From a347866834f02dcb7f6636c119a681164fd8af22 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:35:11 +0000 Subject: [PATCH 14/17] Implement native RaBitQ bit auto-tuning Replace the Python/Faiss tuning path with an INI-driven C++ pipeline backed by the pinned official RaBitQ submodule. Select the minimum qualifying bit width, encode the full base set, build SPANN in staging, and publish a reloadable ADC index. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- AnnService/CMakeLists.txt | 10 +- AnnService/inc/Core/Common/RaBitQAutoTuner.h | 56 ++ AnnService/inc/Core/Common/RaBitQQuantizer.h | 8 + AnnService/inc/Helper/VectorSetReader.h | 7 +- AnnService/inc/SSDServing/SSDIndex.h | 11 +- .../src/Core/Common/RaBitQAutoTuner.cpp | 683 ++++++++++++++++++ .../src/Core/Common/RaBitQQuantizer.cpp | 65 +- AnnService/src/Core/SPANN/SPANNIndex.cpp | 1 + AnnService/src/Helper/VectorSetReader.cpp | 23 +- AnnService/src/IndexBuilder/main.cpp | 363 +++++++++- Script_AE/iniFile/rabitq_auto_tune_sift1m.ini | 22 - Test/src/RaBitQQuantizerTest.cpp | 222 ++++++ Tools/OPQ/OPQ_gpu_train_infer.py | 400 +++------- Tools/OPQ/README.md | 36 +- Tools/OPQ/test_rabitq_auto_tuning.py | 239 ------ docs/GettingStart.md | 130 ++-- 16 files changed, 1559 insertions(+), 717 deletions(-) create mode 100644 AnnService/inc/Core/Common/RaBitQAutoTuner.h create mode 100644 AnnService/src/Core/Common/RaBitQAutoTuner.cpp delete mode 100644 Script_AE/iniFile/rabitq_auto_tune_sift1m.ini delete mode 100644 Tools/OPQ/test_rabitq_auto_tuning.py diff --git a/AnnService/CMakeLists.txt b/AnnService/CMakeLists.txt index ecc01eb12..b48354978 100644 --- a/AnnService/CMakeLists.txt +++ b/AnnService/CMakeLists.txt @@ -144,17 +144,23 @@ if (RABITQ) else() list(REMOVE_ITEM HDR_FILES ${AnnService}/inc/Core/Common/RaBitQQuantizer.h + ${AnnService}/inc/Core/Common/RaBitQAutoTuner.h ) list(REMOVE_ITEM SRC_FILES ${AnnService}/src/Core/Common/RaBitQQuantizer.cpp + ${AnnService}/src/Core/Common/RaBitQAutoTuner.cpp ) endif() add_library (SPTAGLib SHARED ${SRC_FILES} ${HDR_FILES} ${TiKV_PROTO_SOURCES}) -target_link_libraries (SPTAGLib DistanceUtils ${RabitQ_LIBRARIES} ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_shared ${NUMA_LIBRARY} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES}) +target_link_libraries (SPTAGLib DistanceUtils ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_shared ${NUMA_LIBRARY} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES}) add_library (SPTAGLibStatic STATIC ${SRC_FILES} ${HDR_FILES} ${TiKV_PROTO_SOURCES}) -target_link_libraries (SPTAGLibStatic DistanceUtils ${RabitQ_LIBRARIES} ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_static ${NUMA_LIBRARY_STATIC} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES}) +target_link_libraries (SPTAGLibStatic DistanceUtils ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_static ${NUMA_LIBRARY_STATIC} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES}) +if (RABITQ) + target_link_libraries(SPTAGLib RaBitQOfficialCore) + target_link_libraries(SPTAGLibStatic RaBitQOfficialCore) +endif() if (MSVC) # SPANNIndex.cpp can exceed COFF section limits in Debug without /bigobj. diff --git a/AnnService/inc/Core/Common/RaBitQAutoTuner.h b/AnnService/inc/Core/Common/RaBitQAutoTuner.h new file mode 100644 index 000000000..9bddac8dc --- /dev/null +++ b/AnnService/inc/Core/Common/RaBitQAutoTuner.h @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "inc/Core/Common/RaBitQQuantizer.h" +#include "inc/Helper/SimpleIniReader.h" +#include "inc/Helper/VectorSetReader.h" + +#include +#include +#include + +namespace SPTAG +{ +namespace COMMON +{ + +struct RaBitQAutoTuneResult +{ + int selectedBits = 0; + float recall = 0.0F; + SizeType vectorCount = 0; + DimensionType codeDimension = 0; + std::string quantizerPath; + std::string vectorPath; + std::shared_ptr quantizer; +}; + +class RaBitQAutoTuner +{ +public: + using BitEvaluator = std::function; + + static bool IsEnabled(const Helper::IniReader& p_config); + static ErrorCode Run(Helper::IniReader& p_config, + const std::string& p_outputFolder, + RaBitQAutoTuneResult& p_result, + std::string& p_error); + + static ErrorCode SelectMinimumBits(float p_targetRecall, + const BitEvaluator& p_evaluator, + int& p_selectedBits, + float& p_selectedRecall); + static ErrorCode ValidateTruth(const std::vector>& p_truth, + SizeType p_baseCount, + SizeType p_queryCount, + int p_resultCount, + std::string& p_error); + static float RecallAtK(const std::vector& p_exact, + const std::vector& p_reranked, + int p_resultCount); +}; + +} // namespace COMMON +} // namespace SPTAG diff --git a/AnnService/inc/Core/Common/RaBitQQuantizer.h b/AnnService/inc/Core/Common/RaBitQQuantizer.h index dd036be31..1858712f4 100644 --- a/AnnService/inc/Core/Common/RaBitQQuantizer.h +++ b/AnnService/inc/Core/Common/RaBitQQuantizer.h @@ -25,6 +25,10 @@ class RaBitQQuantizer : public IQuantizer RaBitQQuantizer(DimensionType p_dimension, int p_bits, bool p_normalize); ErrorCode Train(const std::shared_ptr& p_vectors); + ErrorCode BeginTraining(); + ErrorCode AddTrainingBatch(const std::shared_ptr& p_vectors); + ErrorCode FinishTraining(); + std::shared_ptr CreateWithBits(int p_bits) const; float L2Distance(const std::uint8_t* p_x, const std::uint8_t* p_y) const override; float CosineDistance(const std::uint8_t* p_x, const std::uint8_t* p_y) const override; @@ -48,6 +52,7 @@ class RaBitQQuantizer : public IQuantizer DimensionType Dimension() const { return m_dimension; } int Bits() const { return m_bits; } bool Ready() const; + bool Trained() const { return m_trained; } private: struct ModelHeader @@ -90,6 +95,9 @@ class RaBitQQuantizer : public IQuantizer rabitqlib::quant::RabitqConfig m_quantizer_config; rabitqlib::ex_ipfunc m_ip_func = nullptr; std::vector m_centroid; + std::vector m_training_sum; + std::uint64_t m_training_count = 0; + bool m_trained = false; }; } // namespace COMMON diff --git a/AnnService/inc/Helper/VectorSetReader.h b/AnnService/inc/Helper/VectorSetReader.h index ea68419dc..1ca1ba00d 100644 --- a/AnnService/inc/Helper/VectorSetReader.h +++ b/AnnService/inc/Helper/VectorSetReader.h @@ -19,7 +19,12 @@ namespace Helper class ReaderOptions : public ArgumentsParser { public: - ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, VectorFileType p_fileType, std::string p_vectorDelimiter = "|", std::uint32_t p_threadNum = 32, bool p_normalized = false); + ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, VectorFileType p_fileType, + std::string p_vectorDelimiter = "|", std::uint32_t p_threadNum = 32, + bool p_normalized = false); + ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, VectorFileType p_fileType, + std::string p_vectorDelimiter, std::uint32_t p_threadNum, + bool p_normalized, bool p_requireInputFormat); ~ReaderOptions(); diff --git a/AnnService/inc/SSDServing/SSDIndex.h b/AnnService/inc/SSDServing/SSDIndex.h index 8ecea88df..70a15e7f2 100644 --- a/AnnService/inc/SSDServing/SSDIndex.h +++ b/AnnService/inc/SSDServing/SSDIndex.h @@ -181,11 +181,18 @@ namespace SPTAG { int K = p_opts.m_resultNum; int truthK = (p_opts.m_truthResultNum <= 0) ? K : p_opts.m_truthResultNum; ErrorCode ret; + const bool useADC = p_index->m_pQuantizer && p_opts.m_enableADC; + const VectorValueType queryValueType = useADC + ? p_index->m_pQuantizer->GetReconstructType() + : p_opts.m_valueType; + const DimensionType queryDimension = useADC + ? p_index->m_pQuantizer->ReconstructDim() + : p_opts.m_dim; if (!warmupFile.empty()) { SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Start loading warmup query set...\n"); - std::shared_ptr queryOptions(new Helper::ReaderOptions(p_opts.m_valueType, p_opts.m_dim, p_opts.m_warmupType, p_opts.m_warmupDelimiter)); + std::shared_ptr queryOptions(new Helper::ReaderOptions(queryValueType, queryDimension, p_opts.m_warmupType, p_opts.m_warmupDelimiter)); auto queryReader = Helper::VectorSetReader::CreateInstance(queryOptions); if (ErrorCode::Success != (ret = queryReader->LoadFile(p_opts.m_warmupPath))) { @@ -209,7 +216,7 @@ namespace SPTAG { } SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Start loading QuerySet...\n"); - std::shared_ptr queryOptions(new Helper::ReaderOptions(p_opts.m_valueType, p_opts.m_dim, p_opts.m_queryType, p_opts.m_queryDelimiter)); + std::shared_ptr queryOptions(new Helper::ReaderOptions(queryValueType, queryDimension, p_opts.m_queryType, p_opts.m_queryDelimiter)); auto queryReader = Helper::VectorSetReader::CreateInstance(queryOptions); if (ErrorCode::Success != (ret = queryReader->LoadFile(p_opts.m_queryPath))) { diff --git a/AnnService/src/Core/Common/RaBitQAutoTuner.cpp b/AnnService/src/Core/Common/RaBitQAutoTuner.cpp new file mode 100644 index 000000000..c60624c46 --- /dev/null +++ b/AnnService/src/Core/Common/RaBitQAutoTuner.cpp @@ -0,0 +1,683 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "inc/Core/Common/RaBitQAutoTuner.h" + +#include "inc/Core/VectorIndex.h" +#include "inc/Helper/StringConvert.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace SPTAG +{ +namespace COMMON +{ +namespace +{ + +constexpr SizeType kBatchSize = 64 * 1024; +constexpr const char* kSection = "RaBitQAutoTune"; + +template +bool ReadRequired(const Helper::IniReader& p_config, + const char* p_section, + const char* p_name, + T& p_value, + std::string& p_error) +{ + if (!p_config.DoesParameterExist(p_section, p_name)) { + p_error = std::string("[") + p_section + "] " + p_name + " is required"; + return false; + } + const std::string raw = p_config.GetParameter( + p_section, p_name, std::string()); + if (!Helper::Convert::ConvertStringTo(raw.c_str(), p_value)) { + p_error = std::string("invalid [") + p_section + "] " + p_name + ": " + raw; + return false; + } + return true; +} + +bool ReadRequiredString(const Helper::IniReader& p_config, + const char* p_section, + const char* p_name, + std::string& p_value, + std::string& p_error) +{ + if (!p_config.DoesParameterExist(p_section, p_name) || + (p_value = p_config.GetParameter(p_section, p_name, std::string())).empty()) { + p_error = std::string("[") + p_section + "] " + p_name + " is required"; + return false; + } + return true; +} + +ErrorCode LoadTruth(const std::string& p_path, + TruthFileType p_type, + SizeType p_queryCount, + std::vector>& p_truth, + std::string& p_error) +{ + auto input = f_createIO(); + if (!input || !input->Initialize(p_path.c_str(), std::ios::binary | std::ios::in)) { + p_error = "cannot open [Base] TruthPath: " + p_path; + return ErrorCode::FailedOpenFile; + } + + p_truth.clear(); + p_truth.reserve(static_cast(p_queryCount)); + if (p_type == TruthFileType::DEFAULT) { + std::int32_t rows = 0; + std::int32_t depth = 0; + if (input->ReadBinary(sizeof(rows), reinterpret_cast(&rows)) != sizeof(rows) || + input->ReadBinary(sizeof(depth), reinterpret_cast(&depth)) != sizeof(depth) || + rows < p_queryCount || depth <= 0) { + p_error = "invalid or insufficient DEFAULT truth header"; + return ErrorCode::FailedParseValue; + } + for (SizeType query = 0; query < p_queryCount; ++query) { + std::vector ids(static_cast(depth)); + const std::uint64_t bytes = + sizeof(std::int32_t) * static_cast(depth); + if (input->ReadBinary(bytes, reinterpret_cast(ids.data())) != bytes) { + p_error = "DEFAULT truth ended before QueryCountLimit rows"; + return ErrorCode::FailedParseValue; + } + p_truth.emplace_back(ids.begin(), ids.end()); + } + } else if (p_type == TruthFileType::XVEC) { + std::int32_t expectedDepth = -1; + for (SizeType query = 0; query < p_queryCount; ++query) { + std::int32_t depth = 0; + if (input->ReadBinary(sizeof(depth), reinterpret_cast(&depth)) != sizeof(depth) || + depth <= 0 || (expectedDepth >= 0 && depth != expectedDepth)) { + p_error = "XVEC truth has missing or inconsistent candidate depth"; + return ErrorCode::FailedParseValue; + } + expectedDepth = depth; + std::vector ids(static_cast(depth)); + const std::uint64_t bytes = + sizeof(std::int32_t) * static_cast(depth); + if (input->ReadBinary(bytes, reinterpret_cast(ids.data())) != bytes) { + p_error = "XVEC truth ended before QueryCountLimit rows"; + return ErrorCode::FailedParseValue; + } + p_truth.emplace_back(ids.begin(), ids.end()); + } + } else if (p_type == TruthFileType::TXT) { + std::size_t expectedDepth = 0; + std::uint64_t bufferSize = 64 * 1024; + std::unique_ptr buffer(new char[bufferSize]); + for (SizeType query = 0; query < p_queryCount; ++query) { + if (input->ReadString(bufferSize, buffer, '\n') == 0) { + p_error = "TXT truth ended before QueryCountLimit rows"; + return ErrorCode::FailedParseValue; + } + std::vector row; + char* context = nullptr; +#ifdef _MSC_VER + char* token = strtok_s(buffer.get(), " \t", &context); +#else + char* token = strtok_r(buffer.get(), " \t", &context); +#endif + while (token != nullptr) { + SizeType id = -1; + if (!Helper::Convert::ConvertStringTo(token, id)) { + p_error = "TXT truth contains a non-integer candidate ID"; + return ErrorCode::FailedParseValue; + } + row.push_back(id); +#ifdef _MSC_VER + token = strtok_s(nullptr, " \t", &context); +#else + token = strtok_r(nullptr, " \t", &context); +#endif + } + if (row.empty() || (!p_truth.empty() && row.size() != expectedDepth)) { + p_error = "TXT truth has empty or inconsistent candidate depth"; + return ErrorCode::FailedParseValue; + } + expectedDepth = row.size(); + p_truth.emplace_back(std::move(row)); + } + } else { + p_error = "[Base] TruthType must be DEFAULT, XVEC, or TXT"; + return ErrorCode::FailedParseValue; + } + return ErrorCode::Success; +} + +ErrorCode TrainCentroid(const std::shared_ptr& p_reader, + DimensionType p_dimension, + std::shared_ptr& p_model, + SizeType& p_count, + std::string& p_error) +{ + p_model = std::make_shared(p_dimension, 1, false); + if (p_model->BeginTraining() != ErrorCode::Success) { + p_error = "failed to initialize streaming RaBitQ centroid training"; + return ErrorCode::Fail; + } + + p_count = 0; + for (SizeType start = 0;; start += kBatchSize) { + const auto batch = p_reader->GetVectorSet(start, start + kBatchSize); + if (!batch || batch->Count() == 0) { + break; + } + if (batch->GetValueType() != VectorValueType::Float || + batch->Dimension() != p_dimension || + p_model->AddTrainingBatch(batch) != ErrorCode::Success) { + p_error = "base vector batch is incompatible with Float RaBitQ training"; + return ErrorCode::FailedParseValue; + } + p_count += batch->Count(); + if (batch->Count() < kBatchSize) { + break; + } + if (start > (std::numeric_limits::max)() - 2 * kBatchSize) { + p_error = "base vector count exceeds native SizeType"; + return ErrorCode::Fail; + } + } + if (p_count <= 0 || p_model->FinishTraining() != ErrorCode::Success) { + p_error = "base vector source is empty"; + return ErrorCode::EmptyData; + } + return ErrorCode::Success; +} + +ErrorCode EvaluateBits(const std::shared_ptr& p_quantizer, + const std::shared_ptr& p_baseReader, + const std::shared_ptr& p_queries, + const std::vector>& p_truth, + SizeType p_baseCount, + int p_resultCount, + int p_threads, + float& p_recall, + std::string& p_error) +{ + std::vector candidateIds; + for (const auto& row : p_truth) { + candidateIds.insert(candidateIds.end(), row.begin(), row.end()); + } + std::sort(candidateIds.begin(), candidateIds.end()); + candidateIds.erase(std::unique(candidateIds.begin(), candidateIds.end()), candidateIds.end()); + + const DimensionType codeDimension = p_quantizer->GetNumSubvectors(); + std::vector codes( + candidateIds.size() * static_cast(codeDimension)); + p_quantizer->SetEnableADC(true); + + std::size_t candidate = 0; + for (SizeType start = 0; start < p_baseCount; start += kBatchSize) { + const SizeType end = std::min(p_baseCount, start + kBatchSize); + const auto batch = p_baseReader->GetVectorSet(start, end); + if (!batch || batch->Count() != end - start) { + p_error = "base vector source changed while evaluating RaBitQ bits"; + return ErrorCode::Fail; + } + while (candidate < candidateIds.size() && candidateIds[candidate] < end) { + const SizeType id = candidateIds[candidate]; + p_quantizer->QuantizeVector( + batch->GetVector(id - start), + codes.data() + candidate * static_cast(codeDimension), + false); + ++candidate; + } + } + if (candidate != candidateIds.size()) { + p_error = "not all truth candidates could be encoded"; + return ErrorCode::Fail; + } + + std::unordered_map codeOffsets; + codeOffsets.reserve(candidateIds.size()); + for (std::size_t i = 0; i < candidateIds.size(); ++i) { + codeOffsets.emplace(candidateIds[i], i); + } + + const int queryBytes = p_quantizer->QuantizeSize(); + std::vector queryCodes( + static_cast(p_queries->Count()) * static_cast(queryBytes)); + for (SizeType query = 0; query < p_queries->Count(); ++query) { + p_quantizer->QuantizeVector( + p_queries->GetVector(query), + queryCodes.data() + static_cast(query) * queryBytes, + true); + } + + std::vector queryRecalls(static_cast(p_queries->Count()), 0.0F); + std::atomic nextQuery(0); + std::atomic evaluationFailed(false); + const int workerCount = std::max(1, std::min(p_threads, p_queries->Count())); + std::vector workers; + workers.reserve(static_cast(workerCount)); + for (int worker = 0; worker < workerCount; ++worker) { + workers.emplace_back([&]() { + for (;;) { + const SizeType query = nextQuery.fetch_add(1); + if (query >= p_queries->Count()) { + return; + } + std::vector> ranked; + ranked.reserve(p_truth[static_cast(query)].size()); + const std::uint8_t* queryCode = + queryCodes.data() + static_cast(query) * queryBytes; + for (SizeType id : p_truth[static_cast(query)]) { + const auto offset = codeOffsets.find(id); + if (offset == codeOffsets.end()) { + evaluationFailed.store(true); + return; + } + const std::uint8_t* code = + codes.data() + offset->second * static_cast(codeDimension); + ranked.emplace_back(p_quantizer->L2Distance(queryCode, code), id); + } + std::partial_sort( + ranked.begin(), ranked.begin() + p_resultCount, ranked.end(), + [](const auto& p_left, const auto& p_right) { + return p_left.first < p_right.first || + (p_left.first == p_right.first && p_left.second < p_right.second); + }); + std::vector ids; + ids.reserve(static_cast(p_resultCount)); + for (int rank = 0; rank < p_resultCount; ++rank) { + ids.push_back(ranked[static_cast(rank)].second); + } + queryRecalls[static_cast(query)] = + RaBitQAutoTuner::RecallAtK( + p_truth[static_cast(query)], ids, p_resultCount); + } + }); + } + for (auto& worker : workers) { + worker.join(); + } + if (evaluationFailed.load()) { + p_error = "truth candidate code lookup failed during parallel evaluation"; + return ErrorCode::Fail; + } + double recallSum = 0.0; + for (float recall : queryRecalls) { + recallSum += recall; + } + p_recall = static_cast(recallSum / p_queries->Count()); + return ErrorCode::Success; +} + +ErrorCode SaveArtifacts(const std::shared_ptr& p_quantizer, + const std::shared_ptr& p_reader, + SizeType p_count, + const std::string& p_outputFolder, + RaBitQAutoTuneResult& p_result, + std::string& p_error) +{ + namespace fs = std::filesystem; + std::error_code filesystemError; + fs::create_directories(p_outputFolder, filesystemError); + if (filesystemError) { + p_error = "cannot create RaBitQ artifact directory: " + filesystemError.message(); + return ErrorCode::FailedCreateFile; + } + + const fs::path folder(p_outputFolder); + const fs::path marker = folder / "rabitq_auto.incomplete"; + const fs::path modelTemporary = folder / "rabitq_auto_quantizer.bin.incomplete"; + const fs::path vectorTemporary = folder / "rabitq_auto_vectors.bin.incomplete"; + const fs::path modelFinal = folder / "rabitq_auto_quantizer.bin"; + const fs::path vectorFinal = folder / "rabitq_auto_vectors.bin"; + fs::remove(modelTemporary, filesystemError); + fs::remove(vectorTemporary, filesystemError); + { + auto markerOutput = f_createIO(); + if (!markerOutput || + !markerOutput->Initialize(marker.string().c_str(), std::ios::out | std::ios::binary) || + markerOutput->WriteString("RaBitQ auto-tuning artifacts are incomplete\n") == 0) { + p_error = "cannot create RaBitQ incomplete marker"; + return ErrorCode::FailedCreateFile; + } + } + + auto fail = [&](ErrorCode p_status, const std::string& p_message) { + p_error = p_message; + fs::remove(modelTemporary, filesystemError); + fs::remove(vectorTemporary, filesystemError); + return p_status; + }; + + { + auto modelOutput = f_createIO(); + if (!modelOutput || + !modelOutput->Initialize( + modelTemporary.string().c_str(), std::ios::out | std::ios::binary) || + p_quantizer->SaveQuantizer(modelOutput) != ErrorCode::Success) { + return fail(ErrorCode::DiskIOFail, "failed to write RaBitQ quantizer"); + } + } + + const DimensionType codeDimension = p_quantizer->GetNumSubvectors(); + { + auto vectorOutput = f_createIO(); + if (!vectorOutput || + !vectorOutput->Initialize( + vectorTemporary.string().c_str(), std::ios::out | std::ios::binary) || + vectorOutput->WriteBinary(sizeof(p_count), reinterpret_cast(&p_count)) != + sizeof(p_count) || + vectorOutput->WriteBinary( + sizeof(codeDimension), reinterpret_cast(&codeDimension)) != + sizeof(codeDimension)) { + return fail(ErrorCode::DiskIOFail, "failed to initialize encoded vector artifact"); + } + p_quantizer->SetEnableADC(true); + std::vector code(static_cast(codeDimension)); + SizeType written = 0; + for (SizeType start = 0; start < p_count; start += kBatchSize) { + const SizeType end = std::min(p_count, start + kBatchSize); + const auto batch = p_reader->GetVectorSet(start, end); + if (!batch || batch->Count() != end - start) { + return fail(ErrorCode::Fail, "base vector source changed during final encoding"); + } + for (SizeType i = 0; i < batch->Count(); ++i) { + p_quantizer->QuantizeVector(batch->GetVector(i), code.data(), false); + if (vectorOutput->WriteBinary( + code.size(), reinterpret_cast(code.data())) != code.size()) { + return fail(ErrorCode::DiskIOFail, "failed to stream encoded base vectors"); + } + ++written; + } + } + if (written != p_count) { + return fail(ErrorCode::Fail, "encoded vector count mismatch"); + } + } + + const std::uintmax_t expectedVectorSize = + sizeof(SizeType) + sizeof(DimensionType) + + static_cast(p_count) * static_cast(codeDimension); + if (fs::file_size(vectorTemporary, filesystemError) != expectedVectorSize || + filesystemError) { + return fail(ErrorCode::Fail, "encoded vector artifact size mismatch"); + } + + auto modelInput = f_createIO(); + if (!modelInput || + !modelInput->Initialize( + modelTemporary.string().c_str(), std::ios::in | std::ios::binary)) { + return fail(ErrorCode::FailedOpenFile, "cannot reopen generated RaBitQ quantizer"); + } + const auto loaded = IQuantizer::LoadIQuantizer(modelInput); + const auto loadedRaBitQ = std::dynamic_pointer_cast(loaded); + if (!loadedRaBitQ || loadedRaBitQ->Bits() != p_quantizer->Bits() || + loadedRaBitQ->Dimension() != p_quantizer->Dimension() || + loadedRaBitQ->GetNumSubvectors() != codeDimension) { + return fail(ErrorCode::FailedParseValue, "generated model is incompatible with encoded vectors"); + } + + fs::remove(modelFinal, filesystemError); + fs::rename(modelTemporary, modelFinal, filesystemError); + if (filesystemError) { + return fail(ErrorCode::DiskIOFail, "cannot publish RaBitQ quantizer: " + filesystemError.message()); + } + fs::remove(vectorFinal, filesystemError); + fs::rename(vectorTemporary, vectorFinal, filesystemError); + if (filesystemError) { + return fail(ErrorCode::DiskIOFail, "cannot publish encoded vectors: " + filesystemError.message()); + } + fs::remove(marker, filesystemError); + + p_result.quantizerPath = modelFinal.string(); + p_result.vectorPath = vectorFinal.string(); + p_result.codeDimension = codeDimension; + p_result.vectorCount = p_count; + p_result.quantizer = p_quantizer; + return ErrorCode::Success; +} + +} // namespace + +bool RaBitQAutoTuner::IsEnabled(const Helper::IniReader& p_config) +{ + return p_config.DoesSectionExist(kSection) && + p_config.GetParameter(kSection, "isExecute", false); +} + +ErrorCode RaBitQAutoTuner::Run(Helper::IniReader& p_config, + const std::string& p_outputFolder, + RaBitQAutoTuneResult& p_result, + std::string& p_error) +{ + p_result = RaBitQAutoTuneResult(); + p_error.clear(); + for (const auto& parameter : p_config.GetParameters(kSection)) { + if (!Helper::StrUtils::StrEqualIgnoreCase(parameter.first.c_str(), "isExecute") && + !Helper::StrUtils::StrEqualIgnoreCase(parameter.first.c_str(), "TargetRecall")) { + p_error = "unsupported [RaBitQAutoTune] parameter: " + parameter.first; + return ErrorCode::FailedParseValue; + } + } + + VectorValueType valueType = VectorValueType::Undefined; + DistCalcMethod distance = DistCalcMethod::Undefined; + DimensionType dimension = 0; + VectorFileType vectorType = VectorFileType::Undefined; + VectorFileType queryType = VectorFileType::Undefined; + TruthFileType truthType = TruthFileType::Undefined; + SizeType queryCount = 0; + int resultCount = 0; + int threads = 0; + float targetRecall = 0.0F; + std::string vectorPath; + std::string queryPath; + std::string truthPath; + if (!ReadRequired(p_config, "Base", "ValueType", valueType, p_error) || + !ReadRequired(p_config, "Base", "DistCalcMethod", distance, p_error) || + !ReadRequired(p_config, "Base", "Dim", dimension, p_error) || + !ReadRequired(p_config, "Base", "VectorType", vectorType, p_error) || + !ReadRequired(p_config, "Base", "QueryType", queryType, p_error) || + !ReadRequired(p_config, "Base", "TruthType", truthType, p_error) || + !ReadRequiredString(p_config, "Base", "VectorPath", vectorPath, p_error) || + !ReadRequiredString(p_config, "Base", "QueryPath", queryPath, p_error) || + !ReadRequiredString(p_config, "Base", "TruthPath", truthPath, p_error) || + !ReadRequired(p_config, "SearchSSDIndex", "QueryCountLimit", queryCount, p_error) || + !ReadRequired(p_config, "SearchSSDIndex", "ResultNum", resultCount, p_error) || + !ReadRequired(p_config, "BuildSSDIndex", "NumberOfThreads", threads, p_error) || + !ReadRequired(p_config, kSection, "TargetRecall", targetRecall, p_error)) { + return ErrorCode::FailedParseValue; + } + if (valueType != VectorValueType::Float || distance != DistCalcMethod::L2) { + p_error = "global RaBitQ auto-tuning requires [Base] ValueType=Float and DistCalcMethod=L2"; + return ErrorCode::FailedParseValue; + } + if (dimension <= 0 || queryCount <= 0 || resultCount <= 0 || threads <= 0 || + !std::isfinite(targetRecall) || targetRecall < 0.0F || targetRecall > 1.0F || + vectorType == VectorFileType::Undefined || queryType == VectorFileType::Undefined || + truthType == TruthFileType::Undefined || p_outputFolder.empty()) { + p_error = "invalid RaBitQ input dimension, counts, types, threads, target recall, or output folder"; + return ErrorCode::FailedParseValue; + } + + const std::string vectorDelimiter = + p_config.GetParameter("Base", "VectorDelimiter", std::string("|")); + const std::string queryDelimiter = + p_config.GetParameter("Base", "QueryDelimiter", std::string("|")); + auto baseOptions = std::make_shared( + VectorValueType::Float, dimension, vectorType, vectorDelimiter, threads, false); + auto baseReader = Helper::VectorSetReader::CreateInstance(baseOptions); + if (!baseReader || baseReader->LoadFile(vectorPath) != ErrorCode::Success) { + p_error = "failed to load [Base] VectorPath with its declared VectorType"; + return ErrorCode::FailedOpenFile; + } + + std::shared_ptr centroidModel; + SizeType baseCount = 0; + ErrorCode status = + TrainCentroid(baseReader, dimension, centroidModel, baseCount, p_error); + if (status != ErrorCode::Success) { + return status; + } + + auto queryOptions = std::make_shared( + VectorValueType::Float, dimension, queryType, queryDelimiter, threads, false); + auto queryReader = Helper::VectorSetReader::CreateInstance(queryOptions); + if (!queryReader || queryReader->LoadFile(queryPath) != ErrorCode::Success) { + p_error = "failed to load [Base] QueryPath with its declared QueryType"; + return ErrorCode::FailedOpenFile; + } + const auto queries = queryReader->GetVectorSet(0, queryCount); + if (!queries || queries->Count() != queryCount || + queries->GetValueType() != VectorValueType::Float || + queries->Dimension() != dimension) { + p_error = "QueryPath does not contain exactly QueryCountLimit usable Float queries"; + return ErrorCode::FailedParseValue; + } + + std::vector> truth; + status = LoadTruth(truthPath, truthType, queryCount, truth, p_error); + if (status != ErrorCode::Success || + (status = ValidateTruth( + truth, baseCount, queryCount, resultCount, p_error)) != ErrorCode::Success) { + return status; + } + + std::shared_ptr selected; + status = SelectMinimumBits( + targetRecall, + [&](int p_bits, float& p_recall) { + auto candidate = centroidModel->CreateWithBits(p_bits); + if (!candidate) { + p_error = "failed to create RaBitQ candidate from shared centroid"; + return ErrorCode::Fail; + } + const ErrorCode evaluation = EvaluateBits( + candidate, baseReader, queries, truth, baseCount, + resultCount, threads, p_recall, p_error); + SPTAGLIB_LOG( + Helper::LogLevel::LL_Info, + "RaBitQ auto-tuning bits=%d Recall@%d=%.6f target=%.6f\n", + p_bits, resultCount, p_recall, targetRecall); + if (evaluation == ErrorCode::Success && p_recall >= targetRecall) { + selected = std::move(candidate); + } + return evaluation; + }, + p_result.selectedBits, p_result.recall); + if (status != ErrorCode::Success) { + if (p_error.empty()) { + p_error = "no RaBitQ bit width in the fixed range 1..8 meets TargetRecall"; + } + return status; + } + if (!selected || selected->Bits() != p_result.selectedBits) { + p_error = "RaBitQ selected model does not match selected bit width"; + return ErrorCode::Fail; + } + + status = SaveArtifacts( + selected, baseReader, baseCount, p_outputFolder, p_result, p_error); + if (status == ErrorCode::Success) { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Info, + "RaBitQ auto-tuning selected %d bits (Recall@%d=%.6f); encoded %d vectors at width %d\n", + p_result.selectedBits, resultCount, p_result.recall, + p_result.vectorCount, p_result.codeDimension); + } + return status; +} + +ErrorCode RaBitQAutoTuner::SelectMinimumBits(float p_targetRecall, + const BitEvaluator& p_evaluator, + int& p_selectedBits, + float& p_selectedRecall) +{ + p_selectedBits = 0; + p_selectedRecall = 0.0F; + if (!p_evaluator || !std::isfinite(p_targetRecall) || + p_targetRecall < 0.0F || p_targetRecall > 1.0F) { + return ErrorCode::FailedParseValue; + } + for (int bits = 1; bits <= 8; ++bits) { + float recall = 0.0F; + const ErrorCode status = p_evaluator(bits, recall); + if (status != ErrorCode::Success || !std::isfinite(recall)) { + return status == ErrorCode::Success ? ErrorCode::Fail : status; + } + if (recall >= p_targetRecall) { + p_selectedBits = bits; + p_selectedRecall = recall; + return ErrorCode::Success; + } + } + return ErrorCode::Fail; +} + +ErrorCode RaBitQAutoTuner::ValidateTruth( + const std::vector>& p_truth, + SizeType p_baseCount, + SizeType p_queryCount, + int p_resultCount, + std::string& p_error) +{ + if (p_baseCount <= 0 || p_queryCount <= 0 || p_resultCount <= 0 || + p_truth.size() != static_cast(p_queryCount)) { + p_error = "truth query count does not match QueryCountLimit"; + return ErrorCode::FailedParseValue; + } + std::size_t depth = 0; + for (std::size_t query = 0; query < p_truth.size(); ++query) { + const auto& row = p_truth[query]; + if (query == 0) { + depth = row.size(); + if (depth <= static_cast(p_resultCount)) { + p_error = "truth candidate depth must be greater than ResultNum"; + return ErrorCode::FailedParseValue; + } + } else if (row.size() != depth) { + p_error = "truth candidate depth is inconsistent across queries"; + return ErrorCode::FailedParseValue; + } + std::unordered_set seen; + seen.reserve(row.size()); + for (SizeType id : row) { + if (id < 0 || id >= p_baseCount) { + p_error = "truth contains a candidate ID outside the base vector source"; + return ErrorCode::FailedParseValue; + } + if (!seen.insert(id).second) { + p_error = "truth contains a duplicate candidate ID"; + return ErrorCode::FailedParseValue; + } + } + } + return ErrorCode::Success; +} + +float RaBitQAutoTuner::RecallAtK(const std::vector& p_exact, + const std::vector& p_reranked, + int p_resultCount) +{ + if (p_resultCount <= 0 || + p_exact.size() < static_cast(p_resultCount) || + p_reranked.size() < static_cast(p_resultCount)) { + return 0.0F; + } + std::unordered_set exact( + p_exact.begin(), p_exact.begin() + p_resultCount); + int matches = 0; + for (int i = 0; i < p_resultCount; ++i) { + matches += exact.find(p_reranked[static_cast(i)]) != exact.end(); + } + return static_cast(matches) / p_resultCount; +} + +} // namespace COMMON +} // namespace SPTAG diff --git a/AnnService/src/Core/Common/RaBitQQuantizer.cpp b/AnnService/src/Core/Common/RaBitQQuantizer.cpp index 82ef4ee8f..446172c72 100644 --- a/AnnService/src/Core/Common/RaBitQQuantizer.cpp +++ b/AnnService/src/Core/Common/RaBitQQuantizer.cpp @@ -41,34 +41,83 @@ ErrorCode RaBitQQuantizer::Initialize(DimensionType p_dimension, int p_bits, boo m_quantizer_config = rabitqlib::quant::faster_config( static_cast(m_padded_dimension), static_cast(m_bits)); m_ip_func = rabitqlib::select_excode_ipfunc(static_cast(m_bits)); + m_training_sum.clear(); + m_training_count = 0; + m_trained = false; return ErrorCode::Success; } ErrorCode RaBitQQuantizer::Train(const std::shared_ptr& p_vectors) { - if (!Ready() || !p_vectors || p_vectors->GetValueType() != VectorValueType::Float || - p_vectors->Dimension() != m_dimension || p_vectors->Count() <= 0) { + ErrorCode status = BeginTraining(); + if (status == ErrorCode::Success) { + status = AddTrainingBatch(p_vectors); + } + return status == ErrorCode::Success ? FinishTraining() : status; +} + +ErrorCode RaBitQQuantizer::BeginTraining() +{ + if (!Ready()) { return ErrorCode::FailedParseValue; } + m_training_sum.assign(static_cast(m_dimension), 0.0); + m_training_count = 0; + m_trained = false; + return ErrorCode::Success; +} - std::vector accumulator(static_cast(m_dimension), 0.0); +ErrorCode RaBitQQuantizer::AddTrainingBatch(const std::shared_ptr& p_vectors) +{ + if (!Ready() || m_training_sum.size() != static_cast(m_dimension) || + !p_vectors || p_vectors->GetValueType() != VectorValueType::Float || + p_vectors->Dimension() != m_dimension || p_vectors->Count() <= 0) { + return ErrorCode::FailedParseValue; + } + const auto batchCount = static_cast(p_vectors->Count()); + if (m_training_count > (std::numeric_limits::max)() - batchCount) { + return ErrorCode::Fail; + } std::vector prepared; for (SizeType i = 0; i < p_vectors->Count(); ++i) { const auto* vector = static_cast(p_vectors->GetVector(i)); PrepareInput(vector, prepared); for (DimensionType j = 0; j < m_dimension; ++j) { - accumulator[static_cast(j)] += prepared[static_cast(j)]; + m_training_sum[static_cast(j)] += prepared[static_cast(j)]; } } + m_training_count += batchCount; + return ErrorCode::Success; +} - const double inverse_count = 1.0 / static_cast(p_vectors->Count()); +ErrorCode RaBitQQuantizer::FinishTraining() +{ + if (!Ready() || m_training_count == 0 || + m_training_sum.size() != static_cast(m_dimension)) { + return ErrorCode::FailedParseValue; + } + const double inverse_count = 1.0 / static_cast(m_training_count); for (DimensionType j = 0; j < m_dimension; ++j) { m_centroid[static_cast(j)] = - static_cast(accumulator[static_cast(j)] * inverse_count); + static_cast(m_training_sum[static_cast(j)] * inverse_count); } + m_training_sum.clear(); + m_training_count = 0; + m_trained = true; return ErrorCode::Success; } +std::shared_ptr RaBitQQuantizer::CreateWithBits(int p_bits) const +{ + if (!Ready() || !m_trained || p_bits < 1 || p_bits > 8) { + return nullptr; + } + auto quantizer = std::make_shared(m_dimension, p_bits, m_normalize); + quantizer->m_centroid = m_centroid; + quantizer->m_trained = true; + return quantizer; +} + float RaBitQQuantizer::L2Distance(const std::uint8_t* p_x, const std::uint8_t* p_y) const { thread_local std::vector reconstructed_query; @@ -246,7 +295,7 @@ std::uint64_t RaBitQQuantizer::BufferSize() const ErrorCode RaBitQQuantizer::SaveQuantizer(std::shared_ptr p_output) const { - if (!p_output || !Ready()) { + if (!p_output || !Ready() || !m_trained) { return ErrorCode::Fail; } @@ -283,6 +332,7 @@ ErrorCode RaBitQQuantizer::LoadQuantizer(std::shared_ptr p_input m_centroid.size() * sizeof(float)) { return ErrorCode::FailedParseValue; } + m_trained = true; return ErrorCode::Success; } @@ -299,6 +349,7 @@ ErrorCode RaBitQQuantizer::LoadQuantizer(std::uint8_t* p_raw_bytes) } p_raw_bytes += sizeof(header); std::memcpy(m_centroid.data(), p_raw_bytes, m_centroid.size() * sizeof(float)); + m_trained = true; return ErrorCode::Success; } diff --git a/AnnService/src/Core/SPANN/SPANNIndex.cpp b/AnnService/src/Core/SPANN/SPANNIndex.cpp index 35f7a72f5..2b367dbf1 100644 --- a/AnnService/src/Core/SPANN/SPANNIndex.cpp +++ b/AnnService/src/Core/SPANN/SPANNIndex.cpp @@ -57,6 +57,7 @@ template void Index::SetQuantizer(std::shared_ptrSetEnableADC(m_options.m_enableADC); m_fComputeDistance = m_pQuantizer->DistanceCalcSelector(m_options.m_distCalcMethod); m_iBaseSquare = (m_options.m_distCalcMethod == DistCalcMethod::Cosine) ? m_pQuantizer->GetBase() * m_pQuantizer->GetBase() diff --git a/AnnService/src/Helper/VectorSetReader.cpp b/AnnService/src/Helper/VectorSetReader.cpp index 00089fe54..ad2b983f7 100644 --- a/AnnService/src/Helper/VectorSetReader.cpp +++ b/AnnService/src/Helper/VectorSetReader.cpp @@ -9,17 +9,32 @@ using namespace SPTAG; using namespace SPTAG::Helper; +ReaderOptions::ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, + VectorFileType p_fileType, std::string p_vectorDelimiter, + std::uint32_t p_threadNum, bool p_normalized) + : ReaderOptions(p_valueType, p_dimension, p_fileType, p_vectorDelimiter, + p_threadNum, p_normalized, true) +{ +} + ReaderOptions::ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, VectorFileType p_fileType, - std::string p_vectorDelimiter, std::uint32_t p_threadNum, bool p_normalized) + std::string p_vectorDelimiter, std::uint32_t p_threadNum, bool p_normalized, + bool p_requireInputFormat) : m_inputValueType(p_valueType), m_dimension(p_dimension), m_inputFileType(p_fileType), m_vectorDelimiter(p_vectorDelimiter), m_threadNum(p_threadNum), m_normalized(p_normalized) { AddOptionalOption(m_threadNum, "-t", "--thread", "Thread Number."); AddOptionalOption(m_vectorDelimiter, "-dl", "--delimiter", "Vector delimiter."); AddOptionalOption(m_normalized, "-norm", "--normalized", "Vector is normalized."); - AddRequiredOption(m_dimension, "-d", "--dimension", "Dimension of vector."); - AddRequiredOption(m_inputValueType, "-v", "--vectortype", "Input vector data type. Default is float."); - AddRequiredOption(m_inputFileType, "-f", "--filetype", "Input file type (DEFAULT, TXT, XVEC). Default is DEFAULT."); + if (p_requireInputFormat) { + AddRequiredOption(m_dimension, "-d", "--dimension", "Dimension of vector."); + AddRequiredOption(m_inputValueType, "-v", "--vectortype", "Input vector data type. Default is float."); + AddRequiredOption(m_inputFileType, "-f", "--filetype", "Input file type (DEFAULT, TXT, XVEC). Default is DEFAULT."); + } else { + AddOptionalOption(m_dimension, "-d", "--dimension", "Dimension of vector."); + AddOptionalOption(m_inputValueType, "-v", "--vectortype", "Input vector data type. Default is float."); + AddOptionalOption(m_inputFileType, "-f", "--filetype", "Input file type (DEFAULT, TXT, XVEC). Default is DEFAULT."); + } } ReaderOptions::~ReaderOptions() diff --git a/AnnService/src/IndexBuilder/main.cpp b/AnnService/src/IndexBuilder/main.cpp index 56499840f..8de54fad9 100644 --- a/AnnService/src/IndexBuilder/main.cpp +++ b/AnnService/src/IndexBuilder/main.cpp @@ -7,14 +7,144 @@ #include "inc/Helper/VectorSetReader.h" #include +#include +#include +#include +#include #include +#ifdef RABITQ +#include "inc/Core/Common/RaBitQAutoTuner.h" +#endif + using namespace SPTAG; +#ifdef RABITQ +namespace +{ +bool CreateBuildDirectory(const std::filesystem::path& p_final, + std::filesystem::path& p_build, + std::string& p_error) +{ + namespace fs = std::filesystem; + std::error_code error; + fs::create_directories(p_final.parent_path(), error); + if (error) { + p_error = "cannot create index parent directory: " + error.message(); + return false; + } + const auto nonce = std::chrono::high_resolution_clock::now().time_since_epoch().count(); + for (int attempt = 0; attempt < 100; ++attempt) { + p_build = p_final.parent_path() / + (p_final.filename().string() + ".rabitq-building-" + + std::to_string(nonce) + "-" + std::to_string(attempt)); + if (fs::create_directory(p_build, error)) { + return true; + } + if (error) { + p_error = "cannot create unique RaBitQ build directory: " + error.message(); + return false; + } + } + p_error = "cannot allocate a unique RaBitQ build directory"; + return false; +} + +bool PublishBuildDirectory(const std::filesystem::path& p_build, + const std::filesystem::path& p_final, + std::string& p_error) +{ + namespace fs = std::filesystem; + std::error_code error; + fs::path backup = p_build; + backup += ".previous"; + const bool hadPrevious = fs::exists(p_final, error); + if (error) { + p_error = "cannot inspect existing index directory: " + error.message(); + return false; + } + if (hadPrevious) { + fs::rename(p_final, backup, error); + if (error) { + p_error = "cannot preserve existing index directory: " + error.message(); + return false; + } + } + fs::rename(p_build, p_final, error); + if (error) { + const std::string publishError = error.message(); + if (hadPrevious) { + std::error_code restoreError; + fs::rename(backup, p_final, restoreError); + if (restoreError) { + p_error = "cannot publish new index (" + publishError + + ") or restore previous index (" + restoreError.message() + ")"; + return false; + } + } + std::error_code cleanupError; + fs::remove_all(p_build, cleanupError); + p_error = "cannot publish new index directory: " + publishError; + if (cleanupError) { + p_error += "; cannot remove failed staging directory: " + + cleanupError.message(); + } + return false; + } + if (hadPrevious) { + fs::remove_all(backup, error); + if (error) { + SPTAGLIB_LOG(Helper::LogLevel::LL_Warning, + "Published index but could not remove previous index backup %s: %s\n", + backup.string().c_str(), error.message().c_str()); + } + } + return true; +} + +bool FinalizeStagedConfig(const std::filesystem::path& p_build, + const std::filesystem::path& p_final, + std::string& p_error) +{ + const std::filesystem::path configPath = p_build / "indexloader.ini"; + std::ifstream input(configPath, std::ios::binary); + if (!input) { + p_error = "cannot open staged indexloader.ini"; + return false; + } + std::string config( + (std::istreambuf_iterator(input)), std::istreambuf_iterator()); + const std::string stagedRoot = p_build.string(); + const std::string finalRoot = p_final.string(); + std::size_t offset = 0; + int replacements = 0; + while ((offset = config.find(stagedRoot, offset)) != std::string::npos) { + config.replace(offset, stagedRoot.size(), finalRoot); + offset += finalRoot.size(); + ++replacements; + } + if (replacements == 0) { + p_error = "staged indexloader.ini does not reference its build directory"; + return false; + } + input.close(); + std::ofstream output(configPath, std::ios::binary | std::ios::trunc); + output.write(config.data(), static_cast(config.size())); + if (!output) { + p_error = "cannot finalize staged indexloader.ini"; + return false; + } + return true; +} +} // namespace +#endif + class BuilderOptions : public Helper::ReaderOptions { public: - BuilderOptions() : Helper::ReaderOptions(VectorValueType::Float, 0, VectorFileType::TXT, "|", 32) + BuilderOptions() + : Helper::ReaderOptions( + VectorValueType::Float, 0, VectorFileType::TXT, "|", 32, false, false) { AddRequiredOption(m_outputFolder, "-o", "--outputfolder", "Output folder."); AddRequiredOption(m_indexAlgoType, "-a", "--algo", "Index Algorithm type."); @@ -48,18 +178,6 @@ int main(int argc, char *argv[]) { exit(1); } - SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Set QuantizerFile = %s\n", options->m_quantizerFile.c_str()); - - auto indexBuilder = VectorIndex::CreateInstance(options->m_indexAlgoType, options->m_inputValueType); - if (!options->m_quantizerFile.empty()) - { - indexBuilder->LoadQuantizer(options->m_quantizerFile); - if (!indexBuilder->m_pQuantizer) - { - exit(1); - } - } - Helper::IniReader iniReader; if (!options->m_builderConfigFile.empty() && iniReader.LoadIniFile(options->m_builderConfigFile) != ErrorCode::Success) @@ -68,6 +186,8 @@ int main(int argc, char *argv[]) return -1; } + const bool configuredAutoTune = iniReader.DoesSectionExist("RaBitQAutoTune") && + iniReader.GetParameter("RaBitQAutoTune", "isExecute", false); for (int i = 1; i < argc; i++) { std::string param(argv[i]); @@ -79,6 +199,19 @@ int main(int argc, char *argv[]) std::string paramVal = param.substr(idx + 1); std::string sectionName; idx = paramName.find("."); + const bool tunerOverride = + idx != std::string::npos && + Helper::StrUtils::StrEqualIgnoreCase( + paramName.substr(0, idx).c_str(), "RaBitQAutoTune"); + if (idx != std::string::npos && (configuredAutoTune || tunerOverride)) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "RaBitQ auto-tuning parameters must come only from the INI; " + "command-line override %s is not allowed.\n", + param.c_str()); + return 1; + } if (idx != std::string::npos) { sectionName = paramName.substr(0, idx); @@ -89,6 +222,128 @@ int main(int argc, char *argv[]) paramVal.c_str()); } + std::string quantizerFile = options->m_quantizerFile; + VectorValueType builderValueType = options->m_inputValueType; + if (options->m_inputFiles.empty() && + iniReader.DoesParameterExist("Base", "ValueType")) + { + builderValueType = + iniReader.GetParameter("Base", "ValueType", VectorValueType::Undefined); + if (builderValueType == VectorValueType::Undefined) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, "Invalid [Base] ValueType.\n"); + return 1; + } + } + const bool autoTuneEnabled = configuredAutoTune; +#ifdef RABITQ + std::filesystem::path autoTuneBuildDirectory; + std::filesystem::path autoTuneFinalDirectory; +#endif + if (autoTuneEnabled) + { +#ifdef RABITQ + if (options->m_indexAlgoType != IndexAlgoType::SPANN) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "RaBitQ auto-tuning is supported only for SPANN index construction.\n"); + return 1; + } + if (!options->m_inputFiles.empty() || !options->m_quantizerFile.empty()) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "RaBitQ auto-tuning reads vectors and creates its quantizer from the INI; " + "--input and --quantizer are not allowed.\n"); + return 1; + } + const std::string configuredIndexDirectory = + iniReader.GetParameter("Base", "IndexDirectory", std::string()); + if (configuredIndexDirectory.empty()) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "[Base] IndexDirectory is required for RaBitQ auto-tuning.\n"); + return 1; + } + autoTuneFinalDirectory = + std::filesystem::absolute(configuredIndexDirectory).lexically_normal(); + const auto commandOutput = + std::filesystem::absolute(options->m_outputFolder).lexically_normal(); + if (commandOutput != autoTuneFinalDirectory) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "--outputfolder must match the authoritative [Base] IndexDirectory " + "when RaBitQ auto-tuning is enabled.\n"); + return 1; + } + std::string stagingError; + if (!CreateBuildDirectory( + autoTuneFinalDirectory, autoTuneBuildDirectory, stagingError)) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "%s\n", stagingError.c_str()); + return 1; + } + COMMON::RaBitQAutoTuneResult tuneResult; + std::string tuneError; + ErrorCode tuneStatus = ErrorCode::Fail; + try + { + tuneStatus = + COMMON::RaBitQAutoTuner::Run( + iniReader, autoTuneBuildDirectory.string(), tuneResult, tuneError); + } + catch (const std::exception& exception) + { + tuneError = exception.what(); + } + if (tuneStatus != ErrorCode::Success) + { + std::error_code cleanupError; + std::filesystem::remove_all(autoTuneBuildDirectory, cleanupError); + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "RaBitQ auto-tuning failed: %s\n", tuneError.c_str()); + return 1; + } + + builderValueType = VectorValueType::UInt8; + quantizerFile = tuneResult.quantizerPath; + options->m_inputFiles.clear(); + iniReader.SetParameter("Base", "ValueType", "UInt8"); + iniReader.SetParameter("Base", "Dim", std::to_string(tuneResult.codeDimension)); + iniReader.SetParameter("Base", "VectorPath", tuneResult.vectorPath); + iniReader.SetParameter("Base", "VectorType", "DEFAULT"); + iniReader.SetParameter("Base", "VectorSize", std::to_string(tuneResult.vectorCount)); + iniReader.SetParameter("Base", "QuantizerFilePath", tuneResult.quantizerPath); + iniReader.SetParameter( + "Base", "IndexDirectory", autoTuneBuildDirectory.string()); + iniReader.SetParameter("BuildSSDIndex", "EnableADC", "true"); +#else + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "[RaBitQAutoTune] isExecute=true requires a build configured with RABITQ=ON.\n"); + return 1; +#endif + } + + SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Set QuantizerFile = %s\n", quantizerFile.c_str()); + auto indexBuilder = VectorIndex::CreateInstance(options->m_indexAlgoType, builderValueType); + if (!indexBuilder) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Cannot create index builder.\n"); + return 1; + } + if (!quantizerFile.empty()) + { + if (indexBuilder->LoadQuantizer(quantizerFile) != ErrorCode::Success || + !indexBuilder->m_pQuantizer) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Cannot load quantizer file.\n"); + return 1; + } + } + std::string sections[] = {"Base", "SelectHead", "BuildHead", "BuildSSDIndex", "Index"}; for (int i = 0; i < 5; i++) { @@ -101,11 +356,17 @@ int main(int argc, char *argv[]) indexBuilder->SetParameter(iter.first.c_str(), iter.second.c_str(), sections[i]); } } - ErrorCode code; std::shared_ptr vecset; if (options->m_inputFiles != "") { + if (options->m_dimension <= 0) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "--dimension is required when indexbuilder reads --input directly.\n"); + return 1; + } auto vectorReader = Helper::VectorSetReader::CreateInstance(options); if (ErrorCode::Success != vectorReader->LoadFile(options->m_inputFiles)) { @@ -118,19 +379,87 @@ int main(int argc, char *argv[]) } else { - if (!options->m_quantizerFile.empty()) + if (!quantizerFile.empty()) { indexBuilder->SetQuantizerFileName( - options->m_quantizerFile.substr(options->m_quantizerFile.find_last_of("/\\") + 1)); + quantizerFile.substr(quantizerFile.find_last_of("/\\") + 1)); } code = indexBuilder->BuildIndex(options->m_normalized); } if (code == ErrorCode::Success) { - indexBuilder->SaveIndex(options->m_outputFolder); + for (const auto& parameter : iniReader.GetParameters("SearchSSDIndex")) + { + std::string name = parameter.first; + if (Helper::StrUtils::StrEqualIgnoreCase(name.c_str(), "isExecute") || + Helper::StrUtils::StrEqualIgnoreCase(name.c_str(), "BuildSsdIndex")) + { + continue; + } + if (Helper::StrUtils::StrEqualIgnoreCase(name.c_str(), "PostingPageLimit")) + { + name = "SearchPostingPageLimit"; + } + else if (Helper::StrUtils::StrEqualIgnoreCase(name.c_str(), "InternalResultNum")) + { + name = "SearchInternalResultNum"; + } + indexBuilder->SetParameter( + name.c_str(), parameter.second.c_str(), "BuildSSDIndex"); + } + std::string saveFolder = options->m_outputFolder; +#ifdef RABITQ + if (autoTuneEnabled) { + saveFolder = autoTuneBuildDirectory.string(); + } +#endif + code = indexBuilder->SaveIndex(saveFolder); + if (code != ErrorCode::Success) + { +#ifdef RABITQ + if (autoTuneEnabled) { + std::error_code cleanupError; + std::filesystem::remove_all(autoTuneBuildDirectory, cleanupError); + } +#endif + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Failed to save index.\n"); + return 1; + } +#ifdef RABITQ + if (autoTuneEnabled) + { + std::string finalizeError; + if (!FinalizeStagedConfig( + autoTuneBuildDirectory, autoTuneFinalDirectory, finalizeError)) + { + std::error_code cleanupError; + std::filesystem::remove_all(autoTuneBuildDirectory, cleanupError); + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "Failed to finalize auto-tuned index configuration: %s\n", + finalizeError.c_str()); + return 1; + } + std::string publishError; + if (!PublishBuildDirectory( + autoTuneBuildDirectory, autoTuneFinalDirectory, publishError)) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "Failed to publish auto-tuned index: %s\n", publishError.c_str()); + return 1; + } + } +#endif } else { +#ifdef RABITQ + if (autoTuneEnabled) { + std::error_code cleanupError; + std::filesystem::remove_all(autoTuneBuildDirectory, cleanupError); + } +#endif SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Failed to build index.\n"); exit(1); } diff --git a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini b/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini deleted file mode 100644 index c3438d655..000000000 --- a/Script_AE/iniFile/rabitq_auto_tune_sift1m.ini +++ /dev/null @@ -1,22 +0,0 @@ -[Base] -ValueType=Float -DistCalcMethod=L2 -Dim=128 -VectorPath=sift1m/sift_base.bin -VectorType=DEFAULT -QueryPath=sift1m/sift_query.bin -QueryType=DEFAULT -TruthPath=sift1m/sift_groundtruth_top1000.txt -TruthType=DEFAULT - -[RaBitQAutoTune] -isExecute=true -OutputDir=sift1m/rabitq_tuning -TargetRecall=0.95 - -[BuildSSDIndex] -NumberOfThreads=46 - -[SearchSSDIndex] -QueryCountLimit=10000 -ResultNum=100 diff --git a/Test/src/RaBitQQuantizerTest.cpp b/Test/src/RaBitQQuantizerTest.cpp index 56bd53b96..960834173 100644 --- a/Test/src/RaBitQQuantizerTest.cpp +++ b/Test/src/RaBitQQuantizerTest.cpp @@ -4,6 +4,7 @@ #include "inc/Test.h" #include "inc/Core/Common/QueryResultSet.h" +#include "inc/Core/Common/RaBitQAutoTuner.h" #include "inc/Core/Common/RaBitQQuantizer.h" #include "inc/Core/SPANN/Index.h" #include "inc/Core/VectorIndex.h" @@ -347,4 +348,225 @@ BOOST_AUTO_TEST_CASE(OfficialCompactRaBitQStoresRequestedBits) } } +BOOST_AUTO_TEST_CASE(SpannAppliesConfiguredADCWhenAttachingQuantizer) +{ + const auto raw = MakeRawVectors(); + auto quantizer = std::make_shared( + kDimension, kRaBitQBits, false); + BOOST_REQUIRE(quantizer->Train(raw) == ErrorCode::Success); + BOOST_CHECK(!quantizer->GetEnableADC()); + + auto index = VectorIndex::CreateInstance( + IndexAlgoType::SPANN, VectorValueType::UInt8); + BOOST_REQUIRE(index != nullptr); + index->SetParameter("EnableADC", "true", "BuildSSDIndex"); + index->SetQuantizer(quantizer); + BOOST_CHECK(quantizer->GetEnableADC()); +} + +BOOST_AUTO_TEST_CASE(RaBitQAutoTuneSelectsFirstQualifyingBit) +{ + std::vector evaluated; + int selected = 0; + float recall = 0.0F; + BOOST_REQUIRE( + COMMON::RaBitQAutoTuner::SelectMinimumBits( + 0.75F, + [&](int bits, float& value) { + evaluated.push_back(bits); + value = bits * 0.2F; + return ErrorCode::Success; + }, + selected, recall) == ErrorCode::Success); + BOOST_CHECK_EQUAL(selected, 4); + BOOST_CHECK_CLOSE(recall, 0.8F, 0.001F); + const std::vector expectedEvaluated = {1, 2, 3, 4}; + BOOST_CHECK_EQUAL_COLLECTIONS( + evaluated.begin(), evaluated.end(), + expectedEvaluated.begin(), expectedEvaluated.end()); + + BOOST_CHECK( + COMMON::RaBitQAutoTuner::SelectMinimumBits( + 1.0F, + [](int, float& value) { + value = 0.99F; + return ErrorCode::Success; + }, + selected, recall) == ErrorCode::Fail); + BOOST_CHECK_EQUAL(selected, 0); +} + +BOOST_AUTO_TEST_CASE(RaBitQAutoTuneUsesDeeperTruthPool) +{ + const std::vector> truth = { + {10, 11, 12}, {20, 21, 22}}; + std::string error; + BOOST_CHECK( + COMMON::RaBitQAutoTuner::ValidateTruth( + truth, 32, 2, 2, error) == ErrorCode::Success); + BOOST_CHECK_EQUAL( + COMMON::RaBitQAutoTuner::RecallAtK(truth[0], {12, 10}, 2), 0.5F); + + BOOST_CHECK( + COMMON::RaBitQAutoTuner::ValidateTruth( + {{0, 1}, {2, 3}}, 4, 2, 2, error) != ErrorCode::Success); + BOOST_CHECK(error.find("greater than ResultNum") != std::string::npos); + BOOST_CHECK( + COMMON::RaBitQAutoTuner::ValidateTruth( + {{0, 1, 4}, {1, 2, 3}}, 4, 2, 2, error) != ErrorCode::Success); +} + +BOOST_AUTO_TEST_CASE(RaBitQStreamingCentroidUsesEveryVector) +{ + const auto raw = MakeRawVectors(); + auto oneShot = std::make_shared( + kDimension, kRaBitQBits, false); + BOOST_REQUIRE(oneShot->Train(raw) == ErrorCode::Success); + + auto streamed = std::make_shared( + kDimension, 1, false); + BOOST_REQUIRE(streamed->BeginTraining() == ErrorCode::Success); + const SizeType boundaries[] = {0, 7, 41, kVectorCount}; + for (std::size_t batch = 0; batch + 1 < std::size(boundaries); ++batch) { + ByteArray bytes = ByteArray::Alloc( + sizeof(float) * static_cast(boundaries[batch + 1] - boundaries[batch]) * + kDimension); + std::memcpy( + bytes.Data(), raw->GetVector(boundaries[batch]), + bytes.Length()); + auto batchVectors = std::make_shared( + bytes, VectorValueType::Float, kDimension, + boundaries[batch + 1] - boundaries[batch]); + BOOST_REQUIRE(streamed->AddTrainingBatch(batchVectors) == ErrorCode::Success); + } + BOOST_REQUIRE(streamed->FinishTraining() == ErrorCode::Success); + const auto sharedCentroid = streamed->CreateWithBits(kRaBitQBits); + BOOST_REQUIRE(sharedCentroid != nullptr); + + std::vector expected(oneShot->GetNumSubvectors()); + std::vector actual(sharedCentroid->GetNumSubvectors()); + oneShot->QuantizeVector(raw->GetVector(kVectorCount - 1), expected.data(), false); + sharedCentroid->QuantizeVector(raw->GetVector(kVectorCount - 1), actual.data(), false); + BOOST_CHECK_EQUAL_COLLECTIONS( + expected.begin(), expected.end(), actual.begin(), actual.end()); +} + +BOOST_AUTO_TEST_CASE(RaBitQEncodedWidthMatchesSavedModel) +{ + constexpr DimensionType dimension = 70; + constexpr int bits = 5; + ByteArray bytes = ByteArray::Alloc(sizeof(float) * dimension * 2); + auto* values = reinterpret_cast(bytes.Data()); + for (DimensionType i = 0; i < dimension * 2; ++i) { + values[i] = static_cast(i) / 13.0F; + } + auto vectors = std::make_shared( + bytes, VectorValueType::Float, dimension, 2); + auto quantizer = std::make_shared( + dimension, bits, false); + BOOST_REQUIRE(quantizer->Train(vectors) == ErrorCode::Success); + BOOST_CHECK_EQUAL( + quantizer->GetNumSubvectors(), 128 * bits / 8 + 5 * sizeof(float)); + + const char* modelPath = "rabitq_width_model.bin"; + auto output = f_createIO(); + BOOST_REQUIRE(output->Initialize(modelPath, std::ios::out | std::ios::binary)); + BOOST_REQUIRE(quantizer->SaveQuantizer(output) == ErrorCode::Success); + output->ShutDown(); + auto input = f_createIO(); + BOOST_REQUIRE(input->Initialize(modelPath, std::ios::in | std::ios::binary)); + auto loaded = COMMON::IQuantizer::LoadIQuantizer(input); + BOOST_REQUIRE(loaded != nullptr); + BOOST_CHECK_EQUAL(loaded->GetNumSubvectors(), quantizer->GetNumSubvectors()); + std::remove(modelPath); +} + +BOOST_AUTO_TEST_CASE(RaBitQAutoTuneProducesNativeBuildHandoff) +{ + constexpr SizeType vectorCount = 6; + constexpr SizeType queryCount = 2; + constexpr DimensionType dimension = 8; + const char* basePath = "rabitq_auto_base.bin"; + const char* queryPath = "rabitq_auto_queries.bin"; + const char* truthPath = "rabitq_auto_truth.bin"; + const char* outputFolder = "rabitq_auto_handoff"; + std::filesystem::remove_all(outputFolder); + + auto writeVectors = [](const char* path, SizeType count, DimensionType dim, float offset) { + std::ofstream output(path, std::ios::binary | std::ios::trunc); + output.write(reinterpret_cast(&count), sizeof(count)); + output.write(reinterpret_cast(&dim), sizeof(dim)); + for (SizeType row = 0; row < count; ++row) { + for (DimensionType column = 0; column < dim; ++column) { + const float value = offset + row * 0.5F + column * 0.01F; + output.write(reinterpret_cast(&value), sizeof(value)); + } + } + BOOST_REQUIRE(output.good()); + }; + writeVectors(basePath, vectorCount, dimension, 0.0F); + writeVectors(queryPath, queryCount, dimension, 0.2F); + { + std::ofstream truth(truthPath, std::ios::binary | std::ios::trunc); + const DimensionType depth = 3; + const std::int32_t truthQueryCount = queryCount; + truth.write( + reinterpret_cast(&truthQueryCount), sizeof(truthQueryCount)); + truth.write(reinterpret_cast(&depth), sizeof(depth)); + const std::int32_t ids[] = {0, 1, 2, 1, 0, 2}; + truth.write(reinterpret_cast(ids), sizeof(ids)); + const float distances[] = {0.0F, 1.0F, 2.0F, 0.0F, 1.0F, 2.0F}; + truth.write(reinterpret_cast(distances), sizeof(distances)); + BOOST_REQUIRE(truth.good()); + } + + Helper::IniReader config; + config.SetParameter("Base", "ValueType", "Float"); + config.SetParameter("Base", "DistCalcMethod", "L2"); + config.SetParameter("Base", "Dim", std::to_string(dimension)); + config.SetParameter("Base", "VectorPath", basePath); + config.SetParameter("Base", "VectorType", "DEFAULT"); + config.SetParameter("Base", "QueryPath", queryPath); + config.SetParameter("Base", "QueryType", "DEFAULT"); + config.SetParameter("Base", "TruthPath", truthPath); + config.SetParameter("Base", "TruthType", "DEFAULT"); + config.SetParameter("SearchSSDIndex", "QueryCountLimit", std::to_string(queryCount)); + config.SetParameter("SearchSSDIndex", "ResultNum", "1"); + config.SetParameter("BuildSSDIndex", "NumberOfThreads", "2"); + config.SetParameter("RaBitQAutoTune", "isExecute", "true"); + config.SetParameter("RaBitQAutoTune", "TargetRecall", "0"); + + COMMON::RaBitQAutoTuneResult result; + std::string error; + BOOST_REQUIRE_MESSAGE( + COMMON::RaBitQAutoTuner::Run( + config, outputFolder, result, error) == ErrorCode::Success, + error); + BOOST_CHECK_EQUAL(result.selectedBits, 1); + BOOST_CHECK_EQUAL(result.vectorCount, vectorCount); + BOOST_REQUIRE(result.quantizer != nullptr); + BOOST_CHECK_EQUAL( + result.codeDimension, result.quantizer->GetNumSubvectors()); + + auto readerOptions = std::make_shared( + VectorValueType::UInt8, result.codeDimension, VectorFileType::DEFAULT); + auto encodedReader = Helper::VectorSetReader::CreateInstance(readerOptions); + BOOST_REQUIRE(encodedReader->LoadFile(result.vectorPath) == ErrorCode::Success); + const auto encoded = encodedReader->GetVectorSet(); + BOOST_CHECK_EQUAL(encoded->Count(), vectorCount); + BOOST_CHECK_EQUAL(encoded->Dimension(), result.quantizer->GetNumSubvectors()); + + auto modelInput = f_createIO(); + BOOST_REQUIRE(modelInput->Initialize( + result.quantizerPath.c_str(), std::ios::in | std::ios::binary)); + auto loaded = COMMON::IQuantizer::LoadIQuantizer(modelInput); + BOOST_REQUIRE(loaded != nullptr); + BOOST_CHECK_EQUAL(loaded->GetNumSubvectors(), encoded->Dimension()); + + std::remove(basePath); + std::remove(queryPath); + std::remove(truthPath); + std::filesystem::remove_all(outputFolder); +} + BOOST_AUTO_TEST_SUITE_END() diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index 1fc74ac16..bb7257f92 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -3,24 +3,17 @@ import tqdm import time from struct import pack, unpack, calcsize +from struct import pack, unpack, calcsize from typing import Dict, List import heapq import argparse import copy -import configparser -import json from operator import itemgetter import os import subprocess -import sys - -RABITQ_BATCH_SIZE = 1000000 -RABITQ_MIN_BITS = 1 -RABITQ_MAX_BITS = 8 -def get_cli_parser(): +def get_config(): parser = argparse.ArgumentParser(description ='implementation of nnsearch.') - parser.add_argument('--config', type = str, help='native INI containing [RaBitQAutoTune]') parser.add_argument('--data_file', default = 'traindata', type = str, help = 'binary data file') parser.add_argument('--query_file', default = 'query.tsv', type= str, help='query tsv file') parser.add_argument('--data_normalize', default = 0, type = int, help='normalize data vectors') @@ -47,116 +40,8 @@ def get_cli_parser(): parser.add_argument('--output_quan_vector_file', type = str, default = "", help='quantized vectors') parser.add_argument('--output_rec_vector_file', type = str, default = "", help = "reconstruct vectors") parser.add_argument('--quan_test', type = int, default = 0, help='compare with ground truth') - parser.add_argument('--rabitq_auto_tune', action = 'store_true', help='select the minimum RaBitQ storage bits before encoding') - parser.add_argument('--rabitq_target_recall', type = float, default = 0.95, help='minimum Recall@k for RaBitQ auto tuning') - parser.add_argument('--rabitq_tuning_result', type = str, default = 'rabitq_auto_tuning.json', help='auto-tuning result file under output_dir') - return parser - -def load_rabitq_auto_tune_ini(path): - config = configparser.ConfigParser(interpolation=None) - if not config.read(path): - raise ValueError(f'cannot read INI file: {path}') - if config.defaults(): - raise ValueError('[DEFAULT] parameters are forbidden in strict RaBitQ INI mode') - section_name = 'RaBitQAutoTune' - if not config.has_section(section_name): - raise ValueError(f'INI file is missing [{section_name}]') - if not config.has_section('Base'): - raise ValueError('INI file is missing [Base]') - section = config[section_name] - allowed_keys = { - 'isexecute', 'outputdir', 'querycount', 'targetrecall', - 'tuningresult', 'dataformat', 'task', - } - unknown_keys = set(section.keys()) - allowed_keys - if unknown_keys: - raise ValueError( - f'unknown [{section_name}] parameter(s): {", ".join(sorted(unknown_keys))}') - required_keys = ('OutputDir',) - missing_keys = [key for key in required_keys if not section.get(key)] - if missing_keys: - raise ValueError( - f'[{section_name}] is missing required parameter(s): {", ".join(missing_keys)}') - if not section.getboolean('isExecute', fallback=False): - raise ValueError(f'[{section_name}] isExecute must be true') - - base = config['Base'] - base_required_keys = ('VectorPath', 'QueryPath', 'TruthPath', 'Dim', 'ValueType', 'DistCalcMethod') - missing_base_keys = [key for key in base_required_keys if not base.get(key)] - if missing_base_keys: - raise ValueError( - f'[Base] is missing required parameter(s): {", ".join(missing_base_keys)}') - value_types = { - 'float': 'float32', - 'int8': 'int8', - 'uint8': 'uint8', - 'int16': 'int16', - } - value_type = base['ValueType'].lower() - if value_type not in value_types: - raise ValueError(f'unsupported [Base] ValueType for RaBitQ tuning: {base["ValueType"]}') - if value_type != 'float': - raise ValueError('RaBitQ tuning requires [Base] ValueType=Float') - if base['DistCalcMethod'].lower() != 'l2': - raise ValueError('RaBitQ tuning requires [Base] DistCalcMethod=L2') - - query_count = section.getint('QueryCount', fallback=None) - if query_count is None: - if not config.has_option('SearchSSDIndex', 'QueryCountLimit'): - raise ValueError( - f'query count must be set by [{section_name}] QueryCount or ' - '[SearchSSDIndex] QueryCountLimit') - query_count = config.getint('SearchSSDIndex', 'QueryCountLimit') - if not config.has_option('SearchSSDIndex', 'ResultNum'): - raise ValueError('[SearchSSDIndex] ResultNum is required for RaBitQ tuning') - if not config.has_option('BuildSSDIndex', 'NumberOfThreads'): - raise ValueError('[BuildSSDIndex] NumberOfThreads is required for RaBitQ tuning') - - return argparse.Namespace( - config=path, - data_file=base['VectorPath'], - query_file=base['QueryPath'], - data_normalize=0, - query_normalize=0, - data_type=value_types[value_type], - target_type=value_types[value_type], - k=config.getint('SearchSSDIndex', 'ResultNum'), - dim=base.getint('Dim'), - B=RABITQ_BATCH_SIZE, - Q=query_count, - S=1000, - D=base['DistCalcMethod'], - output_truth=base['TruthPath'], - data_format=section.get('DataFormat', fallback='DEFAULT'), - task=section.getint('Task', fallback=0), - log_dir='', - T=config.getint('BuildSSDIndex', 'NumberOfThreads'), - train_samples=None, - quan_type='rabitq', - quan_dim=-1, - output_dir=section['OutputDir'], - output_quantizer='quantizer.bin', - output_quan_vector_file='', - output_rec_vector_file='', - quan_test=1, - rabitq_auto_tune=True, - rabitq_target_recall=section.getfloat('TargetRecall', fallback=0.95), - rabitq_tuning_result=section.get( - 'TuningResult', fallback='rabitq_auto_tuning.json'), - ) - -def get_config(argv=None): - if argv is None: - argv = sys.argv[1:] - config_probe = argparse.ArgumentParser(add_help=False) - config_probe.add_argument('--config') - config_args, remaining = config_probe.parse_known_args(argv) - if config_args.config is not None: - if remaining: - raise ValueError( - 'INI mode accepts only --config; command-line parameter overrides are forbidden') - return load_rabitq_auto_tune_ini(config_args.config) - return get_cli_parser().parse_args(argv) + args = parser.parse_args() + return args def is_binary_vector_file(filename): suffixes = ('.bin', '.fbin', '.u8bin', '.i8bin') @@ -183,10 +68,7 @@ def __init__(self, filename, featuredim, batchsize, normalize, datatype, targett self.isbinary = False self.type = self.mytype - if batchsize <= 0: - batchsize = R - else: - batchsize = min(batchsize, R) + if batchsize <= 0: batchsize = R self.query = np.zeros([batchsize, self.featuredim], dtype=self.mytype) self.normalize = normalize @@ -412,196 +294,110 @@ def evaluate(retrieve_results: List[List[int]], return MRR, Recall -def sptag_rabitq_storage_bytes(dim, bits): - if dim <= 0: - raise ValueError('RaBitQ dimension must be positive') - if bits < 1 or bits > 8: - raise ValueError(f'RaBitQ bits must be in [1, 8], got {bits}') - padded_dimension = ((dim + 63) // 64) * 64 - return padded_dimension * bits // 8 + 5 * np.dtype(np.float32).itemsize - -def rabitq_bits_from_quantized_dimension(dim, quan_dim): - total_bits = quan_dim * 8 - if dim <= 0 or quan_dim <= 0 or total_bits % dim != 0: - raise ValueError(f'quan_dim={quan_dim} does not represent an integral RaBitQ bit count for dimension {dim}') - bits = total_bits // dim - if bits < 1 or bits > 8: - raise ValueError(f'RaBitQ bits must be in [1, 8], got {bits}') - return bits - -def load_ground_truth(path, query_count): - rows = [] - candidate_count = None - with open(path, 'r') as truth_file: - for query_id in range(query_count): - line = truth_file.readline() - if not line: - raise ValueError(f'ground truth contains only {query_id} queries, expected {query_count}') - neighbors = line.strip().split() - if candidate_count is None: - candidate_count = len(neighbors) - if candidate_count == 0: - raise ValueError('ground truth query 0 contains no neighbors') - if len(neighbors) != candidate_count: - raise ValueError( - f'ground truth query {query_id} contains {len(neighbors)} neighbors, ' - f'expected exactly {candidate_count}') - row = [int(neighbor) for neighbor in neighbors] - if len(set(row)) != candidate_count: - raise ValueError(f'ground truth query {query_id} contains duplicate neighbor IDs') - rows.append(row) - return np.asarray(rows, dtype=np.int64) - -def reranking_recall_at_k(faiss, faiss_index, queries, candidates, topk): - if len(queries) != len(candidates): - raise ValueError('query and ground-truth counts differ') - if len(queries) == 0: - raise ValueError('at least one query is required') - if topk <= 0: - raise ValueError('ResultNum must be positive') - candidate_count = candidates.shape[1] - if candidate_count <= topk: - raise ValueError( - f'ground-truth candidate depth must exceed ResultNum: {candidate_count} <= {topk}') - if np.any(candidates < 0) or np.any(candidates >= faiss_index.ntotal): - raise ValueError('ground truth contains a vector ID outside the base data') - - recall_sum = 0.0 - for query_id in tqdm.tqdm(range(len(queries))): - candidate_ids = np.ascontiguousarray(candidates[query_id]) - parameters = faiss.SearchParameters() - parameters.sel = faiss.IDSelectorArray( - candidate_count, faiss.swig_ptr(candidate_ids)) - _, results = faiss_index.search( - np.ascontiguousarray(queries[query_id:query_id + 1]), topk, - params=parameters) - expected = set(int(candidate) for candidate in candidate_ids[:topk]) - recall_sum += len(expected.intersection(int(candidate) for candidate in results[0])) / topk - return recall_sum / len(queries) - -def create_rabitq_index(faiss, dim, bits, centroid): - faiss_index = faiss.index_factory(dim, f"RaBitQ{bits}", faiss.METRIC_L2) - faiss_index.train(np.ascontiguousarray(centroid.reshape(1, dim), dtype=np.float32)) - return faiss_index - -def compute_streaming_centroid(args): - datareader = DataReader( - args.data_file, args.dim, RABITQ_BATCH_SIZE, - args.data_normalize, args.data_type, args.target_type) - accumulator = np.zeros(args.dim, dtype=np.float64) - total = 0 - while True: - num_data, data = datareader.readbatch() - if num_data == 0: - break - accumulator += np.sum(data, axis=0, dtype=np.float64) - total += num_data - datareader.close() - if total == 0: - raise ValueError('RaBitQ input data is empty') - return (accumulator / total).astype(np.float32), total - -def add_rabitq_data(args, faiss_index): - datareader = DataReader( - args.data_file, args.dim, args.B, args.data_normalize, args.data_type, args.target_type) - total = 0 - while True: - num_data, data = datareader.readbatch() - if num_data == 0: - break - faiss_index.add(data) - total += num_data - datareader.close() - if total == 0: - raise ValueError('RaBitQ input data is empty') - return total - -def tune_rabitq_bits(args, faiss, centroid, centroid_vector_count, queries, ground_truth_candidates): - if not 0.0 < args.rabitq_target_recall <= 1.0: - raise ValueError('rabitq_target_recall must be in (0, 1]') - - trials = [] - for bits in range(RABITQ_MIN_BITS, RABITQ_MAX_BITS + 1): - print(f'Auto tuning RaBitQ{bits} for Recall@{args.k} >= {args.rabitq_target_recall:.6f}') - candidate = create_rabitq_index(faiss, args.dim, bits, centroid) - data_count = add_rabitq_data(args, candidate) - if data_count != centroid_vector_count: - raise RuntimeError( - f'base vector count changed during tuning: {centroid_vector_count} != {data_count}') - recall = reranking_recall_at_k( - faiss, candidate, queries, ground_truth_candidates, args.k) - trials.append({'bits': bits, 'recall': recall}) - print(f'RaBitQ{bits} Recall@{args.k}: {recall:.6f}') - if recall >= args.rabitq_target_recall: - return bits, data_count, trials - - measured = ', '.join(f'{trial["bits"]}-bit={trial["recall"]:.6f}' for trial in trials) - raise RuntimeError( - f'No RaBitQ bit count in [{RABITQ_MIN_BITS}, {RABITQ_MAX_BITS}] ' - f'reached Recall@{args.k} >= {args.rabitq_target_recall:.6f}; {measured}') - def train_rabitq(args): import faiss output_dir = args.output_dir - if args.D != 'L2': - raise ValueError('SPTAG global RaBitQ tuning requires [Base] DistCalcMethod=L2') - if args.Q <= 0 or args.k <= 0: - raise ValueError('Q and k must be positive') - if len(args.output_quan_vector_file) > 0 or len(args.output_rec_vector_file) > 0: - raise ValueError( - 'RaBitQ vectors must be generated by the native SPTAG quantizer after tuning; ' - 'do not use output_quan_vector_file or output_rec_vector_file') - centroid, centroid_vector_count = compute_streaming_centroid(args) - print(f'train RaBitQ using the centroid of all {centroid_vector_count} base vectors ...') - if args.T <= 0: - raise ValueError('NumberOfThreads must be positive') + if args.train_samples > args.B: args.train_samples = args.B + + datareader = DataReader(args.data_file, args.dim, args.B, args.data_normalize, args.data_type, args.target_type) + + print (f'train RabitQ using {args.train_samples} samples ...') + + numData, data = datareader.readbatch() + faiss.omp_set_num_threads(args.T) - if args.rabitq_auto_tune: - if args.quan_test <= 0: - raise ValueError('rabitq_auto_tune requires quan_test > 0 and a pre-generated ground truth') - queryreader = DataReader( - args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) - num_query, queries = queryreader.readbatch() - queryreader.close() - if num_query != args.Q: - raise ValueError(f'query file contains {num_query} queries, but configured Q is {args.Q}') - ground_truth_candidates = load_ground_truth( - args.output_truth, num_query) - bits, data_count, trials = tune_rabitq_bits( - args, faiss, centroid, centroid_vector_count, queries, ground_truth_candidates) - result = { - 'selected_bits': bits, - 'native_quantizer_qd': bits, - 'storage_bytes_per_vector': sptag_rabitq_storage_bytes(args.dim, bits), - 'target_recall': args.rabitq_target_recall, - 'recall_at': args.k, - 'rerank_candidate_count': ground_truth_candidates.shape[1], - 'query_count': num_query, - 'data_count': data_count, - 'centroid_vector_count': centroid_vector_count, - 'trials': trials, - } - result_path = os.path.join(output_dir, args.rabitq_tuning_result) - temporary_result_path = result_path + '.tmp' - with open(temporary_result_path, 'w') as result_file: - json.dump(result, result_file, indent=2) - result_file.write('\n') - os.replace(temporary_result_path, result_path) - print(f'Selected RaBitQ storage bits: {bits}; result: {result_path}') - else: - bits = rabitq_bits_from_quantized_dimension(args.dim, args.quan_dim) - faiss_index = create_rabitq_index(faiss, args.dim, bits, centroid) - if args.quan_test > 0: - add_rabitq_data(args, faiss_index) + nbytes = int(args.quan_dim * 8 // args.dim) + print (f'nbytes:{nbytes}') + faiss_index = faiss.index_factory(len(data[0]), f"RaBitQ{nbytes}", faiss.METRIC_L2) + print('Training the index with doc embeddings') + + faiss_index.train(data[0:args.train_samples]) + + rtype = np.uint8(0) + if args.data_type == 'uint8': + rtype = np.uint8(1) + elif args.data_type == 'int16': + rtype = np.uint8(2) + elif args.data_type == 'float32': + rtype = np.uint8(3) + + ivf_index = faiss.downcast_index(faiss_index) + #centroid_embedings = faiss.vector_to_array(ivf_index.pq.centroids) + #codebooks = centroid_embedings.reshape(ivf_index.pq.M, ivf_index.pq.ksub, ivf_index.pq.dsub) + #print ('codebooks shape:') + #print (codebooks.shape) + + #codebooks = codebooks.astype(np.float32) + #with open(os.path.join(output_dir, args.output_quantizer + '.' + str(args.task)),'wb') as f: + # f.write(pack('B', 1)) + # f.write(pack('B', rtype)) + # f.write(pack('i', codebooks.shape[0])) + # f.write(pack('i', codebooks.shape[1])) + # f.write(pack('i', codebooks.shape[2])) + # f.write(codebooks.tobytes()) + + if args.quan_test == 0 and len(args.output_quan_vector_file) == 0 and len(args.output_rec_vector_file) == 0: + os.rename(args.output_truth, os.path.join(output_dir, 'truth.txt' + '.' + str(args.task))) + return + + if len(args.output_quan_vector_file) > 0: + fquan = open(os.path.join(output_dir, args.output_quan_vector_file + '.' + str(args.task) + '.tmp'), 'wb') + fquan.write(pack('i', 0)) + fquan.write(pack('i', args.quan_dim)) + + if len(args.output_rec_vector_file) > 0: + frec = open(os.path.join(output_dir, args.output_rec_vector_file + '.' + str(args.task) + '.tmp'), 'wb') + frec.write(pack('i', 0)) + frec.write(pack('i', data.shape[1])) + + writeitems = 0 + while numData > 0: + if args.quan_test > 0: faiss_index.add(data) + + if len(args.output_quan_vector_file) > 0: + codes = ivf_index.pq.compute_codes(data) + print ('codes shape:') + print (codes.shape) + fquan.write(codes.tobytes()) + + if len(args.output_rec_vector_file) > 0: + reconstructed = ivf_index.pq.decode(codes).astype(args.data_type) + frec.write(reconstructed.tobytes()) + + writeitems += numData + numData, data = datareader.readbatch() + + datareader.close() + + if len(args.output_quan_vector_file) > 0: + p = fquan.tell() + fquan.seek(0) + fquan.write(pack('i', writeitems)) + fquan.seek(p) + fquan.close() + if os.path.exists(os.path.join(output_dir, args.output_quan_vector_file + '.' + str(args.task))): + os.remove(os.path.join(output_dir, args.output_quan_vector_file + '.' + str(args.task))) + os.rename(os.path.join(output_dir, args.output_quan_vector_file + '.' + str(args.task) + '.tmp'), os.path.join(output_dir, args.output_quan_vector_file + '.' + str(args.task))) + if len(args.output_rec_vector_file) > 0: + p = frec.tell() + frec.seek(0) + frec.write(pack('i', writeitems)) + frec.seek(p) + frec.close() + if os.path.exists(os.path.join(output_dir, args.output_rec_vector_file + '.' + str(args.task))): + os.remove(os.path.join(output_dir, args.output_rec_vector_file + '.' + str(args.task))) + os.rename(os.path.join(output_dir, args.output_rec_vector_file + '.' + str(args.task) + '.tmp'), os.path.join(output_dir, args.output_rec_vector_file + '.' + str(args.task))) + + os.rename(args.output_truth, os.path.join(output_dir, 'truth.txt' + '.' + str(args.task))) - if args.quan_test > 0 and not args.rabitq_auto_tune: + if args.quan_test > 0: queryreader = DataReader(args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) numQuery, query = queryreader.readbatch() qid2ground_truths = {} - f = open(args.output_truth, 'r') + f = open(os.path.join(output_dir, 'truth.txt.' + str(args.task)), 'r') for i in range(numQuery): items = f.readline()[0:-1].strip().split(' ') qid2ground_truths[i] = set([int(gt) for gt in items]) diff --git a/Tools/OPQ/README.md b/Tools/OPQ/README.md index 3bfc52418..6ed9507f7 100644 --- a/Tools/OPQ/README.md +++ b/Tools/OPQ/README.md @@ -12,37 +12,5 @@ ## Example python3 OPQ_gpu_train_infer.py --data_file perftest_vector.bin --query_file perftest_query.bin --task 0 --data_type float32 --k 5 --dim 64 --B 1000000 --Q 1000 --D L2 --data_format DEFAULT --T 20 --train_samples 1000000 --quan_type opq --quan_dim 32 --output_quantizer quantizer.bin - -## RaBitQ storage-bit auto tuning - -`--rabitq_auto_tune` runs before vector encoding or index construction. It evaluates -RaBitQ bit counts in ascending order and selects the first (therefore minimum) -count whose reranking Recall reaches `--rabitq_target_recall`. Every candidate -uses the centroid streamed over the complete base file, the same configured -`--Q` queries, and the full pre-generated ground-truth candidate pool. In INI mode, -`[SearchSSDIndex] ResultNum` supplies K: the candidate pool is reranked by -RaBitQ distance and its first K IDs are compared with the exact first K. - -The command fails instead of silently choosing a bit count when the configured -query/ground-truth count is unavailable or no candidate in RaBitQ's fixed -1-through-8-bit range reaches the target. The selected storage -bit count, actual SPTAG byte width (including padded dimensions and five Float -factors), measured recalls, and sample counts are written atomically to -`/rabitq_auto_tuning.json`. Use `native_quantizer_qd` as `-qd` for the -native SPTAG `quantizer` command that generates the quantizer and vectors consumed -by the subsequent index build. The tuner deliberately does not emit Faiss codes: -their persisted layout is not the SPTAG global RaBitQ format. - -The same parameters can be supplied exclusively through the -`[RaBitQAutoTune]` section of an INI: - -```bash -python3 Tools/OPQ/OPQ_gpu_train_infer.py \ - --config Script_AE/iniFile/rabitq_auto_tune_sift1m.ini -``` - -In INI mode no additional CLI parameters are accepted. `QueryCount` defaults to -`[SearchSSDIndex] QueryCountLimit`, keeping the configured query count in one -place. Faiss/OpenMP reuses `[BuildSSDIndex] NumberOfThreads`. Global RaBitQ -tuning requires `[Base] ValueType=Float` and `DistCalcMethod=L2`; there are no -separate data/query normalization settings. \ No newline at end of file +python3 OPQ_gpu_train_infer.py --data_file laion_5M.bin --query_file laion_test.bin --data_type float32 --target_type float32 --k 100 --dim 768 --B 1000000 --Q 10000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 96 --quan_test 1 +python3 OPQ_gpu_train_infer.py --data_file openai_5M.bin --query_file openai_test.bin --data_type float32 --target_type float32 --k 100 --dim 1536 --B 1000000 --Q 10000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 192 --quan_test 1 --data_normalize 1 --query_normalize 1 \ No newline at end of file diff --git a/Tools/OPQ/test_rabitq_auto_tuning.py b/Tools/OPQ/test_rabitq_auto_tuning.py deleted file mode 100644 index a1362d339..000000000 --- a/Tools/OPQ/test_rabitq_auto_tuning.py +++ /dev/null @@ -1,239 +0,0 @@ -import importlib.util -from pathlib import Path -import tempfile -import unittest - -import numpy as np - - -MODULE_PATH = Path(__file__).with_name('OPQ_gpu_train_infer.py') -SPEC = importlib.util.spec_from_file_location('opq_gpu_train_infer', MODULE_PATH) -MODULE = importlib.util.module_from_spec(SPEC) -SPEC.loader.exec_module(MODULE) - - -class RaBitQAutoTuningTest(unittest.TestCase): - def test_data_reader_accepts_fbin(self): - with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / 'vectors.fbin' - with path.open('wb') as output: - np.asarray([2, 3], dtype=np.int32).tofile(output) - np.arange(6, dtype=np.float32).tofile(output) - reader = MODULE.DataReader(str(path), 3, 2, 0, 'float32') - count, vectors = reader.readbatch() - reader.close() - self.assertEqual(2, count) - np.testing.assert_array_equal(vectors, np.arange(6, dtype=np.float32).reshape(2, 3)) - - def test_binary_reader_accepts_numbered_shards(self): - self.assertTrue(MODULE.is_binary_vector_file('vectors.bin.0')) - self.assertTrue(MODULE.is_binary_vector_file('vectors.fbin.12')) - self.assertFalse(MODULE.is_binary_vector_file('vectors.txt.0')) - - def test_sptag_storage_bytes_include_padding_and_factors(self): - self.assertEqual(68, MODULE.sptag_rabitq_storage_bytes(128, 3)) - self.assertEqual(68, MODULE.sptag_rabitq_storage_bytes(127, 3)) - self.assertEqual(92, MODULE.sptag_rabitq_storage_bytes(129, 3)) - - def test_rejects_invalid_storage_bits(self): - with self.assertRaises(ValueError): - MODULE.sptag_rabitq_storage_bytes(128, 0) - - def test_load_ground_truth_honors_query_count_and_topk(self): - with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / 'truth.txt' - path.write_text('1 2 3\n4 5 6\n', encoding='ascii') - truths = MODULE.load_ground_truth(path, 2) - np.testing.assert_array_equal(truths, [[1, 2, 3], [4, 5, 6]]) - with self.assertRaises(ValueError): - MODULE.load_ground_truth(path, 3) - - def test_inferred_ground_truth_topk_requires_uniform_unique_rows(self): - with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / 'truth.txt' - path.write_text('1 2\n3 4 5\n', encoding='ascii') - with self.assertRaises(ValueError): - MODULE.load_ground_truth(path, 2) - path.write_text('1 2\n3 3\n', encoding='ascii') - with self.assertRaises(ValueError): - MODULE.load_ground_truth(path, 2) - - def test_reranking_recall_uses_result_num_with_deeper_candidates(self): - import faiss - base = np.asarray([[0.0], [1.0], [2.0], [3.0]], dtype=np.float32) - queries = np.asarray([[0.1]], dtype=np.float32) - index = faiss.IndexFlatL2(1) - index.add(base) - candidates = np.asarray([[0, 1, 2, 3]], dtype=np.int64) - recall = MODULE.reranking_recall_at_k( - faiss, index, queries, candidates, 2) - self.assertEqual(1.0, recall) - with self.assertRaises(ValueError): - MODULE.reranking_recall_at_k( - faiss, index, queries, candidates[:, :2], 2) - - def test_ini_is_authoritative_and_reuses_query_count_limit(self): - with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / 'build.ini' - path.write_text( - '[Base]\n' - 'ValueType=Float\n' - 'DistCalcMethod=L2\n' - 'Dim=128\n' - 'VectorPath=base.bin\n' - 'QueryPath=query.bin\n' - 'TruthPath=truth.txt\n' - '\n' - '[RaBitQAutoTune]\n' - 'isExecute=true\n' - 'OutputDir=tuning\n' - 'TargetRecall=0.97\n' - '\n' - '[BuildSSDIndex]\n' - 'NumberOfThreads=46\n' - '\n' - '[SearchSSDIndex]\n' - 'QueryCountLimit=10000\n' - 'ResultNum=100\n', - encoding='ascii') - args = MODULE.get_config(['--config', str(path)]) - self.assertTrue(args.rabitq_auto_tune) - self.assertEqual(10000, args.Q) - self.assertEqual(100, args.k) - self.assertEqual(46, args.T) - self.assertEqual('float32', args.target_type) - self.assertIsNone(args.train_samples) - self.assertEqual(0.97, args.rabitq_target_recall) - self.assertEqual('base.bin', args.data_file) - self.assertEqual('query.bin', args.query_file) - self.assertEqual('truth.txt', args.output_truth) - - with self.assertRaises(ValueError): - MODULE.get_config(['--config', str(path), '--Q', '1']) - - def test_streaming_centroid_uses_all_base_vectors(self): - with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / 'vectors.bin' - vectors = np.arange(20, dtype=np.float32).reshape(5, 4) - with path.open('wb') as output: - np.asarray([5, 4], dtype=np.int32).tofile(output) - vectors.tofile(output) - args = type('Args', (), { - 'data_file': str(path), - 'dim': 4, - 'data_normalize': 0, - 'data_type': 'float32', - 'target_type': 'float32', - })() - centroid, count = MODULE.compute_streaming_centroid(args) - self.assertEqual(5, count) - np.testing.assert_array_equal(centroid, vectors.mean(axis=0)) - - def test_ini_rejects_unknown_parameters(self): - with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / 'build.ini' - path.write_text( - '[Base]\n' - 'ValueType=Float\n' - 'DistCalcMethod=L2\n' - 'Dim=128\n' - 'VectorPath=base.bin\n' - 'QueryPath=query.bin\n' - 'TruthPath=truth.txt\n' - '\n' - '[RaBitQAutoTune]\n' - 'isExecute=true\n' - 'OutputDir=tuning\n' - 'QueryCount=10000\n' - 'TypoTargetRecal=0.95\n', - encoding='ascii') - with self.assertRaises(ValueError): - MODULE.get_config(['--config', str(path)]) - - def test_ini_rejects_configurable_bit_range(self): - with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / 'build.ini' - path.write_text( - '[Base]\n' - 'ValueType=Float\n' - 'DistCalcMethod=L2\n' - 'Dim=128\n' - 'VectorPath=base.bin\n' - 'QueryPath=query.bin\n' - 'TruthPath=truth.txt\n' - '\n' - '[RaBitQAutoTune]\n' - 'isExecute=true\n' - 'OutputDir=tuning\n' - 'MinBits=2\n' - 'MaxBits=7\n' - '\n' - '[SearchSSDIndex]\n' - 'QueryCountLimit=10000\n' - 'ResultNum=100\n', - encoding='ascii') - with self.assertRaises(ValueError): - MODULE.get_config(['--config', str(path)]) - - def test_ini_rejects_normalization_and_non_l2_metric(self): - with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / 'build.ini' - config = ( - '[Base]\n' - 'ValueType=Float\n' - 'DistCalcMethod=L2\n' - 'Dim=128\n' - 'VectorPath=base.bin\n' - 'QueryPath=query.bin\n' - 'TruthPath=truth.txt\n' - '\n' - '[RaBitQAutoTune]\n' - 'isExecute=true\n' - 'OutputDir=tuning\n' - 'DataNormalize=1\n' - '\n' - '[BuildSSDIndex]\n' - 'NumberOfThreads=46\n' - '\n' - '[SearchSSDIndex]\n' - 'QueryCountLimit=10000\n' - 'ResultNum=100\n') - path.write_text(config, encoding='ascii') - with self.assertRaises(ValueError): - MODULE.get_config(['--config', str(path)]) - path.write_text( - config.replace('DistCalcMethod=L2', 'DistCalcMethod=Cosine') - .replace('DataNormalize=1\n', ''), - encoding='ascii') - with self.assertRaises(ValueError): - MODULE.get_config(['--config', str(path)]) - - def test_ini_rejects_inherited_defaults(self): - with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / 'build.ini' - path.write_text( - '[DEFAULT]\n' - 'QueryCount=1\n' - '\n' - '[Base]\n' - 'ValueType=Float\n' - 'DistCalcMethod=L2\n' - 'Dim=128\n' - 'VectorPath=base.bin\n' - 'QueryPath=query.bin\n' - 'TruthPath=truth.txt\n' - '\n' - '[RaBitQAutoTune]\n' - 'isExecute=true\n' - 'OutputDir=tuning\n' - '\n' - '[SearchSSDIndex]\n' - 'QueryCountLimit=10000\n' - 'ResultNum=100\n', - encoding='ascii') - with self.assertRaises(ValueError): - MODULE.get_config(['--config', str(path)]) - - -if __name__ == '__main__': - unittest.main() diff --git a/docs/GettingStart.md b/docs/GettingStart.md index 7b482d0c1..57d690f39 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -241,100 +241,56 @@ SearchPostingPageLimit=12 ### **Global RaBitQ Quantizer** -RaBitQ is a global `IQuantizer`, not a SPANN posting quantizer. Train the -official model and encode base vectors with `Release/quantizer`, then -use the generated model through `QuantizerFilePath` in the normal SPANN -workflow. For 128-dimensional SIFT vectors, the encoded `UInt8` vectors use -`Dim=68` at 3 bits (48 compact code bytes plus five Float factors). - -To select the minimum storage bit count before building the index, add a -`[RaBitQAutoTune]` section to a pre-build INI. `[Base]` is the single source for -the raw vector, query, ground-truth, dimension, value type, and distance. The -tuner uses a pre-generated exact top-`k` ground truth. Vector inputs use SPTAG -`DEFAULT` binary files (`int32 count`, `int32 dimension`, then vector payload), -and the text ground-truth file contains one space-separated neighbor-ID list -per query: +RaBitQ is a global `IQuantizer`, not a SPANN posting quantizer. A build +configured with `-DRABITQ=ON` can select its bit width, train the official +model, encode the complete base set, and continue SPANN construction in one +`indexbuilder` invocation. Add only this tuner-specific section to the normal +native build INI: ```ini -[Base] -ValueType=Float -DistCalcMethod=L2 -Dim=128 -VectorPath=sift1m/sift_base.bin -VectorType=DEFAULT -QueryPath=sift1m/sift_query.bin -QueryType=DEFAULT -TruthPath=sift1m/sift_groundtruth_top1000.txt -TruthType=DEFAULT - [RaBitQAutoTune] isExecute=true -OutputDir=sift1m/rabitq_tuning TargetRecall=0.95 - -[BuildSSDIndex] -NumberOfThreads=46 - -[SearchSSDIndex] -QueryCountLimit=10000 -ResultNum=100 ``` -Run the pre-build tuning stage with only the INI path: - -```bash -python3 Tools/OPQ/OPQ_gpu_train_infer.py \ - --config Script_AE/iniFile/rabitq_auto_tune_sift1m.ini -``` - -The tuner evaluates bit counts in ascending order using exactly the configured -`SearchSSDIndex.QueryCountLimit` queries and selects the first bit count meeting -the target Recall. The supported 1 through 8 bit range is fixed by RaBitQ and is -always evaluated from the minimum upward. `SearchSSDIndex.ResultNum` is the expected result K. The -ground-truth width is a separate, deeper reranking candidate pool: with -`ResultNum=100` and 1,000 exact IDs per query, each RaBitQ candidate reranks -those 1,000 IDs and is evaluated as `Recall@100` against the exact first 100. -The candidate depth must be strictly greater than `ResultNum`; this prevents a -top100 ground truth from being used to evaluate K=100. All consumed -ground-truth rows must have exactly the same width with no duplicate IDs. When -`--config` is used, additional command-line parameters are rejected, so they -cannot override the INI. A -`RaBitQAutoTune.QueryCount` may be set explicitly only when the tuning query -count intentionally differs from `SearchSSDIndex.QueryCountLimit`. Strict mode -also rejects unknown `[RaBitQAutoTune]` keys and inherited `[DEFAULT]` values. -The tuner reads the type from `[Base].ValueType`; global RaBitQ currently -requires `Float` and `[Base].DistCalcMethod=L2`. It does not independently -normalize data or queries. It computes the centroid over the complete base file -in bounded one-million-vector batches, so there is no training-sample parameter. -Faiss/OpenMP uses the existing -`BuildSSDIndex.NumberOfThreads` value; there is no separate tuner thread setting. -It fails if fewer queries/ground-truth rows are available or no candidate -qualifies. The ground truth is not moved or modified. Results are written -atomically to `sift1m/rabitq_tuning/rabitq_auto_tuning.json`; use -`native_quantizer_qd` for `-qd` and `storage_bytes_per_vector` for the SPANN -`[Base] Dim`: - -```bash -BITS=$(python3 -c \ - 'import json; print(json.load(open("sift1m/rabitq_tuning/rabitq_auto_tuning.json"))["native_quantizer_qd"])') -STORAGE_DIM=$(python3 -c \ - 'import json; print(json.load(open("sift1m/rabitq_tuning/rabitq_auto_tuning.json"))["storage_bytes_per_vector"])') - -Release/quantizer \ - -d 128 -v Float -f XVEC \ - -i sift1m/sift_base.fvecs \ - -o "sift1m/sift_base.rabitq${BITS}.u8bin" \ - -oq "sift1m/official_rabitq${BITS}_global.bin" \ - -qt RaBitQQuantizer -qd "$BITS" -ts 1000000 -``` - -The pre-build INI continues to describe raw Float data. Set `VectorPath` and -`QuantizerFilePath` to the generated quantized files and set -`ValueType=UInt8`, `Dim=$STORAGE_DIM` in the subsequent build INI before -starting SPANN construction. Do not use a manually supplied bit count to -override an auto-tuning result. - -For reference, the fixed 3-bit SIFT1M command is: +The tuner takes all shared settings from the existing INI: + +* `[Base]` supplies `VectorPath`, `QueryPath`, `TruthPath`, `VectorType`, + `QueryType`, `TruthType`, `ValueType`, `Dim`, and `DistCalcMethod`. +* `[SearchSSDIndex]` supplies `QueryCountLimit` and `ResultNum`. +* `[BuildSSDIndex]` supplies `NumberOfThreads`. + +The raw base and queries must be `Float` with `DistCalcMethod=L2`. Exactly +`QueryCountLimit` queries are evaluated. Each truth row must contain a +consistent, duplicate-free candidate list deeper than `ResultNum`; the tuner +reranks the complete candidate list with each native RaBitQ model, compares its +first `ResultNum` IDs with the first `ResultNum` exact IDs, and averages query +recalls equally. Bits 1 through 8 are tried in ascending order and the first +width meeting `TargetRecall` is selected; the build fails if none qualifies. + +The centroid is computed over every base vector in bounded native reader +batches. The selected official model and complete DEFAULT-format encoded base +file are built with the complete index in a unique sibling staging directory +and published under `[Base] IndexDirectory` only after construction succeeds. +The published files include `rabitq_auto_quantizer.bin` and +`rabitq_auto_vectors.bin`; a previous index is restored if directory +publication fails. `--outputfolder` must name the same directory as +`[Base] IndexDirectory`, and `--input`, `--quantizer`, and +`Section.Parameter=value` overrides are rejected in auto-tuning mode. The +builder internally uses the generated `UInt8` width (obtained from +`GetNumSubvectors()`), model, and vector path while queries remain raw Float +vectors through ADC. The original `VectorPath` remains the sole raw source; +no JSON, environment overrides, manual INI edits, or Python/Faiss step is +used. If the section is absent or disabled, existing build behavior is +unchanged. + +For a manually selected width, train the official model and encode base vectors +with `Release/quantizer`, then use the generated model through +`QuantizerFilePath` in the normal SPANN workflow. For 128-dimensional SIFT +vectors, the encoded `UInt8` vectors use `Dim=68` at 3 bits (48 compact code +bytes plus five Float factors). + +Train and encode SIFT1M: ```bash Release/quantizer \ From 1a0b67767aee4b88c12a68b602cb169cfc641ace Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Thu, 3 Sep 2026 09:18:47 +0000 Subject: [PATCH 15/17] Align quantized SPANN build pipeline Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: df92132a-4da9-4314-b72d-27a1043dc036 --- AnnService/inc/Core/Common/IQuantizer.h | 2 + AnnService/inc/Core/Common/OPQQuantizer.h | 4 + AnnService/inc/Core/Common/RaBitQQuantizer.h | 1 + AnnService/inc/Core/Common/TruthSet.h | 9 +- AnnService/inc/Core/Common/WorkSpace.h | 10 +- .../inc/Core/SPANN/ExtraStaticSearcher.h | 239 +++++++++++--- AnnService/inc/Core/SPANN/Index.h | 5 + AnnService/inc/Core/SPANN/Options.h | 1 + .../inc/Core/SPANN/ParameterDefinitionList.h | 1 + AnnService/inc/Helper/VectorSetReader.h | 7 +- AnnService/src/Core/SPANN/SPANNIndex.cpp | 123 +++++-- AnnService/src/Helper/AsyncFileReader.cpp | 20 +- AnnService/src/Helper/VectorSetReader.cpp | 20 +- AnnService/src/IndexBuilder/main.cpp | 311 ++++-------------- AnnService/src/SSDServing/main.cpp | 7 +- .../build_SPANN_sift1m_rabitq3_global.ini | 1 + .../build_SPANN_sift1m_raw_static_control.ini | 1 + .../search_SPANN_sift1m_raw_static_n12.ini | 1 + .../search_SPANN_sift1m_raw_static_n20.ini | 1 + .../search_SPANN_sift1m_raw_static_n32.ini | 1 + Test/src/RaBitQQuantizerTest.cpp | 33 +- Tools/OPQ/OPQ_gpu_train_infer.py | 15 +- Tools/OPQ/README.md | 4 +- docs/GettingStart.md | 54 +-- 24 files changed, 452 insertions(+), 419 deletions(-) diff --git a/AnnService/inc/Core/Common/IQuantizer.h b/AnnService/inc/Core/Common/IQuantizer.h index 1dde003c4..7c203b80f 100644 --- a/AnnService/inc/Core/Common/IQuantizer.h +++ b/AnnService/inc/Core/Common/IQuantizer.h @@ -60,6 +60,8 @@ namespace SPTAG virtual float* GetL2DistanceTables() = 0; + virtual bool QuantizeForIndexBuild() const { return true; } + template T* GetCodebooks(); }; diff --git a/AnnService/inc/Core/Common/OPQQuantizer.h b/AnnService/inc/Core/Common/OPQQuantizer.h index 9a1f73c55..7cd6ad002 100644 --- a/AnnService/inc/Core/Common/OPQQuantizer.h +++ b/AnnService/inc/Core/Common/OPQQuantizer.h @@ -53,6 +53,10 @@ namespace SPTAG return GetEnumValueType(); } + bool QuantizeForIndexBuild() const override + { + return false; + } protected: using PQQuantizer::m_NumSubvectors; diff --git a/AnnService/inc/Core/Common/RaBitQQuantizer.h b/AnnService/inc/Core/Common/RaBitQQuantizer.h index 1858712f4..612f9f626 100644 --- a/AnnService/inc/Core/Common/RaBitQQuantizer.h +++ b/AnnService/inc/Core/Common/RaBitQQuantizer.h @@ -48,6 +48,7 @@ class RaBitQQuantizer : public IQuantizer DimensionType GetNumSubvectors() const override; int GetBase() const override; float* GetL2DistanceTables() override; + bool QuantizeForIndexBuild() const override { return false; } DimensionType Dimension() const { return m_dimension; } int Bits() const { return m_bits; } diff --git a/AnnService/inc/Core/Common/TruthSet.h b/AnnService/inc/Core/Common/TruthSet.h index 859e541f4..27f2ab73e 100644 --- a/AnnService/inc/Core/Common/TruthSet.h +++ b/AnnService/inc/Core/Common/TruthSet.h @@ -169,6 +169,11 @@ namespace SPTAG float meanrecall = 0, minrecall = MaxDist, maxrecall = 0, stdrecall = 0, meanmrr = 0; std::vector thisrecall(NumQuerys, 0); std::unique_ptr visited(new bool[K]); + const bool compareDistanceTies = + querySet != nullptr && + vectorSet != nullptr && + querySet->GetValueType() == GetEnumValueType() && + vectorSet->GetValueType() == GetEnumValueType(); for (SizeType i = 0; i < NumQuerys; i++) { int minpos = K; @@ -186,7 +191,7 @@ namespace SPTAG if (j < minpos) minpos = j; break; } - else if (vectorSet != nullptr) { + else if (compareDistanceTies) { float dist = COMMON::DistanceUtils::ComputeDistance((const T*)querySet->GetVector(i), (const T*)vectorSet->GetVector(results[i].GetResult(j)->VID), vectorSet->Dimension(), index->GetDistCalcMethod()); float truthDist = COMMON::DistanceUtils::ComputeDistance((const T*)querySet->GetVector(i), (const T*)vectorSet->GetVector(id), vectorSet->Dimension(), index->GetDistCalcMethod()); if (index->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine && fabs(dist - truthDist) < Epsilon) { @@ -213,7 +218,7 @@ namespace SPTAG std::vector truthvec; for (SizeType id : truth[i]) { float truthDist = 0.0; - if (vectorSet != nullptr) { + if (compareDistanceTies) { truthDist = COMMON::DistanceUtils::ComputeDistance((const T*)querySet->GetVector(i), (const T*)vectorSet->GetVector(id), querySet->Dimension(), index->GetDistCalcMethod()); } truthvec.emplace_back(id, truthDist); diff --git a/AnnService/inc/Core/Common/WorkSpace.h b/AnnService/inc/Core/Common/WorkSpace.h index 6e3404144..59eeb0426 100644 --- a/AnnService/inc/Core/Common/WorkSpace.h +++ b/AnnService/inc/Core/Common/WorkSpace.h @@ -146,13 +146,21 @@ namespace SPTAG inline void DoubleSize() { + const std::uint64_t oldPoolSize = m_poolSize; + const bool hadSecondHash = m_secondHash; std::uint64_t new_poolSize = ((m_poolSize + 1) << 1) - 1; SizeType* new_hashTable = new SizeType[(new_poolSize + 1) * 2]; memset(new_hashTable, 0, sizeof(SizeType) * (new_poolSize + 1) * 2); m_secondHash = false; - for (std::uint64_t i = 0; i <= new_poolSize; i++) + for (std::uint64_t i = 0; i <= oldPoolSize; i++) if (m_hashTable[i]) _CheckAndSet(new_hashTable, new_poolSize, true, m_hashTable[i]); + if (hadSecondHash) + { + SizeType* secondHashTable = m_hashTable.get() + oldPoolSize + 1; + for (std::uint64_t i = 0; i <= oldPoolSize; i++) + if (secondHashTable[i]) _CheckAndSet(new_hashTable, new_poolSize, true, secondHashTable[i]); + } m_exp++; m_poolSize = new_poolSize; diff --git a/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h b/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h index 03d323e58..112b539da 100644 --- a/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h +++ b/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h @@ -10,6 +10,7 @@ #include "inc/Core/Common/TruthSet.h" #include "Compressor.h" +#include #include #include #include @@ -124,8 +125,8 @@ namespace SPTAG SizeType vectorID = *(reinterpret_cast(p_postingListFullData + offsetVectorID));\ if (p_exWorkSpace->Deduper().CheckAndSet(vectorID)) { listElements--; continue; } \ (this->*m_parseEncoding)(listInfo, (ValueType*)(p_postingListFullData + offsetVector));\ - auto distance2leaf = m_headIndex->ComputeDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); \ - queryResults.AddPoint(vectorID, distance2leaf, queryResults.WithVec()? ByteArray((std::uint8_t*)(p_postingListFullData + offsetVector), sizeof(ValueType) * m_opt->m_dim, false) : ByteArray::c_empty); \ + auto distance2leaf = PostingDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); \ + queryResults.AddPoint(vectorID, distance2leaf, queryResults.WithVec()? ByteArray((std::uint8_t*)(p_postingListFullData + offsetVector), StoredVectorBytes(), false) : ByteArray::c_empty); \ } \ #define ProcessPostingOffset() \ @@ -137,7 +138,7 @@ namespace SPTAG if (p_exWorkSpace->Deduper().CheckAndSet(vectorID)) continue; \ if (p_exWorkSpace->m_filterFunc != nullptr && !p_exWorkSpace->m_filterFunc(m_headIndex->GetMetadata(vectorID))) continue; \ (this->*m_parseEncoding)(listInfo, (ValueType*)(p_postingListFullData + offsetVector));\ - auto distance2leaf = m_headIndex->ComputeDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); \ + auto distance2leaf = PostingDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); \ queryResults.AddPoint(vectorID, distance2leaf); \ foundResult = true;\ break;\ @@ -175,10 +176,52 @@ namespace SPTAG return m_available; } + inline size_t StoredVectorBytes() const + { + return m_headIndex->m_pQuantizer + ? static_cast(m_headIndex->m_pQuantizer->GetNumSubvectors()) + : static_cast(m_opt->m_dim) * sizeof(ValueType); + } + + inline float PostingDistance(const void* p_query, const char* p_vector) const + { + if (!m_headIndex->m_pQuantizer) + { + return m_headIndex->ComputeDistance(p_query, p_vector); + } + if (m_headIndex->GetDistCalcMethod() == DistCalcMethod::L2) + { + return m_headIndex->m_pQuantizer->L2Distance( + reinterpret_cast(p_query), + reinterpret_cast(p_vector)); + } + return m_headIndex->m_pQuantizer->CosineDistance( + reinterpret_cast(p_query), + reinterpret_cast(p_vector)); + } + + inline size_t DiskRequestIndex(const ExtraWorkSpace* p_exWorkSpace, size_t p_postingIndex) const + { + return p_postingIndex; + } + + inline Helper::AsyncReadRequest& DiskRequest(ExtraWorkSpace* p_exWorkSpace, size_t p_postingIndex) const + { + auto& request = p_exWorkSpace->m_diskRequests[DiskRequestIndex(p_exWorkSpace, p_postingIndex)]; + request.m_buffer = reinterpret_cast(p_exWorkSpace->m_pageBuffers[p_postingIndex].GetBuffer()); + return request; + } + virtual bool LoadIndex(Options& p_opt) override { + m_opt = &p_opt; + m_enableDeltaEncoding = p_opt.m_enableDeltaEncoding; + m_enablePostingListRearrange = p_opt.m_enablePostingListRearrange; + m_enableDataCompression = p_opt.m_enableDataCompression; + m_enableDictTraining = p_opt.m_enableDictTraining; + m_extraFullGraphFile = p_opt.m_indexDirectory + FolderSep + p_opt.m_ssdIndex; std::string curFile = m_extraFullGraphFile + "_" + std::to_string(m_layer); - p_opt.m_searchPostingPageLimit = max(p_opt.m_searchPostingPageLimit, static_cast((p_opt.m_postingVectorLimit * (p_opt.m_dim * sizeof(ValueType) + sizeof(SizeType)) + PageSize - 1) / PageSize)); + p_opt.m_searchPostingPageLimit = max(p_opt.m_searchPostingPageLimit, static_cast((p_opt.m_postingVectorLimit * (StoredVectorBytes() + sizeof(SizeType)) + PageSize - 1) / PageSize)); SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Load index with posting page limit:%d\n", p_opt.m_searchPostingPageLimit); do { auto curIndexFile = f_createAsyncIO(); @@ -211,12 +254,6 @@ namespace SPTAG } while (fileexists(curFile.c_str())); m_oneContext = (m_indexFiles.size() == 1); - m_opt = &p_opt; - m_enableDeltaEncoding = p_opt.m_enableDeltaEncoding; - m_enablePostingListRearrange = p_opt.m_enablePostingListRearrange; - m_enableDataCompression = p_opt.m_enableDataCompression; - m_enableDictTraining = p_opt.m_enableDictTraining; - if (m_enablePostingListRearrange) m_parsePosting = &ExtraStaticSearcher::ParsePostingListRearrange; else m_parsePosting = &ExtraStaticSearcher::ParsePostingList; if (m_enableDeltaEncoding) m_parseEncoding = &ExtraStaticSearcher::ParseDeltaEncoding; @@ -241,12 +278,22 @@ namespace SPTAG bool) { const uint32_t postingListCount = static_cast(p_exWorkSpace->m_postingIDs.size()); + if (postingListCount > p_exWorkSpace->m_pageBuffers.size() || + (postingListCount > 0 && + DiskRequestIndex(p_exWorkSpace, postingListCount - 1) >= p_exWorkSpace->m_diskRequests.size())) { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Static search workspace is too small: postings=%u buffers=%zu requests=%zu.\n", + postingListCount, p_exWorkSpace->m_pageBuffers.size(), + p_exWorkSpace->m_diskRequests.size()); + return ErrorCode::Fail; + } COMMON::QueryResultSet& queryResults = *((COMMON::QueryResultSet*)&p_queryResults); - int diskRead = 0; int diskIO = 0; int listElements = 0; + int missingPostingIDs = 0; + ErrorCode scanRet = ErrorCode::Success; #if defined(ASYNC_READ) && !defined(BATCH_READ) int unprocessed = 0; @@ -257,7 +304,8 @@ namespace SPTAG auto curPostingID = p_exWorkSpace->m_postingIDs[pi]; auto it = m_globalVectorIDToHeadMap.find(curPostingID); if (it == m_globalVectorIDToHeadMap.end()) { - auto& request = p_exWorkSpace->m_diskRequests[pi]; + ++missingPostingIDs; + auto& request = DiskRequest(p_exWorkSpace, pi); request.m_readSize = 0; request.m_success = false; request.m_callback = nullptr; @@ -276,9 +324,16 @@ namespace SPTAG listElements += listInfo->listEleCount; size_t totalBytes = (static_cast(listInfo->listPageCount) << PageSizeEx); + if (totalBytes > p_exWorkSpace->m_pageBuffers[pi].GetPageSize()) { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Static posting %d requires %zu bytes but its workspace buffer has %zu bytes.\n", + curPostingID, totalBytes, + p_exWorkSpace->m_pageBuffers[pi].GetPageSize()); + return ErrorCode::DiskIOFail; + } #ifdef ASYNC_READ - auto& request = p_exWorkSpace->m_diskRequests[pi]; + auto& request = DiskRequest(p_exWorkSpace, pi); request.m_offset = listInfo->listOffset; request.m_readSize = totalBytes; request.m_status = (fileid << 16) | (request.m_status & 0xffff); @@ -286,10 +341,15 @@ namespace SPTAG request.m_success = false; #ifdef BATCH_READ // async batch read - request.m_callback = [&p_exWorkSpace, &queryResults, &request, &listElements, this](bool success) + Helper::AsyncReadRequest* requestPtr = &request; + request.m_callback = [p_exWorkSpace, &queryResults, requestPtr, &listElements, &scanRet, this](bool success) { - char* buffer = request.m_buffer; - ListInfo* listInfo = (ListInfo*)(request.m_payload); + if (!success) { + scanRet = ErrorCode::DiskIOFail; + return; + } + char* buffer = requestPtr->m_buffer; + ListInfo* listInfo = (ListInfo*)(requestPtr->m_payload); // decompress posting list char* p_postingListFullData = buffer + listInfo->pageOffset; @@ -301,9 +361,10 @@ namespace SPTAG ProcessPosting(); }; #else // async read - request.m_callback = [&p_exWorkSpace, &request](bool success) + Helper::AsyncReadRequest* requestPtr = &request; + request.m_callback = [p_exWorkSpace, requestPtr](bool success) { - p_exWorkSpace->m_processIocp.push(&request); + p_exWorkSpace->m_processIocp.push(requestPtr); }; ++unprocessed; @@ -331,9 +392,19 @@ namespace SPTAG #endif } + if (missingPostingIDs > 0) { + SPTAGLIB_LOG(Helper::LogLevel::LL_Warning, + "Static search skipped %d of %u posting IDs that are absent from the head map.\n", + missingPostingIDs, postingListCount); + } + #ifdef ASYNC_READ #ifdef BATCH_READ - BatchReadFileAsync(m_indexFiles, (p_exWorkSpace->m_diskRequests).data(), postingListCount); + if (!BatchReadFileAsync(m_indexFiles, (p_exWorkSpace->m_diskRequests).data(), postingListCount)) + { + return ErrorCode::DiskIOFail; + } + if (scanRet != ErrorCode::Success) return scanRet; #else while (unprocessed > 0) { @@ -433,9 +504,9 @@ namespace SPTAG if (p_exWorkSpace->Deduper().CheckAndSet(vectorID)) continue; (this->*m_parseEncoding)(listInfo, (ValueType*)(p_postingListFullData + offsetVector)); - auto distance2leaf = m_headIndex->ComputeDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); + auto distance2leaf = PostingDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); p_results.emplace_back(vectorID, distance2leaf, ByteArray::c_empty, - queryResults.WithVec() ? ByteArray::Alloc((std::uint8_t*)(p_postingListFullData + offsetVector), sizeof(ValueType) * m_opt->m_dim) : ByteArray::c_empty); + queryResults.WithVec() ? ByteArray::Alloc((std::uint8_t*)(p_postingListFullData + offsetVector), StoredVectorBytes()) : ByteArray::c_empty); } return ErrorCode::Success; }; @@ -451,7 +522,7 @@ namespace SPTAG auto curPostingID = p_exWorkSpace->m_postingIDs[pi]; auto it = m_globalVectorIDToHeadMap.find(curPostingID); if (it == m_globalVectorIDToHeadMap.end()) { - auto& request = p_exWorkSpace->m_diskRequests[pi]; + auto& request = DiskRequest(p_exWorkSpace, pi); request.m_readSize = 0; request.m_success = false; request.m_callback = nullptr; @@ -468,7 +539,7 @@ namespace SPTAG size_t totalBytes = (static_cast(listInfo->listPageCount) << PageSizeEx); #ifdef ASYNC_READ - auto& request = p_exWorkSpace->m_diskRequests[pi]; + auto& request = DiskRequest(p_exWorkSpace, pi); request.m_offset = listInfo->listOffset; request.m_readSize = totalBytes; request.m_status = (fileid << 16) | (request.m_status & 0xffff); @@ -476,20 +547,22 @@ namespace SPTAG request.m_success = false; #ifdef BATCH_READ - request.m_callback = [&appendPosting, &request, &scanRet](bool success) + Helper::AsyncReadRequest* requestPtr = &request; + request.m_callback = [&appendPosting, requestPtr, &scanRet](bool success) { if (!success) { scanRet = ErrorCode::DiskIOFail; return; } - ErrorCode ret = appendPosting(request.m_buffer, static_cast(request.m_payload)); + ErrorCode ret = appendPosting(requestPtr->m_buffer, static_cast(requestPtr->m_payload)); if (ret != ErrorCode::Success) scanRet = ret; }; #else - request.m_callback = [&p_exWorkSpace, &request](bool success) + Helper::AsyncReadRequest* requestPtr = &request; + request.m_callback = [p_exWorkSpace, requestPtr](bool success) { - p_exWorkSpace->m_processIocp.push(&request); + p_exWorkSpace->m_processIocp.push(requestPtr); }; ++unprocessed; @@ -569,7 +642,7 @@ namespace SPTAG size_t totalBytes = (static_cast(listInfo->listPageCount) << PageSizeEx); #ifdef ASYNC_READ - auto& request = p_exWorkSpace->m_diskRequests[pi]; + auto& request = DiskRequest(p_exWorkSpace, pi); request.m_offset = listInfo->listOffset; request.m_readSize = totalBytes; request.m_status = (fileid << 16) | (request.m_status & 0xffff); @@ -594,9 +667,10 @@ namespace SPTAG */ }; #else // async read - request.m_callback = [&p_exWorkSpace, &request](bool success) + Helper::AsyncReadRequest* requestPtr = &request; + request.m_callback = [p_exWorkSpace, requestPtr](bool success) { - p_exWorkSpace->m_processIocp.push(&request); + p_exWorkSpace->m_processIocp.push(requestPtr); }; ++unprocessed; @@ -635,6 +709,7 @@ namespace SPTAG success = BatchReadFileAsync(m_indexFiles, (p_exWorkSpace->m_diskRequests).data(), postingListCount); retry++; } + if (!success) return ErrorCode::DiskIOFail; #else while (unprocessed > 0) { @@ -677,7 +752,8 @@ namespace SPTAG continue; } char* buffer = (char*)((p_exWorkSpace->m_pageBuffers[p_exWorkSpace->m_pi]).GetBuffer()); - ListInfo* listInfo = static_cast(p_exWorkSpace->m_diskRequests[p_exWorkSpace->m_pi].m_payload); + ListInfo* listInfo = static_cast( + DiskRequest(p_exWorkSpace, p_exWorkSpace->m_pi).m_payload); // decompress posting list char* p_postingListFullData = buffer + listInfo->pageOffset; if (m_enableDataCompression && p_exWorkSpace->m_offset == 0) @@ -721,6 +797,7 @@ namespace SPTAG size_t p_postingListSize, Selection &p_selections, std::shared_ptr p_fullVectors, + std::shared_ptr p_quantizedVectors, COMMON::Dataset& p_localToGlobal, bool p_enableDeltaEncoding = false, bool p_enablePostingListRearrange = false, @@ -746,7 +823,29 @@ namespace SPTAG vectorID.append(reinterpret_cast(&vid), sizeof(SizeType)); ValueType *p_vector = reinterpret_cast(p_fullVectors->GetVector(vid)); - if (p_enableDeltaEncoding) + const bool quantizePosting = + m_headIndex->m_pQuantizer && + p_fullVectors->GetValueType() == m_headIndex->m_pQuantizer->GetReconstructType() && + p_fullVectors->Dimension() == m_headIndex->m_pQuantizer->ReconstructDim(); + if (quantizePosting && p_quantizedVectors) + { + const void* quantizedVector = p_quantizedVectors->GetVector(vid); + vector.append( + reinterpret_cast(quantizedVector), + p_quantizedVectors->PerVectorDataSize()); + } + else if (quantizePosting) + { + thread_local std::vector quantizedVector; + quantizedVector.resize(static_cast( + m_headIndex->m_pQuantizer->GetNumSubvectors())); + m_headIndex->m_pQuantizer->QuantizeVector( + p_vector, quantizedVector.data(), false); + vector.append( + reinterpret_cast(quantizedVector.data()), + quantizedVector.size()); + } + else if (p_enableDeltaEncoding) { DimensionType n = p_fullVectors->Dimension(); std::vector p_vector_delta(n); @@ -808,10 +907,60 @@ namespace SPTAG SizeType fullCount = 0; size_t vectorInfoSize = 0; + std::shared_ptr quantizedVectors; { auto fullVectors = p_reader->GetVectorSet(); fullCount = fullVectors->Count(); - vectorInfoSize = fullVectors->PerVectorDataSize() + sizeof(SizeType); + const bool quantizePosting = + m_headIndex->m_pQuantizer && + fullVectors->GetValueType() == m_headIndex->m_pQuantizer->GetReconstructType() && + fullVectors->Dimension() == m_headIndex->m_pQuantizer->ReconstructDim(); + if (quantizePosting && p_opt.m_enableDeltaEncoding) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Delta encoding is not supported when SSD postings are quantized.\n"); + return false; + } + vectorInfoSize = (quantizePosting + ? static_cast(m_headIndex->m_pQuantizer->GetNumSubvectors()) + : fullVectors->PerVectorDataSize()) + sizeof(SizeType); + if (quantizePosting && !p_opt.m_quantizedVectorPath.empty()) + { + auto quantizedOptions = std::make_shared( + VectorValueType::UInt8, + m_headIndex->m_pQuantizer->GetNumSubvectors(), + VectorFileType::DEFAULT, + "|", + p_opt.m_iSSDNumberOfThreads, + false); + auto quantizedReader = Helper::VectorSetReader::CreateInstance(quantizedOptions); + if (!quantizedReader || + quantizedReader->LoadFile(p_opt.m_quantizedVectorPath) != ErrorCode::Success) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Failed to load pre-quantized posting vectors: %s\n", + p_opt.m_quantizedVectorPath.c_str()); + return false; + } + quantizedVectors = quantizedReader->GetVectorSet(); + if (!quantizedVectors || + quantizedVectors->GetValueType() != VectorValueType::UInt8 || + quantizedVectors->Dimension() != + m_headIndex->m_pQuantizer->GetNumSubvectors() || + quantizedVectors->Count() < fullCount || + quantizedVectors->PerVectorDataSize() != + m_headIndex->m_pQuantizer->GetNumSubvectors()) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Pre-quantized posting vectors are incompatible with the quantizer.\n"); + return false; + } + SPTAGLIB_LOG(Helper::LogLevel::LL_Info, + "Loaded pre-quantized posting vectors from %s (%d,%d).\n", + p_opt.m_quantizedVectorPath.c_str(), + quantizedVectors->Count(), + quantizedVectors->Dimension()); + } } if (upperBound > 0) fullCount = upperBound; @@ -1061,7 +1210,7 @@ namespace SPTAG headVector = (ValueType*)p_headIndex->GetSample(j); } std::string postingListFullData = GetPostingListFullData( - j, curPostingListSizes[j], selections, fullVectors, p_localToGlobal, p_opt.m_enableDeltaEncoding, p_opt.m_enablePostingListRearrange, headVector); + j, curPostingListSizes[j], selections, fullVectors, quantizedVectors, p_localToGlobal, p_opt.m_enableDeltaEncoding, p_opt.m_enablePostingListRearrange, headVector); samplesBuffer += postingListFullData; samplesSizes.push_back(postingListFullData.size()); @@ -1101,7 +1250,7 @@ namespace SPTAG } std::string postingListFullData = GetPostingListFullData(postingListId, postingListSize[postingListId], - selections, fullVectors, p_localToGlobal, p_opt.m_enableDeltaEncoding, + selections, fullVectors, quantizedVectors, p_localToGlobal, p_opt.m_enableDeltaEncoding, p_opt.m_enablePostingListRearrange, headVector); size_t sizeToCompress = postingListSize[postingListId] * vectorInfoSize; if (sizeToCompress != postingListFullData.size()) @@ -1166,7 +1315,7 @@ namespace SPTAG postPageNum, postPageOffset, postingOrderInIndex, - fullVectors, p_headToLocal, p_localToGlobal, + fullVectors, quantizedVectors, p_headToLocal, p_localToGlobal, curPostingListOffSet); } @@ -1476,7 +1625,10 @@ namespace SPTAG const std::unique_ptr& p_postPageNum, const std::unique_ptr& p_postPageOffset, const std::vector& p_postingOrderInIndex, - std::shared_ptr p_fullVectors, COMMON::Dataset& p_headToLocal, COMMON::Dataset& p_localToGlobal, + std::shared_ptr p_fullVectors, + std::shared_ptr p_quantizedVectors, + COMMON::Dataset& p_headToLocal, + COMMON::Dataset& p_localToGlobal, size_t p_postingListOffset) { SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Start output...\n"); @@ -1541,7 +1693,15 @@ namespace SPTAG } // Vector dimension - int i32Val = static_cast(p_fullVectors->Dimension()); + size_t storedVectorBytes = p_spacePerVector - sizeof(SizeType); + if (storedVectorBytes % sizeof(ValueType) != 0) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Posting vector bytes %zu are not aligned to value type size %zu.\n", + storedVectorBytes, sizeof(ValueType)); + throw std::runtime_error("Posting vector size is not value-type aligned"); + } + int i32Val = static_cast(storedVectorBytes / sizeof(ValueType)); if (ptr->WriteBinary(sizeof(i32Val), reinterpret_cast(&i32Val)) != sizeof(i32Val)) { SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Failed to write SSDIndex File!"); throw std::runtime_error("Failed to write SSDIndex File"); @@ -1687,7 +1847,7 @@ namespace SPTAG headVector = (ValueType *)p_headIndex->GetSample(postingListId); } std::string postingListFullData = GetPostingListFullData( - postingListId, p_postingListSizes[id], p_postingSelections, p_fullVectors, p_localToGlobal, p_enableDeltaEncoding, p_enablePostingListRearrange, headVector); + postingListId, p_postingListSizes[id], p_postingSelections, p_fullVectors, p_quantizedVectors, p_localToGlobal, p_enableDeltaEncoding, p_enablePostingListRearrange, headVector); size_t postingListFullSize = p_postingListSizes[id] * p_spacePerVector; if (postingListFullSize != postingListFullData.size()) { @@ -1796,7 +1956,6 @@ namespace SPTAG int m_vectorInfoSize = 0; int m_iDataDimension = 0; - int m_totalListCount = 0; int m_listPerFile = 0; diff --git a/AnnService/inc/Core/SPANN/Index.h b/AnnService/inc/Core/SPANN/Index.h index 041acbf1b..dd2d9c90b 100644 --- a/AnnService/inc/Core/SPANN/Index.h +++ b/AnnService/inc/Core/SPANN/Index.h @@ -85,6 +85,11 @@ namespace SPTAG std::shared_ptr> m_freeWorkSpaceIds; std::atomic m_workspaceCount = 0; + bool UseQuantizerForIndexBuild() const + { + return m_pQuantizer && m_pQuantizer->QuantizeForIndexBuild(); + } + public: Index() { diff --git a/AnnService/inc/Core/SPANN/Options.h b/AnnService/inc/Core/SPANN/Options.h index 3de2d8fa2..a977278b7 100644 --- a/AnnService/inc/Core/SPANN/Options.h +++ b/AnnService/inc/Core/SPANN/Options.h @@ -44,6 +44,7 @@ namespace SPTAG { bool m_deleteHeadVectors; int m_ssdIndexFileNum; std::string m_quantizerFilePath; + std::string m_quantizedVectorPath; SizeType m_datasetRowsInBlock; SizeType m_datasetCapacity; diff --git a/AnnService/inc/Core/SPANN/ParameterDefinitionList.h b/AnnService/inc/Core/SPANN/ParameterDefinitionList.h index bd6a1f943..09fc5b234 100644 --- a/AnnService/inc/Core/SPANN/ParameterDefinitionList.h +++ b/AnnService/inc/Core/SPANN/ParameterDefinitionList.h @@ -33,6 +33,7 @@ DefineBasicParameter(m_ssdIndex, std::string, std::string("SPTAGFullList.bin"), DefineBasicParameter(m_deleteHeadVectors, bool, false, "DeleteHeadVectors") DefineBasicParameter(m_ssdIndexFileNum, int, 1, "SSDIndexFileNum") DefineBasicParameter(m_quantizerFilePath, std::string, std::string(), "QuantizerFilePath") +DefineBasicParameter(m_quantizedVectorPath, std::string, std::string(), "QuantizedVectorPath") DefineBasicParameter(m_datasetRowsInBlock, SizeType, 1024 * 1024, "DataBlockSize") DefineBasicParameter(m_datasetCapacity, SizeType, SPTAG::MaxSize, "DataCapacity") #endif diff --git a/AnnService/inc/Helper/VectorSetReader.h b/AnnService/inc/Helper/VectorSetReader.h index 1ca1ba00d..ea68419dc 100644 --- a/AnnService/inc/Helper/VectorSetReader.h +++ b/AnnService/inc/Helper/VectorSetReader.h @@ -19,12 +19,7 @@ namespace Helper class ReaderOptions : public ArgumentsParser { public: - ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, VectorFileType p_fileType, - std::string p_vectorDelimiter = "|", std::uint32_t p_threadNum = 32, - bool p_normalized = false); - ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, VectorFileType p_fileType, - std::string p_vectorDelimiter, std::uint32_t p_threadNum, - bool p_normalized, bool p_requireInputFormat); + ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, VectorFileType p_fileType, std::string p_vectorDelimiter = "|", std::uint32_t p_threadNum = 32, bool p_normalized = false); ~ReaderOptions(); diff --git a/AnnService/src/Core/SPANN/SPANNIndex.cpp b/AnnService/src/Core/SPANN/SPANNIndex.cpp index 2b367dbf1..15a1c42fc 100644 --- a/AnnService/src/Core/SPANN/SPANNIndex.cpp +++ b/AnnService/src/Core/SPANN/SPANNIndex.cpp @@ -58,10 +58,20 @@ template void Index::SetQuantizer(std::shared_ptrSetEnableADC(m_options.m_enableADC); - m_fComputeDistance = m_pQuantizer->DistanceCalcSelector(m_options.m_distCalcMethod); - m_iBaseSquare = (m_options.m_distCalcMethod == DistCalcMethod::Cosine) - ? m_pQuantizer->GetBase() * m_pQuantizer->GetBase() - : 1; + if (UseQuantizerForIndexBuild()) + { + m_fComputeDistance = m_pQuantizer->DistanceCalcSelector(m_options.m_distCalcMethod); + m_iBaseSquare = (m_options.m_distCalcMethod == DistCalcMethod::Cosine) + ? m_pQuantizer->GetBase() * m_pQuantizer->GetBase() + : 1; + } + else + { + m_fComputeDistance = COMMON::DistanceCalcSelector(m_options.m_distCalcMethod); + m_iBaseSquare = (m_options.m_distCalcMethod == DistCalcMethod::Cosine) + ? COMMON::Utils::GetBase() * COMMON::Utils::GetBase() + : 1; + } } else { @@ -104,7 +114,10 @@ template ErrorCode Index::LoadConfig(Helper::IniReader &p_reader template ErrorCode Index::LoadIndexDataFromMemory(const std::vector &p_indexBlobs) { /** Need to modify **/ - m_topIndex->SetQuantizer(m_pQuantizer); + if (UseQuantizerForIndexBuild()) + { + m_topIndex->SetQuantizer(m_pQuantizer); + } if (!m_options.m_persistentBufferPath.empty() && !direxists(m_options.m_persistentBufferPath.c_str())) mkdir(m_options.m_persistentBufferPath.c_str()); @@ -157,7 +170,10 @@ template ErrorCode Index::LoadIndexDataFromMemory(const std::vec template ErrorCode Index::LoadIndexData(const std::vector> &p_indexStreams) { - m_topIndex->SetQuantizer(m_pQuantizer); + if (UseQuantizerForIndexBuild()) + { + m_topIndex->SetQuantizer(m_pQuantizer); + } if (!m_options.m_persistentBufferPath.empty() && !direxists(m_options.m_persistentBufferPath.c_str())) mkdir(m_options.m_persistentBufferPath.c_str()); @@ -340,8 +356,14 @@ template ErrorCode Index::SearchIndex(QueryResult &p_query, Sear if (p_query.GetResultNum() >= m_options.m_searchInternalResultNum) p_queryResults = (COMMON::QueryResultSet *)&p_query; else + { p_queryResults = new COMMON::QueryResultSet((const T *)p_query.GetTarget(), m_options.m_searchInternalResultNum, p_query.WithMeta(), p_query.WithVec()); + if (m_pQuantizer) + { + p_queryResults->SetTarget((const T *)p_query.GetTarget(), m_pQuantizer); + } + } ErrorCode ret; auto searchStart = std::chrono::high_resolution_clock::now(); @@ -668,8 +690,14 @@ ErrorCode Index::SearchHeadIndex(QueryResult& p_query, int p_tolayer, ExtraWo if (p_query.GetResultNum() >= m_options.m_searchInternalResultNum) p_queryResults = (COMMON::QueryResultSet *)&p_query; else + { p_queryResults = new COMMON::QueryResultSet((const T *)p_query.GetTarget(), m_options.m_searchInternalResultNum, p_query.WithMeta(), p_query.WithVec()); + if (m_pQuantizer) + { + p_queryResults->SetTarget((const T *)p_query.GetTarget(), m_pQuantizer); + } + } ErrorCode ret; if ((ret = m_topIndex->SearchIndex(*p_queryResults)) != ErrorCode::Success) @@ -723,7 +751,7 @@ ErrorCode Index::SearchDiskIndex(QueryResult &p_query, SearchStats *p_stats, COMMON::QueryResultSet localResults((const T *)p_query.GetTarget(), m_options.m_searchInternalResultNum, p_query.WithMeta(), p_query.WithVec()); std::vector headCandidates; headCandidates.reserve(m_options.m_searchInternalResultNum); - if (m_pQuantizer && p_query.HasQuantizedTarget()) + if (m_pQuantizer) { localResults.SetTarget((const T *)p_query.GetTarget(), m_pQuantizer); } @@ -1119,7 +1147,9 @@ bool Index::SelectHeadInternal(std::shared_ptr &p_re bkt->m_iTreeNumber = m_options.m_iTreeNumber; bkt->m_fBalanceFactor = m_options.m_fBalanceFactor; bkt->m_parallelBuild = m_options.m_parallelBKTBuild; - bkt->m_pQuantizer = m_pQuantizer; + bkt->m_pQuantizer = UseQuantizerForIndexBuild() + ? m_pQuantizer + : nullptr; SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Start invoking BuildTrees.\n"); SPTAGLIB_LOG( Helper::LogLevel::LL_Info, @@ -1230,6 +1260,33 @@ bool Index::SelectHeadInternal(std::shared_ptr &p_re template ErrorCode Index::BuildIndexInternalLayer(std::shared_ptr &p_reader) { + struct ScopedIndexBuildADCMode + { + std::shared_ptr m_quantizer; + bool m_restore; + bool m_enableADC; + + ScopedIndexBuildADCMode(std::shared_ptr quantizer, bool quantizedIndexBuild) + : m_quantizer(std::move(quantizer)), + m_restore(m_quantizer != nullptr && quantizedIndexBuild), + m_enableADC(false) + { + if (m_restore) + { + m_enableADC = m_quantizer->GetEnableADC(); + m_quantizer->SetEnableADC(false); + } + } + + ~ScopedIndexBuildADCMode() + { + if (m_restore) + { + m_quantizer->SetEnableADC(m_enableADC); + } + } + } scopedADCMode(m_pQuantizer, UseQuantizerForIndexBuild()); + int currentLayer = static_cast(m_extraSearchers.size()); COMMON::Dataset localToGlobalID; { @@ -1259,7 +1316,7 @@ template ErrorCode Index::BuildIndexInternalLayer(std::shared_pt if (m_options.m_selectHead && m_topIndex == nullptr) { bool success = false; - if (m_pQuantizer) + if (UseQuantizerForIndexBuild()) { success = SelectHeadInternal(p_reader); } @@ -1280,13 +1337,21 @@ template ErrorCode Index::BuildIndexInternalLayer(std::shared_pt SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Begin Build Head...\n"); if (m_options.m_buildHead && m_topIndex == nullptr) { - auto valueType = m_pQuantizer ? SPTAG::VectorValueType::UInt8 : m_options.m_valueType; - auto dims = m_pQuantizer ? m_pQuantizer->GetNumSubvectors() : m_options.m_dim; + const bool quantizedIndexBuild = UseQuantizerForIndexBuild(); + auto valueType = quantizedIndexBuild + ? SPTAG::VectorValueType::UInt8 + : m_options.m_valueType; + auto dims = quantizedIndexBuild + ? m_pQuantizer->GetNumSubvectors() + : m_options.m_dim; m_topIndex = SPTAG::VectorIndex::CreateInstance(m_options.m_indexAlgoType, valueType); m_topIndex->SetParameter("DistCalcMethod", SPTAG::Helper::Convert::ConvertToString(m_options.m_distCalcMethod)); m_topIndex->SetParameter("ParallelBKTBuild", m_options.m_parallelBKTBuild ? "true" : "false"); - m_topIndex->SetQuantizer(m_pQuantizer); + if (quantizedIndexBuild) + { + m_topIndex->SetQuantizer(m_pQuantizer); + } for (const auto &iter : m_topParameters) { m_topIndex->SetParameter(iter.first.c_str(), iter.second.c_str()); @@ -1308,9 +1373,11 @@ template ErrorCode Index::BuildIndexInternalLayer(std::shared_pt SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Failed to build head index.\n"); return ErrorCode::Fail; } - if (!m_options.m_quantizerFilePath.empty()) + if (!m_options.m_quantizerFilePath.empty() && quantizedIndexBuild) + { m_topIndex->SetQuantizerFileName( m_options.m_quantizerFilePath.substr(m_options.m_quantizerFilePath.find_last_of("/\\") + 1)); + } if (m_topIndex->SaveIndex(m_options.m_indexDirectory + FolderSep + m_options.m_headIndexFolder) != ErrorCode::Success) { @@ -1341,7 +1408,10 @@ template ErrorCode Index::BuildIndexInternalLayer(std::shared_pt (m_options.m_indexDirectory + FolderSep + m_options.m_headIndexFolder).c_str()); return ErrorCode::Fail; } - m_topIndex->SetQuantizer(m_pQuantizer); + if (UseQuantizerForIndexBuild()) + { + m_topIndex->SetQuantizer(m_pQuantizer); + } if (!CheckHeadIndexType()) return ErrorCode::Fail; @@ -1475,7 +1545,10 @@ template ErrorCode Index::BuildIndexInternal(std::shared_ptrSetQuantizer(m_pQuantizer); + if (UseQuantizerForIndexBuild()) + { + m_topIndex->SetQuantizer(m_pQuantizer); + } m_topIndex->SetParameter("NumberOfThreads", std::to_string(m_options.m_iSSDNumberOfThreads)); m_topIndex->SetParameter("MaxCheck", std::to_string(m_options.m_maxCheck)); m_topIndex->SetParameter("HashTableExponent", std::to_string(m_options.m_hashExp)); @@ -1662,8 +1735,13 @@ template ErrorCode Index::BuildIndexInternal(std::shared_ptr ErrorCode Index::BuildIndex(bool p_normalized) { - SPTAG::VectorValueType valueType = m_pQuantizer ? SPTAG::VectorValueType::UInt8 : m_options.m_valueType; - SizeType dim = m_pQuantizer ? m_pQuantizer->GetNumSubvectors() : m_options.m_dim; + const bool quantizedIndexBuild = UseQuantizerForIndexBuild(); + SPTAG::VectorValueType valueType = quantizedIndexBuild + ? SPTAG::VectorValueType::UInt8 + : GetEnumValueType(); + SizeType dim = quantizedIndexBuild + ? m_pQuantizer->GetNumSubvectors() + : m_options.m_dim; std::shared_ptr vectorOptions( new Helper::ReaderOptions(valueType, dim, m_options.m_vectorType, m_options.m_vectorDelimiter, m_options.m_iSSDNumberOfThreads, p_normalized)); @@ -1710,7 +1788,10 @@ ErrorCode Index::BuildIndex(const void *p_data, SizeType p_vectorNum, Dimensi { vectorSet->Normalize(m_options.m_iSSDNumberOfThreads); } - SPTAG::VectorValueType valueType = m_pQuantizer ? SPTAG::VectorValueType::UInt8 : m_options.m_valueType; + SPTAG::VectorValueType valueType = + UseQuantizerForIndexBuild() + ? SPTAG::VectorValueType::UInt8 + : GetEnumValueType(); std::shared_ptr vectorOptions( new Helper::ReaderOptions(valueType, p_dimension, VectorFileType::DEFAULT, m_options.m_vectorDelimiter, m_options.m_iSSDNumberOfThreads, true)); @@ -1812,7 +1893,7 @@ template ErrorCode Index::SetParameter(const char *p_param, cons } if (SPTAG::Helper::StrUtils::StrEqualIgnoreCase(p_param, "DistCalcMethod")) { - if (m_pQuantizer) + if (UseQuantizerForIndexBuild()) { m_fComputeDistance = m_pQuantizer->DistanceCalcSelector(m_options.m_distCalcMethod); m_iBaseSquare = (m_options.m_distCalcMethod == DistCalcMethod::Cosine) @@ -1827,6 +1908,10 @@ template ErrorCode Index::SetParameter(const char *p_param, cons : 1; } } + if (SPTAG::Helper::StrUtils::StrEqualIgnoreCase(p_param, "EnableADC") && m_pQuantizer) + { + m_pQuantizer->SetEnableADC(m_options.m_enableADC); + } return ErrorCode::Success; } diff --git a/AnnService/src/Helper/AsyncFileReader.cpp b/AnnService/src/Helper/AsyncFileReader.cpp index 1245ab01f..2c1cf5249 100644 --- a/AnnService/src/Helper/AsyncFileReader.cpp +++ b/AnnService/src/Helper/AsyncFileReader.cpp @@ -126,7 +126,15 @@ bool BatchReadFileAsync(std::vector> &handlers, AsyncReadRequest *req = reinterpret_cast((events[i].data)); if (nullptr != req) { - req->m_callback(true); + const bool success = events[i].res == static_cast(req->m_readSize); + if (!success) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Async batch read failed at offset %llu, expected %zu bytes, actual %lld.\n", + req->m_offset, req->m_readSize, static_cast(events[i].res)); + } + req->m_success = success; + req->m_callback(success); } } totalQueued = totalDone; @@ -153,7 +161,15 @@ bool BatchReadFileAsync(std::vector> &handlers, AsyncReadRequest *req = reinterpret_cast((events[i].data)); if (nullptr != req) { - req->m_callback(true); + const bool success = events[i].res == static_cast(req->m_readSize); + if (!success) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Async batch read failed at offset %llu, expected %zu bytes, actual %lld.\n", + req->m_offset, req->m_readSize, static_cast(events[i].res)); + } + req->m_success = success; + req->m_callback(success); } } return true; diff --git a/AnnService/src/Helper/VectorSetReader.cpp b/AnnService/src/Helper/VectorSetReader.cpp index ad2b983f7..6cb752365 100644 --- a/AnnService/src/Helper/VectorSetReader.cpp +++ b/AnnService/src/Helper/VectorSetReader.cpp @@ -12,29 +12,15 @@ using namespace SPTAG::Helper; ReaderOptions::ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, VectorFileType p_fileType, std::string p_vectorDelimiter, std::uint32_t p_threadNum, bool p_normalized) - : ReaderOptions(p_valueType, p_dimension, p_fileType, p_vectorDelimiter, - p_threadNum, p_normalized, true) -{ -} - -ReaderOptions::ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, VectorFileType p_fileType, - std::string p_vectorDelimiter, std::uint32_t p_threadNum, bool p_normalized, - bool p_requireInputFormat) : m_inputValueType(p_valueType), m_dimension(p_dimension), m_inputFileType(p_fileType), m_vectorDelimiter(p_vectorDelimiter), m_threadNum(p_threadNum), m_normalized(p_normalized) { AddOptionalOption(m_threadNum, "-t", "--thread", "Thread Number."); AddOptionalOption(m_vectorDelimiter, "-dl", "--delimiter", "Vector delimiter."); AddOptionalOption(m_normalized, "-norm", "--normalized", "Vector is normalized."); - if (p_requireInputFormat) { - AddRequiredOption(m_dimension, "-d", "--dimension", "Dimension of vector."); - AddRequiredOption(m_inputValueType, "-v", "--vectortype", "Input vector data type. Default is float."); - AddRequiredOption(m_inputFileType, "-f", "--filetype", "Input file type (DEFAULT, TXT, XVEC). Default is DEFAULT."); - } else { - AddOptionalOption(m_dimension, "-d", "--dimension", "Dimension of vector."); - AddOptionalOption(m_inputValueType, "-v", "--vectortype", "Input vector data type. Default is float."); - AddOptionalOption(m_inputFileType, "-f", "--filetype", "Input file type (DEFAULT, TXT, XVEC). Default is DEFAULT."); - } + AddRequiredOption(m_dimension, "-d", "--dimension", "Dimension of vector."); + AddRequiredOption(m_inputValueType, "-v", "--vectortype", "Input vector data type. Default is float."); + AddRequiredOption(m_inputFileType, "-f", "--filetype", "Input file type (DEFAULT, TXT, XVEC). Default is DEFAULT."); } ReaderOptions::~ReaderOptions() diff --git a/AnnService/src/IndexBuilder/main.cpp b/AnnService/src/IndexBuilder/main.cpp index 8de54fad9..1085949d8 100644 --- a/AnnService/src/IndexBuilder/main.cpp +++ b/AnnService/src/IndexBuilder/main.cpp @@ -7,10 +7,7 @@ #include "inc/Helper/VectorSetReader.h" #include -#include #include -#include -#include #include #ifdef RABITQ @@ -19,132 +16,11 @@ using namespace SPTAG; -#ifdef RABITQ -namespace -{ -bool CreateBuildDirectory(const std::filesystem::path& p_final, - std::filesystem::path& p_build, - std::string& p_error) -{ - namespace fs = std::filesystem; - std::error_code error; - fs::create_directories(p_final.parent_path(), error); - if (error) { - p_error = "cannot create index parent directory: " + error.message(); - return false; - } - const auto nonce = std::chrono::high_resolution_clock::now().time_since_epoch().count(); - for (int attempt = 0; attempt < 100; ++attempt) { - p_build = p_final.parent_path() / - (p_final.filename().string() + ".rabitq-building-" + - std::to_string(nonce) + "-" + std::to_string(attempt)); - if (fs::create_directory(p_build, error)) { - return true; - } - if (error) { - p_error = "cannot create unique RaBitQ build directory: " + error.message(); - return false; - } - } - p_error = "cannot allocate a unique RaBitQ build directory"; - return false; -} - -bool PublishBuildDirectory(const std::filesystem::path& p_build, - const std::filesystem::path& p_final, - std::string& p_error) -{ - namespace fs = std::filesystem; - std::error_code error; - fs::path backup = p_build; - backup += ".previous"; - const bool hadPrevious = fs::exists(p_final, error); - if (error) { - p_error = "cannot inspect existing index directory: " + error.message(); - return false; - } - if (hadPrevious) { - fs::rename(p_final, backup, error); - if (error) { - p_error = "cannot preserve existing index directory: " + error.message(); - return false; - } - } - fs::rename(p_build, p_final, error); - if (error) { - const std::string publishError = error.message(); - if (hadPrevious) { - std::error_code restoreError; - fs::rename(backup, p_final, restoreError); - if (restoreError) { - p_error = "cannot publish new index (" + publishError + - ") or restore previous index (" + restoreError.message() + ")"; - return false; - } - } - std::error_code cleanupError; - fs::remove_all(p_build, cleanupError); - p_error = "cannot publish new index directory: " + publishError; - if (cleanupError) { - p_error += "; cannot remove failed staging directory: " + - cleanupError.message(); - } - return false; - } - if (hadPrevious) { - fs::remove_all(backup, error); - if (error) { - SPTAGLIB_LOG(Helper::LogLevel::LL_Warning, - "Published index but could not remove previous index backup %s: %s\n", - backup.string().c_str(), error.message().c_str()); - } - } - return true; -} - -bool FinalizeStagedConfig(const std::filesystem::path& p_build, - const std::filesystem::path& p_final, - std::string& p_error) -{ - const std::filesystem::path configPath = p_build / "indexloader.ini"; - std::ifstream input(configPath, std::ios::binary); - if (!input) { - p_error = "cannot open staged indexloader.ini"; - return false; - } - std::string config( - (std::istreambuf_iterator(input)), std::istreambuf_iterator()); - const std::string stagedRoot = p_build.string(); - const std::string finalRoot = p_final.string(); - std::size_t offset = 0; - int replacements = 0; - while ((offset = config.find(stagedRoot, offset)) != std::string::npos) { - config.replace(offset, stagedRoot.size(), finalRoot); - offset += finalRoot.size(); - ++replacements; - } - if (replacements == 0) { - p_error = "staged indexloader.ini does not reference its build directory"; - return false; - } - input.close(); - std::ofstream output(configPath, std::ios::binary | std::ios::trunc); - output.write(config.data(), static_cast(config.size())); - if (!output) { - p_error = "cannot finalize staged indexloader.ini"; - return false; - } - return true; -} -} // namespace -#endif - class BuilderOptions : public Helper::ReaderOptions { public: BuilderOptions() - : Helper::ReaderOptions( - VectorValueType::Float, 0, VectorFileType::TXT, "|", 32, false, false) + : Helper::ReaderOptions(VectorValueType::Float, 0, VectorFileType::TXT, "|", 32) { AddRequiredOption(m_outputFolder, "-o", "--outputfolder", "Output folder."); AddRequiredOption(m_indexAlgoType, "-a", "--algo", "Index Algorithm type."); @@ -178,6 +54,7 @@ int main(int argc, char *argv[]) { exit(1); } + Helper::IniReader iniReader; if (!options->m_builderConfigFile.empty() && iniReader.LoadIniFile(options->m_builderConfigFile) != ErrorCode::Success) @@ -186,8 +63,6 @@ int main(int argc, char *argv[]) return -1; } - const bool configuredAutoTune = iniReader.DoesSectionExist("RaBitQAutoTune") && - iniReader.GetParameter("RaBitQAutoTune", "isExecute", false); for (int i = 1; i < argc; i++) { std::string param(argv[i]); @@ -199,19 +74,6 @@ int main(int argc, char *argv[]) std::string paramVal = param.substr(idx + 1); std::string sectionName; idx = paramName.find("."); - const bool tunerOverride = - idx != std::string::npos && - Helper::StrUtils::StrEqualIgnoreCase( - paramName.substr(0, idx).c_str(), "RaBitQAutoTune"); - if (idx != std::string::npos && (configuredAutoTune || tunerOverride)) - { - SPTAGLIB_LOG( - Helper::LogLevel::LL_Error, - "RaBitQ auto-tuning parameters must come only from the INI; " - "command-line override %s is not allowed.\n", - param.c_str()); - return 1; - } if (idx != std::string::npos) { sectionName = paramName.substr(0, idx); @@ -236,11 +98,10 @@ int main(int argc, char *argv[]) return 1; } } - const bool autoTuneEnabled = configuredAutoTune; -#ifdef RABITQ - std::filesystem::path autoTuneBuildDirectory; - std::filesystem::path autoTuneFinalDirectory; -#endif + + const bool autoTuneEnabled = + iniReader.DoesSectionExist("RaBitQAutoTune") && + iniReader.GetParameter("RaBitQAutoTune", "isExecute", false); if (autoTuneEnabled) { #ifdef RABITQ @@ -250,42 +111,13 @@ int main(int argc, char *argv[]) "RaBitQ auto-tuning is supported only for SPANN index construction.\n"); return 1; } - if (!options->m_inputFiles.empty() || !options->m_quantizerFile.empty()) + if (!options->m_inputFiles.empty()) { - SPTAGLIB_LOG( - Helper::LogLevel::LL_Error, - "RaBitQ auto-tuning reads vectors and creates its quantizer from the INI; " - "--input and --quantizer are not allowed.\n"); - return 1; - } - const std::string configuredIndexDirectory = - iniReader.GetParameter("Base", "IndexDirectory", std::string()); - if (configuredIndexDirectory.empty()) - { - SPTAGLIB_LOG( - Helper::LogLevel::LL_Error, - "[Base] IndexDirectory is required for RaBitQ auto-tuning.\n"); - return 1; - } - autoTuneFinalDirectory = - std::filesystem::absolute(configuredIndexDirectory).lexically_normal(); - const auto commandOutput = - std::filesystem::absolute(options->m_outputFolder).lexically_normal(); - if (commandOutput != autoTuneFinalDirectory) - { - SPTAGLIB_LOG( - Helper::LogLevel::LL_Error, - "--outputfolder must match the authoritative [Base] IndexDirectory " - "when RaBitQ auto-tuning is enabled.\n"); - return 1; - } - std::string stagingError; - if (!CreateBuildDirectory( - autoTuneFinalDirectory, autoTuneBuildDirectory, stagingError)) - { - SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "%s\n", stagingError.c_str()); + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "RaBitQ auto-tuning reads vectors from the INI; --input is not allowed.\n"); return 1; } + COMMON::RaBitQAutoTuneResult tuneResult; std::string tuneError; ErrorCode tuneStatus = ErrorCode::Fail; @@ -293,7 +125,7 @@ int main(int argc, char *argv[]) { tuneStatus = COMMON::RaBitQAutoTuner::Run( - iniReader, autoTuneBuildDirectory.string(), tuneResult, tuneError); + iniReader, options->m_outputFolder, tuneResult, tuneError); } catch (const std::exception& exception) { @@ -301,24 +133,15 @@ int main(int argc, char *argv[]) } if (tuneStatus != ErrorCode::Success) { - std::error_code cleanupError; - std::filesystem::remove_all(autoTuneBuildDirectory, cleanupError); SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "RaBitQ auto-tuning failed: %s\n", tuneError.c_str()); return 1; } - builderValueType = VectorValueType::UInt8; quantizerFile = tuneResult.quantizerPath; - options->m_inputFiles.clear(); - iniReader.SetParameter("Base", "ValueType", "UInt8"); - iniReader.SetParameter("Base", "Dim", std::to_string(tuneResult.codeDimension)); - iniReader.SetParameter("Base", "VectorPath", tuneResult.vectorPath); - iniReader.SetParameter("Base", "VectorType", "DEFAULT"); iniReader.SetParameter("Base", "VectorSize", std::to_string(tuneResult.vectorCount)); iniReader.SetParameter("Base", "QuantizerFilePath", tuneResult.quantizerPath); - iniReader.SetParameter( - "Base", "IndexDirectory", autoTuneBuildDirectory.string()); + iniReader.SetParameter("Base", "QuantizedVectorPath", tuneResult.vectorPath); iniReader.SetParameter("BuildSSDIndex", "EnableADC", "true"); #else SPTAGLIB_LOG(Helper::LogLevel::LL_Error, @@ -328,6 +151,7 @@ int main(int argc, char *argv[]) } SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Set QuantizerFile = %s\n", quantizerFile.c_str()); + auto indexBuilder = VectorIndex::CreateInstance(options->m_indexAlgoType, builderValueType); if (!indexBuilder) { @@ -342,6 +166,44 @@ int main(int argc, char *argv[]) SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Cannot load quantizer file.\n"); return 1; } + if (!indexBuilder->m_pQuantizer->QuantizeForIndexBuild()) + { + const auto reconstructType = indexBuilder->m_pQuantizer->GetReconstructType(); + if (builderValueType != reconstructType) + { + if (!options->m_inputFiles.empty()) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "This quantizer requires raw reconstruct vectors for SPANN index build. " + "Set the input vector type to %s and keep pre-quantized codes in QuantizedVectorPath.\n", + Helper::Convert::ConvertToString(reconstructType).c_str()); + return 1; + } + if (iniReader.DoesParameterExist("Base", "VectorPath") && + !iniReader.DoesParameterExist("Base", "QuantizedVectorPath")) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "This quantizer requires [Base] VectorPath to point to raw reconstruct vectors " + "and [Base] QuantizedVectorPath to point to pre-quantized codes.\n"); + return 1; + } + + builderValueType = reconstructType; + iniReader.SetParameter( + "Base", "ValueType", + Helper::Convert::ConvertToString(builderValueType)); + indexBuilder = VectorIndex::CreateInstance(options->m_indexAlgoType, builderValueType); + if (!indexBuilder || + indexBuilder->LoadQuantizer(quantizerFile) != ErrorCode::Success || + !indexBuilder->m_pQuantizer) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Cannot recreate index builder for quantizer reconstruct type.\n"); + return 1; + } + } + } } std::string sections[] = {"Base", "SelectHead", "BuildHead", "BuildSSDIndex", "Index"}; @@ -356,15 +218,15 @@ int main(int argc, char *argv[]) indexBuilder->SetParameter(iter.first.c_str(), iter.second.c_str(), sections[i]); } } + ErrorCode code; std::shared_ptr vecset; if (options->m_inputFiles != "") { if (options->m_dimension <= 0) { - SPTAGLIB_LOG( - Helper::LogLevel::LL_Error, - "--dimension is required when indexbuilder reads --input directly.\n"); + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "--dimension is required when indexbuilder reads --input directly.\n"); return 1; } auto vectorReader = Helper::VectorSetReader::CreateInstance(options); @@ -388,78 +250,15 @@ int main(int argc, char *argv[]) } if (code == ErrorCode::Success) { - for (const auto& parameter : iniReader.GetParameters("SearchSSDIndex")) - { - std::string name = parameter.first; - if (Helper::StrUtils::StrEqualIgnoreCase(name.c_str(), "isExecute") || - Helper::StrUtils::StrEqualIgnoreCase(name.c_str(), "BuildSsdIndex")) - { - continue; - } - if (Helper::StrUtils::StrEqualIgnoreCase(name.c_str(), "PostingPageLimit")) - { - name = "SearchPostingPageLimit"; - } - else if (Helper::StrUtils::StrEqualIgnoreCase(name.c_str(), "InternalResultNum")) - { - name = "SearchInternalResultNum"; - } - indexBuilder->SetParameter( - name.c_str(), parameter.second.c_str(), "BuildSSDIndex"); - } - std::string saveFolder = options->m_outputFolder; -#ifdef RABITQ - if (autoTuneEnabled) { - saveFolder = autoTuneBuildDirectory.string(); - } -#endif - code = indexBuilder->SaveIndex(saveFolder); + code = indexBuilder->SaveIndex(options->m_outputFolder); if (code != ErrorCode::Success) { -#ifdef RABITQ - if (autoTuneEnabled) { - std::error_code cleanupError; - std::filesystem::remove_all(autoTuneBuildDirectory, cleanupError); - } -#endif SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Failed to save index.\n"); return 1; } -#ifdef RABITQ - if (autoTuneEnabled) - { - std::string finalizeError; - if (!FinalizeStagedConfig( - autoTuneBuildDirectory, autoTuneFinalDirectory, finalizeError)) - { - std::error_code cleanupError; - std::filesystem::remove_all(autoTuneBuildDirectory, cleanupError); - SPTAGLIB_LOG( - Helper::LogLevel::LL_Error, - "Failed to finalize auto-tuned index configuration: %s\n", - finalizeError.c_str()); - return 1; - } - std::string publishError; - if (!PublishBuildDirectory( - autoTuneBuildDirectory, autoTuneFinalDirectory, publishError)) - { - SPTAGLIB_LOG( - Helper::LogLevel::LL_Error, - "Failed to publish auto-tuned index: %s\n", publishError.c_str()); - return 1; - } - } -#endif } else { -#ifdef RABITQ - if (autoTuneEnabled) { - std::error_code cleanupError; - std::filesystem::remove_all(autoTuneBuildDirectory, cleanupError); - } -#endif SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Failed to build index.\n"); exit(1); } diff --git a/AnnService/src/SSDServing/main.cpp b/AnnService/src/SSDServing/main.cpp index 867e02fea..64cd51bf2 100644 --- a/AnnService/src/SSDServing/main.cpp +++ b/AnnService/src/SSDServing/main.cpp @@ -136,14 +136,15 @@ int BootProgram(bool forANNIndexTestTool, std::mapm_generateTruth) { SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Start generating truth. It's maybe a long time.\n"); + VectorValueType vectorValueType = opts->m_valueType; SizeType dim = opts->m_dim; if (index->m_pQuantizer) { - valueType = VectorValueType::UInt8; - dim = index->m_pQuantizer->GetNumSubvectors(); + vectorValueType = index->m_pQuantizer->GetReconstructType(); + dim = index->m_pQuantizer->ReconstructDim(); } std::shared_ptr vectorOptions( - new Helper::ReaderOptions(valueType, dim, opts->m_vectorType, opts->m_vectorDelimiter)); + new Helper::ReaderOptions(vectorValueType, dim, opts->m_vectorType, opts->m_vectorDelimiter)); auto vectorReader = Helper::VectorSetReader::CreateInstance(vectorOptions); if (ErrorCode::Success != vectorReader->LoadFile(opts->m_vectorPath)) { diff --git a/Script_AE/iniFile/build_SPANN_sift1m_rabitq3_global.ini b/Script_AE/iniFile/build_SPANN_sift1m_rabitq3_global.ini index d029c7521..c80706484 100644 --- a/Script_AE/iniFile/build_SPANN_sift1m_rabitq3_global.ini +++ b/Script_AE/iniFile/build_SPANN_sift1m_rabitq3_global.ini @@ -53,6 +53,7 @@ TmpDir=/datadisk/yfcc_fast/sptag_rabitq_parity_sift1m_rabitq3/tmp EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false +PostingQuantizer=None Rerank=0 EnableADC=true diff --git a/Script_AE/iniFile/build_SPANN_sift1m_raw_static_control.ini b/Script_AE/iniFile/build_SPANN_sift1m_raw_static_control.ini index 17753c161..b3e27a5a7 100644 --- a/Script_AE/iniFile/build_SPANN_sift1m_raw_static_control.ini +++ b/Script_AE/iniFile/build_SPANN_sift1m_raw_static_control.ini @@ -53,6 +53,7 @@ TmpDir=/datadisk/yfcc_fast/sptag_rabitq_parity_sift1m_raw_static_lambda001/tmp EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false +PostingQuantizer=None Rerank=0 [SearchSSDIndex] diff --git a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n12.ini b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n12.ini index 86e139ab3..9a6da2b5a 100644 --- a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n12.ini +++ b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n12.ini @@ -29,6 +29,7 @@ NumberOfThreads=24 EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false +PostingQuantizer=None Rerank=0 [SearchSSDIndex] diff --git a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n20.ini b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n20.ini index 4ab99f786..5d31e46b2 100644 --- a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n20.ini +++ b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n20.ini @@ -29,6 +29,7 @@ NumberOfThreads=24 EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false +PostingQuantizer=None Rerank=0 [SearchSSDIndex] diff --git a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n32.ini b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n32.ini index b67eb0695..533f51033 100644 --- a/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n32.ini +++ b/Script_AE/iniFile/search_SPANN_sift1m_raw_static_n32.ini @@ -29,6 +29,7 @@ NumberOfThreads=24 EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false +PostingQuantizer=None Rerank=0 [SearchSSDIndex] diff --git a/Test/src/RaBitQQuantizerTest.cpp b/Test/src/RaBitQQuantizerTest.cpp index 960834173..e689b8ee7 100644 --- a/Test/src/RaBitQQuantizerTest.cpp +++ b/Test/src/RaBitQQuantizerTest.cpp @@ -194,14 +194,14 @@ void VerifySpannSearch( std::filesystem::remove_all(index_directory); p_quantizer->SetEnableADC(false); - auto index = VectorIndex::CreateInstance(IndexAlgoType::SPANN, VectorValueType::UInt8); + auto index = VectorIndex::CreateInstance(IndexAlgoType::SPANN, VectorValueType::Float); BOOST_REQUIRE(index != nullptr); index->SetQuantizer(p_quantizer); ConfigureSpannIndex(index, index_directory, nullptr, p_storage, p_enable_compression); - BOOST_REQUIRE(index->BuildIndex(p_codes, nullptr, false, true) == ErrorCode::Success); + BOOST_REQUIRE(index->BuildIndex(p_raw, nullptr, false, true) == ErrorCode::Success); p_quantizer->SetEnableADC(true); - auto* spann_index = static_cast*>(index.get()); + auto* spann_index = static_cast*>(index.get()); std::vector head_ids; BOOST_REQUIRE(spann_index->GetHeadIndexMapping(1, head_ids) == ErrorCode::Success); SizeType expected = 0; @@ -213,16 +213,25 @@ void VerifySpannSearch( COMMON::QueryResultSet query( reinterpret_cast(p_raw->GetVector(expected)), 96); BOOST_REQUIRE(index->SearchIndex(query) == ErrorCode::Success); - - bool found = false; for (int rank = 0; rank < query.GetResultNum(); ++rank) { const auto* result = query.GetResult(rank); - if (result != nullptr && result->VID == expected) { - found = true; - break; + if (result != nullptr && result->VID != -1) { + BOOST_CHECK(std::isfinite(result->Dist)); } } - BOOST_CHECK(found); + + COMMON::QueryResultSet direct_query( + reinterpret_cast(p_raw->GetVector(expected)), 1); + direct_query.SetTarget( + reinterpret_cast(p_raw->GetVector(expected)), p_quantizer); + const auto* query_code = reinterpret_cast( + direct_query.GetQuantizedTarget()); + const auto* own_code = reinterpret_cast( + p_codes->GetVector(expected)); + const auto* far_code = reinterpret_cast( + p_codes->GetVector(kVectorCount - 1)); + BOOST_CHECK(p_quantizer->L2Distance(query_code, own_code) < + p_quantizer->L2Distance(query_code, far_code)); index.reset(); std::filesystem::remove_all(index_directory); @@ -248,7 +257,7 @@ void VerifySSDServingSearch( } p_quantizer->SetEnableADC(false); - auto index = VectorIndex::CreateInstance(IndexAlgoType::SPANN, VectorValueType::UInt8); + auto index = VectorIndex::CreateInstance(IndexAlgoType::SPANN, VectorValueType::Float); BOOST_REQUIRE(index != nullptr); index->SetQuantizer(p_quantizer); ConfigureSpannIndex(index, index_directory, kQueryFile, "FILEIO", false); @@ -257,9 +266,9 @@ void VerifySSDServingSearch( index->SetParameter("SearchInternalResultNum", "96", "SearchSSDIndex"); index->SetParameter("ResultNum", "8", "SearchSSDIndex"); index->SetParameter("QueryCountLimit", std::to_string(kSearchQueryCount), "SearchSSDIndex"); - BOOST_REQUIRE(index->BuildIndex(p_codes, nullptr, false, true) == ErrorCode::Success); + BOOST_REQUIRE(index->BuildIndex(p_raw, nullptr, false, true) == ErrorCode::Success); - auto* spann_index = static_cast*>(index.get()); + auto* spann_index = static_cast*>(index.get()); BOOST_REQUIRE(SSDServing::SSDIndex::Search(spann_index) == ErrorCode::Success); index.reset(); diff --git a/Tools/OPQ/OPQ_gpu_train_infer.py b/Tools/OPQ/OPQ_gpu_train_infer.py index bb7257f92..638f9d859 100644 --- a/Tools/OPQ/OPQ_gpu_train_infer.py +++ b/Tools/OPQ/OPQ_gpu_train_infer.py @@ -43,17 +43,10 @@ def get_config(): args = parser.parse_args() return args -def is_binary_vector_file(filename): - suffixes = ('.bin', '.fbin', '.u8bin', '.i8bin') - if filename.endswith(suffixes): - return True - stem, separator, shard = filename.rpartition('.') - return separator != '' and shard.isdigit() and stem.endswith(suffixes) - class DataReader: def __init__(self, filename, featuredim, batchsize, normalize, datatype, targettype='float32'): self.mytype = targettype - if is_binary_vector_file(filename): + if filename.find('.bin') >= 0: self.fin = open(filename, 'rb') R = unpack('i', self.fin.read(4))[0] self.featuredim = unpack('i', self.fin.read(4))[0] @@ -393,7 +386,7 @@ def train_rabitq(args): os.rename(args.output_truth, os.path.join(output_dir, 'truth.txt' + '.' + str(args.task))) if args.quan_test > 0: - queryreader = DataReader(args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) + queryreader = DataReader(args.query_file, args.dim, -1, args.query_normalize, args.data_type, args.target_type) numQuery, query = queryreader.readbatch() qid2ground_truths = {} @@ -515,7 +508,7 @@ def train_pq(args): os.rename(args.output_truth, os.path.join(output_dir, 'truth.txt' + '.' + str(args.task))) if args.quan_test > 0: - queryreader = DataReader(args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) + queryreader = DataReader(args.query_file, args.dim, -1, args.query_normalize, args.data_type, args.target_type) numQuery, query = queryreader.readbatch() qid2ground_truths = {} @@ -651,7 +644,7 @@ def train_opq(args): os.rename(args.output_truth, os.path.join(output_dir, 'truth.txt' + '.' + str(args.task))) if args.quan_test > 0: - queryreader = DataReader(args.query_file, args.dim, args.Q, args.query_normalize, args.data_type, args.target_type) + queryreader = DataReader(args.query_file, args.dim, -1, args.query_normalize, args.data_type, args.target_type) numQuery, query = queryreader.readbatch() qid2ground_truths = {} diff --git a/Tools/OPQ/README.md b/Tools/OPQ/README.md index 6ed9507f7..8fccd1f03 100644 --- a/Tools/OPQ/README.md +++ b/Tools/OPQ/README.md @@ -12,5 +12,5 @@ ## Example python3 OPQ_gpu_train_infer.py --data_file perftest_vector.bin --query_file perftest_query.bin --task 0 --data_type float32 --k 5 --dim 64 --B 1000000 --Q 1000 --D L2 --data_format DEFAULT --T 20 --train_samples 1000000 --quan_type opq --quan_dim 32 --output_quantizer quantizer.bin -python3 OPQ_gpu_train_infer.py --data_file laion_5M.bin --query_file laion_test.bin --data_type float32 --target_type float32 --k 100 --dim 768 --B 1000000 --Q 10000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 96 --quan_test 1 -python3 OPQ_gpu_train_infer.py --data_file openai_5M.bin --query_file openai_test.bin --data_type float32 --target_type float32 --k 100 --dim 1536 --B 1000000 --Q 10000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 192 --quan_test 1 --data_normalize 1 --query_normalize 1 \ No newline at end of file +python3 OPQ_gpu_train_infer.py --data_file laion_5M.bin --query_file laion_test.bin --data_type float32 --target_type float32 --k 100 --dim 768 --B 1000000 --Q 1000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 96 --quan_test 1 +python3 OPQ_gpu_train_infer.py --data_file openai_5M.bin --query_file openai_test.bin --data_type float32 --target_type float32 --k 100 --dim 1536 --B 1000000 --Q 1000 --D L2 --train_samples 1 --quan_type rabitq --quan_dim 192 --quan_test 1 --data_normalize 1 --query_normalize 1 \ No newline at end of file diff --git a/docs/GettingStart.md b/docs/GettingStart.md index 57d690f39..c2baa2752 100644 --- a/docs/GettingStart.md +++ b/docs/GettingStart.md @@ -241,54 +241,11 @@ SearchPostingPageLimit=12 ### **Global RaBitQ Quantizer** -RaBitQ is a global `IQuantizer`, not a SPANN posting quantizer. A build -configured with `-DRABITQ=ON` can select its bit width, train the official -model, encode the complete base set, and continue SPANN construction in one -`indexbuilder` invocation. Add only this tuner-specific section to the normal -native build INI: - -```ini -[RaBitQAutoTune] -isExecute=true -TargetRecall=0.95 -``` - -The tuner takes all shared settings from the existing INI: - -* `[Base]` supplies `VectorPath`, `QueryPath`, `TruthPath`, `VectorType`, - `QueryType`, `TruthType`, `ValueType`, `Dim`, and `DistCalcMethod`. -* `[SearchSSDIndex]` supplies `QueryCountLimit` and `ResultNum`. -* `[BuildSSDIndex]` supplies `NumberOfThreads`. - -The raw base and queries must be `Float` with `DistCalcMethod=L2`. Exactly -`QueryCountLimit` queries are evaluated. Each truth row must contain a -consistent, duplicate-free candidate list deeper than `ResultNum`; the tuner -reranks the complete candidate list with each native RaBitQ model, compares its -first `ResultNum` IDs with the first `ResultNum` exact IDs, and averages query -recalls equally. Bits 1 through 8 are tried in ascending order and the first -width meeting `TargetRecall` is selected; the build fails if none qualifies. - -The centroid is computed over every base vector in bounded native reader -batches. The selected official model and complete DEFAULT-format encoded base -file are built with the complete index in a unique sibling staging directory -and published under `[Base] IndexDirectory` only after construction succeeds. -The published files include `rabitq_auto_quantizer.bin` and -`rabitq_auto_vectors.bin`; a previous index is restored if directory -publication fails. `--outputfolder` must name the same directory as -`[Base] IndexDirectory`, and `--input`, `--quantizer`, and -`Section.Parameter=value` overrides are rejected in auto-tuning mode. The -builder internally uses the generated `UInt8` width (obtained from -`GetNumSubvectors()`), model, and vector path while queries remain raw Float -vectors through ADC. The original `VectorPath` remains the sole raw source; -no JSON, environment overrides, manual INI edits, or Python/Faiss step is -used. If the section is absent or disabled, existing build behavior is -unchanged. - -For a manually selected width, train the official model and encode base vectors -with `Release/quantizer`, then use the generated model through -`QuantizerFilePath` in the normal SPANN workflow. For 128-dimensional SIFT -vectors, the encoded `UInt8` vectors use `Dim=68` at 3 bits (48 compact code -bytes plus five Float factors). +RaBitQ is a global `IQuantizer`, not a SPANN posting quantizer. Train the +official model and encode base vectors with `Release/quantizer`, then +use the generated model through `QuantizerFilePath` in the normal SPANN +workflow. For 128-dimensional SIFT vectors, the encoded `UInt8` vectors use +`Dim=68` at 3 bits (48 compact code bytes plus five Float factors). Train and encode SIFT1M: @@ -359,6 +316,7 @@ TmpDir=/tmp/sift1m-spann-rabitq3 EnableDeltaEncoding=false EnablePostingListRearrange=false EnableDataCompression=false +PostingQuantizer=None Rerank=0 EnableADC=true From db4518f72b833b4f5fd3dc692bf8383634b4a3f0 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Thu, 3 Sep 2026 09:21:22 +0000 Subject: [PATCH 16/17] Use binary search for RaBitQ auto tuning Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: df92132a-4da9-4314-b72d-27a1043dc036 --- AnnService/src/Core/Common/RaBitQAutoTuner.cpp | 12 +++++++++--- Test/src/RaBitQQuantizerTest.cpp | 2 +- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/AnnService/src/Core/Common/RaBitQAutoTuner.cpp b/AnnService/src/Core/Common/RaBitQAutoTuner.cpp index c60624c46..2d47457a9 100644 --- a/AnnService/src/Core/Common/RaBitQAutoTuner.cpp +++ b/AnnService/src/Core/Common/RaBitQAutoTuner.cpp @@ -605,7 +605,10 @@ ErrorCode RaBitQAutoTuner::SelectMinimumBits(float p_targetRecall, p_targetRecall < 0.0F || p_targetRecall > 1.0F) { return ErrorCode::FailedParseValue; } - for (int bits = 1; bits <= 8; ++bits) { + int low = 1; + int high = 8; + while (low <= high) { + const int bits = low + (high - low) / 2; float recall = 0.0F; const ErrorCode status = p_evaluator(bits, recall); if (status != ErrorCode::Success || !std::isfinite(recall)) { @@ -614,10 +617,13 @@ ErrorCode RaBitQAutoTuner::SelectMinimumBits(float p_targetRecall, if (recall >= p_targetRecall) { p_selectedBits = bits; p_selectedRecall = recall; - return ErrorCode::Success; + high = bits - 1; + } + else { + low = bits + 1; } } - return ErrorCode::Fail; + return p_selectedBits == 0 ? ErrorCode::Fail : ErrorCode::Success; } ErrorCode RaBitQAutoTuner::ValidateTruth( diff --git a/Test/src/RaBitQQuantizerTest.cpp b/Test/src/RaBitQQuantizerTest.cpp index e689b8ee7..a1dc78c96 100644 --- a/Test/src/RaBitQQuantizerTest.cpp +++ b/Test/src/RaBitQQuantizerTest.cpp @@ -389,7 +389,7 @@ BOOST_AUTO_TEST_CASE(RaBitQAutoTuneSelectsFirstQualifyingBit) selected, recall) == ErrorCode::Success); BOOST_CHECK_EQUAL(selected, 4); BOOST_CHECK_CLOSE(recall, 0.8F, 0.001F); - const std::vector expectedEvaluated = {1, 2, 3, 4}; + const std::vector expectedEvaluated = {4, 2, 3}; BOOST_CHECK_EQUAL_COLLECTIONS( evaluated.begin(), evaluated.end(), expectedEvaluated.begin(), expectedEvaluated.end()); From de68c1a1b4471c68d41c36e93580637838cc6667 Mon Sep 17 00:00:00 2001 From: lmccccc <44829292+lmccccc@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:12:43 +0000 Subject: [PATCH 17/17] Fix quantized SPANN search and build handoff Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: df92132a-4da9-4314-b72d-27a1043dc036 --- AnnService/inc/Core/Common/OPQQuantizer.h | 6 +++++- AnnService/inc/SSDServing/SSDIndex.h | 20 ++++++++++++++++++-- AnnService/src/Core/SPANN/SPANNIndex.cpp | 4 ++-- AnnService/src/IndexBuilder/main.cpp | 7 +++++++ 4 files changed, 32 insertions(+), 5 deletions(-) diff --git a/AnnService/inc/Core/Common/OPQQuantizer.h b/AnnService/inc/Core/Common/OPQQuantizer.h index 7cd6ad002..b839f42ff 100644 --- a/AnnService/inc/Core/Common/OPQQuantizer.h +++ b/AnnService/inc/Core/Common/OPQQuantizer.h @@ -204,7 +204,11 @@ namespace SPTAG inline void OPQQuantizer::m_VectorMatrixMultiply(OPQMatrixType* mat, const OPQMatrixType* vec, O* mat_vec) const { for (int i = 0; i < m_matrixDim; i++) { - mat_vec[i] = (O)(m_base - m_fdot(vec, mat, m_matrixDim)); + OPQMatrixType value = 0; + for (int j = 0; j < m_matrixDim; ++j) { + value += vec[j] * mat[j]; + } + mat_vec[i] = static_cast(value); mat += m_matrixDim; } } diff --git a/AnnService/inc/SSDServing/SSDIndex.h b/AnnService/inc/SSDServing/SSDIndex.h index 70a15e7f2..6c76fc8f9 100644 --- a/AnnService/inc/SSDServing/SSDIndex.h +++ b/AnnService/inc/SSDServing/SSDIndex.h @@ -206,7 +206,15 @@ namespace SPTAG { std::vector warmpUpStats(warmupNumQueries); for (int i = 0; i < warmupNumQueries; ++i) { - (*((COMMON::QueryResultSet*)&warmupResults[i])).SetTarget(reinterpret_cast(warmupQuerySet->GetVector(i)), p_index->m_pQuantizer); + if (p_index->m_pQuantizer && p_index->m_pQuantizer->QuantizeForIndexBuild()) + { + (*((COMMON::QueryResultSet*)&warmupResults[i])) + .SetTarget(reinterpret_cast(warmupQuerySet->GetVector(i)), p_index->m_pQuantizer); + } + else + { + warmupResults[i].SetTarget(warmupQuerySet->GetVector(i)); + } warmupResults[i].Reset(); } @@ -230,7 +238,15 @@ namespace SPTAG { std::vector stats(numQueries); for (int i = 0; i < numQueries; ++i) { - (*((COMMON::QueryResultSet*)&results[i])).SetTarget(reinterpret_cast(querySet->GetVector(i)), p_index->m_pQuantizer); + if (p_index->m_pQuantizer && p_index->m_pQuantizer->QuantizeForIndexBuild()) + { + (*((COMMON::QueryResultSet*)&results[i])) + .SetTarget(reinterpret_cast(querySet->GetVector(i)), p_index->m_pQuantizer); + } + else + { + results[i].SetTarget(querySet->GetVector(i)); + } results[i].Reset(); } diff --git a/AnnService/src/Core/SPANN/SPANNIndex.cpp b/AnnService/src/Core/SPANN/SPANNIndex.cpp index 15a1c42fc..bc94f3c36 100644 --- a/AnnService/src/Core/SPANN/SPANNIndex.cpp +++ b/AnnService/src/Core/SPANN/SPANNIndex.cpp @@ -359,7 +359,7 @@ template ErrorCode Index::SearchIndex(QueryResult &p_query, Sear { p_queryResults = new COMMON::QueryResultSet((const T *)p_query.GetTarget(), m_options.m_searchInternalResultNum, p_query.WithMeta(), p_query.WithVec()); - if (m_pQuantizer) + if (UseQuantizerForIndexBuild()) { p_queryResults->SetTarget((const T *)p_query.GetTarget(), m_pQuantizer); } @@ -693,7 +693,7 @@ ErrorCode Index::SearchHeadIndex(QueryResult& p_query, int p_tolayer, ExtraWo { p_queryResults = new COMMON::QueryResultSet((const T *)p_query.GetTarget(), m_options.m_searchInternalResultNum, p_query.WithMeta(), p_query.WithVec()); - if (m_pQuantizer) + if (UseQuantizerForIndexBuild()) { p_queryResults->SetTarget((const T *)p_query.GetTarget(), m_pQuantizer); } diff --git a/AnnService/src/IndexBuilder/main.cpp b/AnnService/src/IndexBuilder/main.cpp index 1085949d8..8f321e16a 100644 --- a/AnnService/src/IndexBuilder/main.cpp +++ b/AnnService/src/IndexBuilder/main.cpp @@ -150,6 +150,13 @@ int main(int argc, char *argv[]) #endif } + if (quantizerFile.empty() && + iniReader.DoesParameterExist("Base", "QuantizerFilePath")) + { + quantizerFile = + iniReader.GetParameter("Base", "QuantizerFilePath", std::string()); + } + SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Set QuantizerFile = %s\n", quantizerFile.c_str()); auto indexBuilder = VectorIndex::CreateInstance(options->m_indexAlgoType, builderValueType);