diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 5086dea3c..6d59c35a2 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -11,7 +11,6 @@ on: options: - benchmarks-all - benchmarks-default - - bm-basics-int8-single - bm-basics-fp32-single - bm-basics-fp32-multi - bm-basics-fp64-single @@ -20,6 +19,7 @@ on: - bm-basics-bf16-multi - bm-basics-fp16-single - bm-basics-fp16-multi + - bm-basics-int8-single - bm-batch-iter-fp32-single - bm-batch-iter-fp32-multi - bm-batch-iter-fp64-single @@ -28,6 +28,7 @@ on: - bm-batch-iter-bf16-multi - bm-batch-iter-fp16-single - bm-batch-iter-fp16-multi + - bm-batch-iter-int8-single - bm-updated-fp32-single - bm-spaces description: 'Benchmarks set to run' diff --git a/tests/benchmark/data/hnsw_indices/hnsw_indices_all.txt b/tests/benchmark/data/hnsw_indices/hnsw_indices_all.txt index bafe0b613..8a5696974 100644 --- a/tests/benchmark/data/hnsw_indices/hnsw_indices_all.txt +++ b/tests/benchmark/data/hnsw_indices/hnsw_indices_all.txt @@ -25,5 +25,5 @@ https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/dbpedia-cosine-dim768-fp https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/fashion_images_multi_value-cosine-dim512-M64-efc512-fp16.hnsw_v3 https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/fashion_images_multi_value-cosine-dim512-fp16-test_vectors.raw -https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia-cosine-dim1024-M64-efc512-int8.hnsw_v3 -https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia-cosine-dim1024-int8-test_vectors.raw \ No newline at end of file +https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia_single-cosine-dim1024-M64-efc512-int8.hnsw_v3 +https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia_single-cosine-dim1024-int8-test_vectors.raw diff --git a/tests/benchmark/data/hnsw_indices/hnsw_indices_basic_int8.txt b/tests/benchmark/data/hnsw_indices/hnsw_indices_basic_int8.txt index a4af55a88..18a9454b3 100644 --- a/tests/benchmark/data/hnsw_indices/hnsw_indices_basic_int8.txt +++ b/tests/benchmark/data/hnsw_indices/hnsw_indices_basic_int8.txt @@ -1,2 +1,2 @@ -https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia-cosine-dim1024-M64-efc512-int8.hnsw_v3 -https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia-cosine-dim1024-int8-test_vectors.raw +https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia_single-cosine-dim1024-M64-efc512-int8.hnsw_v3 +https://dev.cto.redis.s3.amazonaws.com/VectorSimilarity/wipedia_single-cosine-dim1024-int8-test_vectors.raw diff --git a/tests/benchmark/data/scripts/download_dataset.py b/tests/benchmark/data/scripts/download_dataset.py new file mode 100644 index 000000000..70ca0a4a1 --- /dev/null +++ b/tests/benchmark/data/scripts/download_dataset.py @@ -0,0 +1,50 @@ +# This file is a template file for downloading datasets +# In this version, it downloads the "wipedia_single" dataset used for int8 +# Refrain from pushing changes unless necessary +from datasets import load_dataset +import numpy as np +import os +import h5py +from math import ceil +from tqdm import tqdm +INT8_KEY = 'emb_int8' +DATASET = 'wikipedia-1024_eng_v3_single' +hdf5_output_file_name = "%s.hdf5" %DATASET + +lang = "en" #Use the English Wikipedia subset + +num_vectors_train = 1_000_000 +num_vectors_test = 10_000 +num_vectors = num_vectors_train + num_vectors_test + +dim = 1024 +docs = load_dataset("Cohere/wikipedia-2023-11-embed-multilingual-v3-int8-binary", lang, split="train", streaming=True) +label_size = 1 +data = np.empty((num_vectors//label_size,label_size, dim), dtype=np.int8) + + +ids = [] +counter = 0 +label_index = 0 + +with tqdm(total=num_vectors) as progress_bar: + for doc in docs: + if counter == num_vectors: + break + ids.append(doc['_id']) + emb = doc['emb_int8'] + data[label_index, counter % label_size] = emb + counter += 1 + if counter % label_size == 0: + label_index += 1 + progress_bar.update(1) + +train_data = data[:num_vectors_train // label_size].reshape(-1, dim) +test_data = data[num_vectors_train // label_size:].reshape(-1, dim) + +print(f"Train data shape: {train_data.shape}") +print(f"Test data shape: {test_data.shape}") + +with h5py.File(hdf5_output_file_name, 'w') as hdf5_file: + hdf5_file.create_dataset('train', data=train_data) + hdf5_file.create_dataset('test', data=test_data) diff --git a/tests/benchmark/data/serializer.py b/tests/benchmark/data/scripts/serializer.py similarity index 93% rename from tests/benchmark/data/serializer.py rename to tests/benchmark/data/scripts/serializer.py index 0f0e60a0a..e6d1c8a46 100644 --- a/tests/benchmark/data/serializer.py +++ b/tests/benchmark/data/scripts/serializer.py @@ -136,6 +136,13 @@ 'metric': VecSimMetric_L2, 'skipRaw': True, }, + { + 'filename': 'wikipedia_single-1024_eng_v3', + 'nickname': 'wipedia_single', + 'dim': 1024, + 'type': VecSimType_INT8, + 'metric': VecSimMetric_Cosine, + }, ] TYPES_ATTR = { @@ -143,6 +150,7 @@ VecSimType_FLOAT64: {"size_in_bytes": 8, "vector_type": np.float64}, VecSimType_BFLOAT16: {"size_in_bytes": 2, "vector_type": bfloat16}, VecSimType_FLOAT16: {"size_in_bytes": 2, "vector_type": np.float16}, + VecSimType_INT8: {"size_in_bytes": 1, "vector_type": np.int8} } @@ -190,7 +198,9 @@ def serialize(files=DEFAULT_FILES): elif hnswparams.type == VecSimType_FLOAT16: serialized_file_name = serialized_file_name + '-fp16' serialized_raw_name = serialized_raw_name + '-fp16' - + elif hnswparams.type == VecSimType_INT8: + serialized_file_name = serialized_file_name + '-int8' + serialized_raw_name = serialized_raw_name + '-int8' print('first, exporting test set to binary') if not file.get('skipRaw', False): test = f['test'] @@ -201,6 +211,8 @@ def serialize(files=DEFAULT_FILES): test = np.array(test_set, dtype=bfloat16) elif hnswparams.type == VecSimType_FLOAT16: test = test.astype(np.float16) + elif hnswparams.type == VecSimType_INT8: + test = test.astype(np.int8) print(f"creating test set of {len(test)} vectors") with open(os.path.join(location, serialized_raw_name + '-test_vectors.raw'), 'wb') as testfile: for vec in test: @@ -222,6 +234,8 @@ def serialize(files=DEFAULT_FILES): data = np.array(data_set, dtype=bfloat16) elif hnswparams.type == VecSimType_FLOAT16: data = data.astype(np.float16) + elif hnswparams.type == VecSimType_INT8: + data = data.astype(np.int8) print(f"creating index with {hnswparams.initialCapacity} vectors") for label, cur in enumerate(data): for vec in cur if hnswparams.multi else [cur]: diff --git a/tests/benchmark/run_files/bm_basics_single_int8.cpp b/tests/benchmark/run_files/bm_basics_single_int8.cpp index 5e874e335..226ad39c2 100644 --- a/tests/benchmark/run_files/bm_basics_single_int8.cpp +++ b/tests/benchmark/run_files/bm_basics_single_int8.cpp @@ -9,16 +9,16 @@ bool BM_VecSimGeneral::is_multi = false; size_t BM_VecSimGeneral::n_queries = 10000; -size_t BM_VecSimGeneral::n_vectors = 999424; +size_t BM_VecSimGeneral::n_vectors = 1000000; size_t BM_VecSimGeneral::dim = 1024; size_t BM_VecSimGeneral::M = 64; size_t BM_VecSimGeneral::EF_C = 512; tieredIndexMock BM_VecSimGeneral::mock_thread_pool{}; const char *BM_VecSimGeneral::hnsw_index_file = - "tests/benchmark/data/wipedia-cosine-dim1024-M64-efc512-int8.hnsw_v3"; + "tests/benchmark/data/wipedia_single-cosine-dim1024-M64-efc512-int8.hnsw_v3"; const char *BM_VecSimGeneral::test_queries_file = - "tests/benchmark/data/wipedia-cosine-dim1024-int8-test_vectors.raw"; + "tests/benchmark/data/wipedia_single-cosine-dim1024-int8-test_vectors.raw"; #define BM_FUNC_NAME(bm_func, algo) bm_func##_##algo##_Single #define BM_ADD_LABEL AddLabel_Single diff --git a/tests/benchmark/run_files/bm_batch_iterator_single_int8.cpp b/tests/benchmark/run_files/bm_batch_iterator_single_int8.cpp index 18859192d..1ad03d5b4 100644 --- a/tests/benchmark/run_files/bm_batch_iterator_single_int8.cpp +++ b/tests/benchmark/run_files/bm_batch_iterator_single_int8.cpp @@ -3,7 +3,7 @@ bool BM_VecSimGeneral::is_multi = false; size_t BM_VecSimGeneral::n_queries = 10000; -size_t BM_VecSimGeneral::n_vectors = 999424; +size_t BM_VecSimGeneral::n_vectors = 1000000; size_t BM_VecSimGeneral::dim = 1024; size_t BM_VecSimGeneral::M = 64; size_t BM_VecSimGeneral::EF_C = 512; @@ -11,9 +11,9 @@ size_t BM_VecSimGeneral::block_size = 1024; tieredIndexMock BM_VecSimGeneral::mock_thread_pool{}; const char *BM_VecSimGeneral::hnsw_index_file = - "tests/benchmark/data/wipedia-cosine-dim1024-M64-efc512-int8.hnsw_v3"; + "tests/benchmark/data/wipedia_single-cosine-dim1024-M64-efc512-int8.hnsw_v3"; const char *BM_VecSimGeneral::test_queries_file = - "tests/benchmark/data/wipedia-cosine-dim1024-int8-test_vectors.raw"; + "tests/benchmark/data/wipedia_single-cosine-dim1024-int8-test_vectors.raw"; #define BM_FUNC_NAME(bm_func, algo) algo##_##bm_func##_Single