Skip to content

Commit a7116ed

Browse files
Mistobaanmeta-codesync[bot]
Mistobaan
authored and
meta-codesync[bot]
committed
Add TurboQuant (CPU) (#5049)
Summary: This PR adds initial TurboQuant (see #4990) support to Faiss and integrates it into the main codepaths needed for local evaluation. Changes in this PR: - add `IndexTurboQuantMSE` and the underlying `TurboQuantizer` implementation - add cloning and serialization support for TurboQuant indexes - expose TurboQuant through the Python SWIG bindings - extend quantizer benchmarking to cover TurboQuant and additional datasets - add unit tests for TurboQuant reconstruction behavior ## Preliminary Benchmarks These are preliminary local results from a macOS CPU-only run. Command: ```bash python bench_quantizer.py glove 100x4 turboquant pq rq eval on glove 100x4 maxtrain=100000 No training set: training on database ===== PQ training time: 0.594 s encode time: 0.317 reconstruction error: 0.010 recall@1: 0.7036 recons_err_compat 0.100 code_size: 50 B/vector ===== RQ training time: 208.554 s max_beam_size=1 encode time: 2.977 reconstruction error: 0.027 recall@1: 0.6034 recons_err_compat 0.162 code_size: 50 B/vector max_beam_size=2 encode time: 5.774 reconstruction error: 0.023 recall@1: 0.6280 recons_err_compat 0.151 code_size: 50 B/vector max_beam_size=4 encode time: 12.271 reconstruction error: 0.020 recall@1: 0.6512 recons_err_compat 0.140 code_size: 50 B/vector max_beam_size=8 encode time: 25.617 reconstruction error: 0.018 recall@1: 0.6751 recons_err_compat 0.131 code_size: 50 B/vector max_beam_size=16 encode time: 50.257 reconstruction error: 0.016 recall@1: 0.6890 recons_err_compat 0.123 code_size: 50 B/vector max_beam_size=32 encode time: 105.272 reconstruction error: 0.014 recall@1: 0.7067 recons_err_compat 0.116 code_size: 50 B/vector ===== TurboQuant training time: 0.002 s encode time: 0.080 reconstruction error: 0.009 recall@1: 0.7189 recons_err_compat 0.095 code_size: 50 B/vector ``` ### Initial takeaway On this macOS CPU-only benchmark, TurboQuant shows: - better `recall@1` than PQ at the same code size - lower reconstruction error than PQ - higher `recall@1` than the best RQ setting tested here - much lower training and encoding cost than RQ ### TODO - [x] add JQL support (only MSE at the moment) - reproduce the paper benchmarks with same datasets (need a beefy machine but script is ready) - ~add CUDA support~ (will be in another PR if this is accepted) Pull Request resolved: #5049 Test Plan: buck build fbcode//faiss:faiss buck test fbcode//faiss/tests:test_factory buck test fbcode//faiss/tests:test_scalar_quantizer buck test fbcode//faiss/tests:test_scalar_quantizer_correctness buck build fbcode//faiss/tests/fb:faiss_fuzzer_helpers Reviewed By: mdouze Differential Revision: D102195204 Pulled By: mnorris11 fbshipit-source-id: a302541a1fc5c229e697f28e639cba633babd072
1 parent dfce6e9 commit a7116ed

18 files changed

Lines changed: 1270 additions & 17 deletions

‎faiss/factory_tools.cpp‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,11 @@ const std::map<faiss::ScalarQuantizer::QuantizerType, std::string> sq_types = {
3838
{faiss::ScalarQuantizer::QT_bf16, "SQbf16"},
3939
{faiss::ScalarQuantizer::QT_8bit_direct_signed, "SQ8_direct_signed"},
4040
{faiss::ScalarQuantizer::QT_8bit_direct, "SQ8_direct"},
41+
{faiss::ScalarQuantizer::QT_1bit_tqmse, "SQtqmse1"},
42+
{faiss::ScalarQuantizer::QT_2bit_tqmse, "SQtqmse2"},
43+
{faiss::ScalarQuantizer::QT_3bit_tqmse, "SQtqmse3"},
44+
{faiss::ScalarQuantizer::QT_4bit_tqmse, "SQtqmse4"},
45+
{faiss::ScalarQuantizer::QT_8bit_tqmse, "SQtqmse8"},
4146
};
4247

4348
int get_hnsw_M(const faiss::IndexHNSW* index) {

‎faiss/impl/ScalarQuantizer.cpp‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,15 +38,29 @@ ScalarQuantizer::ScalarQuantizer() {}
3838

3939
void ScalarQuantizer::set_derived_sizes() {
4040
switch (qtype) {
41+
case QT_1bit_tqmse:
42+
code_size = (d + 7) / 8;
43+
bits = 1;
44+
break;
45+
case QT_2bit_tqmse:
46+
code_size = (d * 2 + 7) / 8;
47+
bits = 2;
48+
break;
49+
case QT_3bit_tqmse:
50+
code_size = (d * 3 + 7) / 8;
51+
bits = 3;
52+
break;
4153
case QT_8bit:
4254
case QT_8bit_uniform:
4355
case QT_8bit_direct:
4456
case QT_8bit_direct_signed:
57+
case QT_8bit_tqmse:
4558
code_size = d;
4659
bits = 8;
4760
break;
4861
case QT_4bit:
4962
case QT_4bit_uniform:
63+
case QT_4bit_tqmse:
5064
code_size = (d + 1) / 2;
5165
bits = 4;
5266
break;
@@ -115,6 +129,21 @@ void ScalarQuantizer::train(size_t n, const float* x) {
115129
case QT_8bit_direct_signed:
116130
// no training necessary
117131
break;
132+
case QT_1bit_tqmse:
133+
scalar_quantizer::train_TurboQuantMSE(d, 1, trained);
134+
break;
135+
case QT_2bit_tqmse:
136+
scalar_quantizer::train_TurboQuantMSE(d, 2, trained);
137+
break;
138+
case QT_3bit_tqmse:
139+
scalar_quantizer::train_TurboQuantMSE(d, 3, trained);
140+
break;
141+
case QT_4bit_tqmse:
142+
scalar_quantizer::train_TurboQuantMSE(d, 4, trained);
143+
break;
144+
case QT_8bit_tqmse:
145+
scalar_quantizer::train_TurboQuantMSE(d, 8, trained);
146+
break;
118147
default:
119148
break;
120149
}

‎faiss/impl/ScalarQuantizer.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,11 @@ struct ScalarQuantizer : Quantizer {
3434
QT_8bit_direct_signed, ///< fast indexing of signed int8s ranging from
3535
///< [-128 to 127]
3636
QT_0bit, ///< 0 bits per component, centroid-only distance (for IVF)
37+
QT_1bit_tqmse, ///< TurboQuant MSE-optimized, 1 bit per component
38+
QT_2bit_tqmse, ///< TurboQuant MSE-optimized, 2 bits per component
39+
QT_3bit_tqmse, ///< TurboQuant MSE-optimized, 3 bits per component
40+
QT_4bit_tqmse, ///< TurboQuant MSE-optimized, 4 bits per component
41+
QT_8bit_tqmse, ///< TurboQuant MSE-optimized, 8 bits per component
3742
QT_count
3843
};
3944

‎faiss/impl/index_read.cpp‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -978,7 +978,7 @@ void read_ScalarQuantizer(
978978
READ1(qtype_int);
979979
FAISS_THROW_IF_NOT_FMT(
980980
qtype_int >= ScalarQuantizer::QT_8bit &&
981-
qtype_int <= ScalarQuantizer::QT_8bit_direct_signed,
981+
qtype_int < ScalarQuantizer::QT_count,
982982
"invalid ScalarQuantizer qtype %d",
983983
qtype_int);
984984
ivsc->qtype = static_cast<ScalarQuantizer::QuantizerType>(qtype_int);
@@ -1017,6 +1017,21 @@ void read_ScalarQuantizer(
10171017
case ScalarQuantizer::QT_count:
10181018
expected = 0;
10191019
break;
1020+
case ScalarQuantizer::QT_1bit_tqmse:
1021+
expected = 2 + 1; // 2^bits centroids + (2^bits - 1) boundaries
1022+
break;
1023+
case ScalarQuantizer::QT_2bit_tqmse:
1024+
expected = 4 + 3;
1025+
break;
1026+
case ScalarQuantizer::QT_3bit_tqmse:
1027+
expected = 8 + 7;
1028+
break;
1029+
case ScalarQuantizer::QT_4bit_tqmse:
1030+
expected = 16 + 15;
1031+
break;
1032+
case ScalarQuantizer::QT_8bit_tqmse:
1033+
expected = 256 + 255;
1034+
break;
10201035
}
10211036
if (ivsc->trained.empty() && expected > 0) {
10221037
// Empty trained is only valid for untrained indices.

‎faiss/impl/scalar_quantizer/quantizers.h‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@
77

88
#pragma once
99

10+
#include <algorithm>
11+
12+
#include <faiss/impl/FaissAssert.h>
1013
#include <faiss/impl/ScalarQuantizer.h>
1114
#include <faiss/impl/simdlib/simdlib_dispatch.h>
1215
#include <faiss/utils/bf16.h>
@@ -113,6 +116,90 @@ struct QuantizerTemplate<
113116
}
114117
};
115118

119+
/*******************************************************************
120+
* TurboQuant MSE quantizer
121+
*******************************************************************/
122+
template <int NBits, SIMDLevel SL>
123+
struct QuantizerTurboQuantMSE;
124+
125+
template <int NBits>
126+
struct QuantizerTurboQuantMSE<NBits, SIMDLevel::NONE>
127+
: ScalarQuantizer::SQuantizer {
128+
static_assert(NBits >= 1 && NBits <= 8);
129+
130+
static constexpr size_t kCentroidsCount = size_t(1) << NBits;
131+
static constexpr uint16_t kIndexMask =
132+
static_cast<uint16_t>((1u << NBits) - 1);
133+
134+
const size_t d;
135+
const float* centroids;
136+
const float* boundaries;
137+
138+
QuantizerTurboQuantMSE(size_t d_in, const std::vector<float>& trained)
139+
: d(d_in), centroids(nullptr), boundaries(nullptr) {
140+
FAISS_THROW_IF_NOT(trained.size() == 2 * kCentroidsCount - 1);
141+
centroids = trained.data();
142+
boundaries = trained.data() + kCentroidsCount;
143+
}
144+
145+
FAISS_ALWAYS_INLINE uint8_t select_index(float x) const {
146+
return static_cast<uint8_t>(
147+
std::upper_bound(
148+
boundaries, boundaries + (kCentroidsCount - 1), x) -
149+
boundaries);
150+
}
151+
152+
FAISS_ALWAYS_INLINE void encode_index(uint8_t idx, uint8_t* code, size_t i)
153+
const {
154+
const size_t bit_offset = i * NBits;
155+
const size_t byte_offset = bit_offset >> 3;
156+
const size_t bit_shift = bit_offset & 7;
157+
const uint16_t packed = static_cast<uint16_t>(idx & kIndexMask)
158+
<< bit_shift;
159+
code[byte_offset] |= packed & 0xff;
160+
if (bit_shift + NBits > 8) {
161+
code[byte_offset + 1] |= packed >> 8;
162+
}
163+
}
164+
165+
FAISS_ALWAYS_INLINE uint8_t
166+
decode_index(const uint8_t* code, size_t i) const {
167+
const size_t bit_offset = i * NBits;
168+
const size_t byte_offset = bit_offset >> 3;
169+
const size_t bit_shift = bit_offset & 7;
170+
171+
uint16_t packed = code[byte_offset];
172+
if (bit_shift + NBits > 8) {
173+
packed |= static_cast<uint16_t>(code[byte_offset + 1]) << 8;
174+
}
175+
return static_cast<uint8_t>((packed >> bit_shift) & kIndexMask);
176+
}
177+
178+
void encode_vector(const float* x, uint8_t* code) const final {
179+
for (size_t i = 0; i < d; i++) {
180+
encode_index(select_index(x[i]), code, i);
181+
}
182+
}
183+
184+
void decode_vector(const uint8_t* code, float* x) const final {
185+
for (size_t i = 0; i < d; i++) {
186+
x[i] = centroids[decode_index(code, i)];
187+
}
188+
}
189+
190+
FAISS_ALWAYS_INLINE float reconstruct_component(
191+
const uint8_t* code,
192+
size_t i) const {
193+
return centroids[decode_index(code, i)];
194+
}
195+
};
196+
197+
template <int NBits, SIMDLevel SL>
198+
struct QuantizerTurboQuantMSE : QuantizerTurboQuantMSE<NBits, SIMDLevel::NONE> {
199+
using QuantizerTurboQuantMSE<NBits, SIMDLevel::NONE>::
200+
QuantizerTurboQuantMSE;
201+
};
202+
116203
/*******************************************************************
117204
* FP16 quantizer
118205
*******************************************************************/

‎faiss/impl/scalar_quantizer/sq-avx2.cpp‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99

1010
#include <faiss/impl/simdlib/simdlib_avx2.h>
1111

12+
#include <cstring>
13+
1214
#include <faiss/impl/scalar_quantizer/codecs.h>
1315
#include <faiss/impl/scalar_quantizer/distance_computers.h>
1416
#include <faiss/impl/scalar_quantizer/quantizers.h>
@@ -21,6 +23,61 @@ namespace scalar_quantizer {
2123

2224
using simd8float32 = faiss::simd8float32_tpl<SIMDLevel::AVX2>;
2325

26+
namespace {
27+
28+
FAISS_ALWAYS_INLINE uint16_t load_u16(const uint8_t* ptr) {
29+
uint16_t value;
30+
std::memcpy(&value, ptr, sizeof(value));
31+
return value;
32+
}
33+
34+
FAISS_ALWAYS_INLINE uint32_t load_u32(const uint8_t* ptr) {
35+
uint32_t value;
36+
std::memcpy(&value, ptr, sizeof(value));
37+
return value;
38+
}
39+
40+
FAISS_ALWAYS_INLINE uint32_t load_u24(const uint8_t* ptr) {
41+
return static_cast<uint32_t>(ptr[0]) |
42+
(static_cast<uint32_t>(ptr[1]) << 8) |
43+
(static_cast<uint32_t>(ptr[2]) << 16);
44+
}
45+
46+
FAISS_ALWAYS_INLINE __m256i unpack_8x1bit_to_u32(const uint8_t* code, int i) {
47+
const uint32_t packed = code[static_cast<size_t>(i) >> 3];
48+
const __m256i shifts = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
49+
const __m256i indices =
50+
_mm256_srlv_epi32(_mm256_set1_epi32(packed), shifts);
51+
return _mm256_and_si256(indices, _mm256_set1_epi32(0x1));
52+
}
53+
54+
FAISS_ALWAYS_INLINE __m256i unpack_8x2bit_to_u32(const uint8_t* code, int i) {
55+
const uint32_t packed = load_u16(code + (static_cast<size_t>(i) >> 2));
56+
const __m256i shifts = _mm256_setr_epi32(0, 2, 4, 6, 8, 10, 12, 14);
57+
const __m256i indices =
58+
_mm256_srlv_epi32(_mm256_set1_epi32(packed), shifts);
59+
return _mm256_and_si256(indices, _mm256_set1_epi32(0x3));
60+
}
61+
62+
FAISS_ALWAYS_INLINE __m256i unpack_8x3bit_to_u32(const uint8_t* code, int i) {
63+
const uint32_t packed =
64+
load_u24(code + ((static_cast<size_t>(i) >> 3) * 3));
65+
const __m256i shifts = _mm256_setr_epi32(0, 3, 6, 9, 12, 15, 18, 21);
66+
const __m256i indices =
67+
_mm256_srlv_epi32(_mm256_set1_epi32(packed), shifts);
68+
return _mm256_and_si256(indices, _mm256_set1_epi32(0x7));
69+
}
70+
71+
FAISS_ALWAYS_INLINE __m256i unpack_8x4bit_to_u32(const uint8_t* code, int i) {
72+
const uint32_t packed = load_u32(code + (static_cast<size_t>(i) >> 1));
73+
const __m256i shifts = _mm256_setr_epi32(0, 4, 8, 12, 16, 20, 24, 28);
74+
const __m256i indices =
75+
_mm256_srlv_epi32(_mm256_set1_epi32(packed), shifts);
76+
return _mm256_and_si256(indices, _mm256_set1_epi32(0xf));
77+
}
78+
79+
} // namespace
80+
2481
/**********************************************************
2582
* Codecs
2683
**********************************************************/
@@ -168,6 +225,56 @@ struct QuantizerTemplate<
168225
}
169226
};
170227

228+
/**********************************************************
229+
* TurboQuant MSE quantizer
230+
**********************************************************/
231+
232+
#define DEFINE_TQMSE_AVX2_SPECIALIZATION(NBITS, INDEX_EXPR) \
233+
template <> \
234+
struct QuantizerTurboQuantMSE<NBITS, SIMDLevel::AVX2> \
235+
: QuantizerTurboQuantMSE<NBITS, SIMDLevel::NONE> { \
236+
using Base = QuantizerTurboQuantMSE<NBITS, SIMDLevel::NONE>; \
237+
\
238+
QuantizerTurboQuantMSE(size_t d, const std::vector<float>& trained) \
239+
: Base(d, trained) { \
240+
assert(d % 8 == 0); \
241+
} \
242+
\
243+
FAISS_ALWAYS_INLINE simd8float32 \
244+
reconstruct_8_components(const uint8_t* code, int i) const { \
245+
const __m256i indices = (INDEX_EXPR); \
246+
return simd8float32(_mm256_i32gather_ps( \
247+
this->centroids, indices, sizeof(float))); \
248+
} \
249+
}
250+
251+
DEFINE_TQMSE_AVX2_SPECIALIZATION(1, unpack_8x1bit_to_u32(code, i));
252+
DEFINE_TQMSE_AVX2_SPECIALIZATION(2, unpack_8x2bit_to_u32(code, i));
253+
DEFINE_TQMSE_AVX2_SPECIALIZATION(3, unpack_8x3bit_to_u32(code, i));
254+
DEFINE_TQMSE_AVX2_SPECIALIZATION(4, unpack_8x4bit_to_u32(code, i));
255+
256+
#undef DEFINE_TQMSE_AVX2_SPECIALIZATION
257+
258+
template <>
259+
struct QuantizerTurboQuantMSE<8, SIMDLevel::AVX2>
260+
: QuantizerTurboQuantMSE<8, SIMDLevel::NONE> {
261+
using Base = QuantizerTurboQuantMSE<8, SIMDLevel::NONE>;
262+
263+
QuantizerTurboQuantMSE(size_t d, const std::vector<float>& trained)
264+
: Base(d, trained) {
265+
assert(d % 8 == 0);
266+
}
267+
268+
FAISS_ALWAYS_INLINE simd8float32
269+
reconstruct_8_components(const uint8_t* code, int i) const {
270+
const __m128i packed = _mm_loadl_epi64(
271+
(const __m128i*)(code + static_cast<size_t>(i)));
272+
const __m256i indices = _mm256_cvtepu8_epi32(packed);
273+
return simd8float32(
274+
_mm256_i32gather_ps(this->centroids, indices, sizeof(float)));
275+
}
276+
};
277+
171278
/**********************************************************
172279
* FP16 Quantizer
173280
**********************************************************/

0 commit comments

Comments
 (0)