Skip to content

Commit 1d6bb52

Browse files
mayya-sharipovamayya-sharipova
authored
Disallow dot_product and max_inner_product for int8_hnsw GPU type (elastic#137043)
* Disallow dot_product and max_inner_product for int8_hnsw GPU type (elastic#136881) For int8_hnsw, during merges we get quantized vectors from Lucene files but dropping for each quantized vector its correction factor. For cosine and euclidean metrics this correction factor is not important, but for dot_product and max_inner_product metrics, they are important. It means that that currently for dot_product and max_inner_product metrics, GPU graph building doesn't work well, and may produce bad recall. This PR does the following: - disallows max_inner_product for int8 - substitutes internally dot_product with cosine Alternatives: for most datasets (really majority), we can substitute dot_product with cosine. But there are some datasets that require max_inner_product, and for this our "int8_hsnw" will not work, and "hnsw" should be used instead * Fix Test failure
1 parent 56fb21a commit 1d6bb52

File tree

9 files changed

+153
-19
lines changed

9 files changed

+153
-19
lines changed

‎docs/changelog/136881.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 136881
2+
summary: Disallow `max_inner_product`, swap `dot_product` for `cosine` for int8_hnsw GPU type
3+
area: Search
4+
type: bug
5+
issues: []

‎server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2267,6 +2267,10 @@ public DenseVectorFieldType(
22672267
this.isSyntheticSource = isSyntheticSource;
22682268
}
22692269

2270+
public VectorSimilarity similarity() {
2271+
return similarity;
2272+
}
2273+
22702274
@Override
22712275
public String typeName() {
22722276
return CONTENT_TYPE;
@@ -2863,7 +2867,7 @@ public KnnVectorsFormat getKnnVectorsFormatForField(KnnVectorsFormat defaultForm
28632867
// if plugins provided alternative KnnVectorsFormat for this indexOptions, use it instead of standard
28642868
KnnVectorsFormat extraKnnFormat = null;
28652869
for (VectorsFormatProvider vectorsFormatProvider : extraVectorsFormatProviders) {
2866-
extraKnnFormat = vectorsFormatProvider.getKnnVectorsFormat(indexSettings, indexOptions);
2870+
extraKnnFormat = vectorsFormatProvider.getKnnVectorsFormat(indexSettings, indexOptions, fieldType().similarity());
28672871
if (extraKnnFormat != null) {
28682872
break;
28692873
}

‎server/src/main/java/org/elasticsearch/index/mapper/vectors/VectorsFormatProvider.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,12 @@ public interface VectorsFormatProvider {
2424
*
2525
* @param indexSettings The index settings.
2626
* @param indexOptions The dense vector index options.
27+
* @param similarity The vector similarity function.
2728
* @return A KnnVectorsFormat instance.
2829
*/
29-
KnnVectorsFormat getKnnVectorsFormat(IndexSettings indexSettings, DenseVectorFieldMapper.DenseVectorIndexOptions indexOptions);
30+
KnnVectorsFormat getKnnVectorsFormat(
31+
IndexSettings indexSettings,
32+
DenseVectorFieldMapper.DenseVectorIndexOptions indexOptions,
33+
DenseVectorFieldMapper.VectorSimilarity similarity
34+
);
3035
}

‎server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,15 @@ public static DenseVectorFieldMapper.DenseVectorIndexOptions randomGpuSupportedI
9696
);
9797
}
9898

99+
public static DenseVectorFieldMapper.VectorSimilarity randomGPUSupportedSimilarity(
100+
DenseVectorFieldMapper.VectorIndexType vectorIndexType
101+
) {
102+
if (vectorIndexType == DenseVectorFieldMapper.VectorIndexType.INT8_HNSW) {
103+
return randomFrom(VectorSimilarity.L2_NORM, VectorSimilarity.COSINE, VectorSimilarity.DOT_PRODUCT);
104+
}
105+
return randomFrom(VectorSimilarity.values());
106+
}
107+
99108
public static DenseVectorFieldMapper.DenseVectorIndexOptions randomIndexOptionsAll() {
100109
List<DenseVectorFieldMapper.DenseVectorIndexOptions> options = new ArrayList<>(
101110
Arrays.asList(

‎x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/plugin/gpu/GPUIndexIT.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
2828
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
2929
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
30+
import static org.hamcrest.Matchers.containsString;
3031

3132
@LuceneTestCase.SuppressCodecs("*") // use our custom codec
3233
public class GPUIndexIT extends ESIntegTestCase {
@@ -160,6 +161,44 @@ public void testSearchWithoutGPU() {
160161
assertSearch(indexName, randomFloatVector(dims), numDocs);
161162
}
162163

164+
public void testInt8HnswMaxInnerProductProductFails() {
165+
String indexName = "index_int8_max_inner_product_fails";
166+
final int dims = randomIntBetween(4, 128);
167+
168+
Settings.Builder settingsBuilder = Settings.builder().put(indexSettings());
169+
settingsBuilder.put("index.number_of_shards", 1);
170+
settingsBuilder.put("index.vectors.indexing.use_gpu", true);
171+
172+
String mapping = String.format(Locale.ROOT, """
173+
{
174+
"properties": {
175+
"my_vector": {
176+
"type": "dense_vector",
177+
"dims": %d,
178+
"similarity": "max_inner_product",
179+
"index_options": {
180+
"type": "int8_hnsw"
181+
}
182+
}
183+
}
184+
}
185+
""", dims);
186+
187+
// Index creation should succeed
188+
assertAcked(prepareCreate(indexName).setSettings(settingsBuilder.build()).setMapping(mapping));
189+
ensureGreen();
190+
191+
// Attempt to index a document and expect it to fail
192+
IllegalArgumentException ex = expectThrows(
193+
IllegalArgumentException.class,
194+
() -> client().prepareIndex(indexName).setId("1").setSource("my_vector", randomFloatVector(dims)).get()
195+
);
196+
assertThat(
197+
ex.getMessage(),
198+
containsString("GPU vector indexing does not support [max_inner_product] similarity for [int8_hnsw] index type.")
199+
);
200+
}
201+
163202
private void createIndex(String indexName, int dims, boolean sorted) {
164203
var settings = Settings.builder().put(indexSettings());
165204
settings.put("index.number_of_shards", 1);

‎x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/plugin/gpu/GPUPluginInitializationIT.java‎

Lines changed: 55 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ public void testFFOff() {
109109
GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class);
110110
VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider();
111111

112-
var format = vectorsFormatProvider.getKnnVectorsFormat(null, null);
112+
var format = vectorsFormatProvider.getKnnVectorsFormat(null, null, null);
113113
assertNull(format);
114114
}
115115

@@ -136,7 +136,11 @@ public void testFFOffGPUFormatNull() {
136136
IndexSettings settings = getIndexSettings();
137137
final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions();
138138

139-
var format = vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions);
139+
var format = vectorsFormatProvider.getKnnVectorsFormat(
140+
settings,
141+
indexOptions,
142+
DenseVectorFieldTypeTests.randomGPUSupportedSimilarity(indexOptions.getType())
143+
);
140144
assertNull(format);
141145
}
142146

@@ -151,7 +155,11 @@ public void testIndexSettingOnIndexTypeSupportedGPUSupported() {
151155
IndexSettings settings = getIndexSettings();
152156
final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions();
153157

154-
var format = vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions);
158+
var format = vectorsFormatProvider.getKnnVectorsFormat(
159+
settings,
160+
indexOptions,
161+
DenseVectorFieldTypeTests.randomGPUSupportedSimilarity(indexOptions.getType())
162+
);
155163
assertNotNull(format);
156164
}
157165

@@ -166,7 +174,14 @@ public void testIndexSettingOnIndexTypeNotSupportedThrows() {
166174
IndexSettings settings = getIndexSettings();
167175
final var indexOptions = DenseVectorFieldTypeTests.randomFlatIndexOptions();
168176

169-
var ex = expectThrows(IllegalArgumentException.class, () -> vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions));
177+
var ex = expectThrows(
178+
IllegalArgumentException.class,
179+
() -> vectorsFormatProvider.getKnnVectorsFormat(
180+
settings,
181+
indexOptions,
182+
DenseVectorFieldTypeTests.randomGPUSupportedSimilarity(indexOptions.getType())
183+
)
184+
);
170185
assertThat(ex.getMessage(), startsWith("[index.vectors.indexing.use_gpu] doesn't support [index_options.type] of"));
171186
}
172187

@@ -181,7 +196,14 @@ public void testIndexSettingOnGPUNotSupportedThrows() {
181196
IndexSettings settings = getIndexSettings();
182197
final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions();
183198

184-
var ex = expectThrows(IllegalArgumentException.class, () -> vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions));
199+
var ex = expectThrows(
200+
IllegalArgumentException.class,
201+
() -> vectorsFormatProvider.getKnnVectorsFormat(
202+
settings,
203+
indexOptions,
204+
DenseVectorFieldTypeTests.randomGPUSupportedSimilarity(indexOptions.getType())
205+
)
206+
);
185207
assertThat(
186208
ex.getMessage(),
187209
equalTo("[index.vectors.indexing.use_gpu] was set to [true], but GPU resources are not accessible on the node.")
@@ -200,7 +222,14 @@ public void testIndexSettingOnGPUSupportThrowsRethrows() {
200222
IndexSettings settings = getIndexSettings();
201223
final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions();
202224

203-
var ex = expectThrows(IllegalArgumentException.class, () -> vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions));
225+
var ex = expectThrows(
226+
IllegalArgumentException.class,
227+
() -> vectorsFormatProvider.getKnnVectorsFormat(
228+
settings,
229+
indexOptions,
230+
DenseVectorFieldTypeTests.randomGPUSupportedSimilarity(indexOptions.getType())
231+
)
232+
);
204233
assertThat(
205234
ex.getMessage(),
206235
equalTo("[index.vectors.indexing.use_gpu] was set to [true], but GPU resources are not accessible on the node.")
@@ -218,7 +247,11 @@ public void testIndexSettingAutoIndexTypeSupportedGPUSupported() {
218247
IndexSettings settings = getIndexSettings();
219248
final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions();
220249

221-
var format = vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions);
250+
var format = vectorsFormatProvider.getKnnVectorsFormat(
251+
settings,
252+
indexOptions,
253+
DenseVectorFieldTypeTests.randomGPUSupportedSimilarity(indexOptions.getType())
254+
);
222255
assertNotNull(format);
223256
}
224257

@@ -233,7 +266,11 @@ public void testIndexSettingAutoGPUNotSupported() {
233266
IndexSettings settings = getIndexSettings();
234267
final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions();
235268

236-
var format = vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions);
269+
var format = vectorsFormatProvider.getKnnVectorsFormat(
270+
settings,
271+
indexOptions,
272+
DenseVectorFieldTypeTests.randomGPUSupportedSimilarity(indexOptions.getType())
273+
);
237274
assertNull(format);
238275
}
239276

@@ -248,7 +285,11 @@ public void testIndexSettingAutoIndexTypeNotSupported() {
248285
IndexSettings settings = getIndexSettings();
249286
final var indexOptions = DenseVectorFieldTypeTests.randomFlatIndexOptions();
250287

251-
var format = vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions);
288+
var format = vectorsFormatProvider.getKnnVectorsFormat(
289+
settings,
290+
indexOptions,
291+
DenseVectorFieldTypeTests.randomGPUSupportedSimilarity(indexOptions.getType())
292+
);
252293
assertNull(format);
253294
}
254295

@@ -263,7 +304,11 @@ public void testIndexSettingOff() {
263304
IndexSettings settings = getIndexSettings();
264305
final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions();
265306

266-
var format = vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions);
307+
var format = vectorsFormatProvider.getKnnVectorsFormat(
308+
settings,
309+
indexOptions,
310+
DenseVectorFieldTypeTests.randomGPUSupportedSimilarity(indexOptions.getType())
311+
);
267312
assertNull(format);
268313
}
269314

‎x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/GPUPlugin.java‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ public List<Setting<?>> getSettings() {
6060

6161
@Override
6262
public VectorsFormatProvider getVectorsFormatProvider() {
63-
return (indexSettings, indexOptions) -> {
63+
return (indexSettings, indexOptions, similarity) -> {
6464
if (GPU_FORMAT.isEnabled()) {
6565
GpuMode gpuMode = indexSettings.getValue(VECTORS_INDEXING_USE_GPU_SETTING);
6666
if (gpuMode == GpuMode.TRUE) {
@@ -74,10 +74,10 @@ public VectorsFormatProvider getVectorsFormatProvider() {
7474
"[index.vectors.indexing.use_gpu] was set to [true], but GPU resources are not accessible on the node."
7575
);
7676
}
77-
return getVectorsFormat(indexOptions);
77+
return getVectorsFormat(indexOptions, similarity);
7878
}
7979
if (gpuMode == GpuMode.AUTO && vectorIndexTypeSupported(indexOptions.getType()) && GPUSupport.isSupported(false)) {
80-
return getVectorsFormat(indexOptions);
80+
return getVectorsFormat(indexOptions, similarity);
8181
}
8282
}
8383
return null;
@@ -88,7 +88,10 @@ private boolean vectorIndexTypeSupported(DenseVectorFieldMapper.VectorIndexType
8888
return type == DenseVectorFieldMapper.VectorIndexType.HNSW || type == DenseVectorFieldMapper.VectorIndexType.INT8_HNSW;
8989
}
9090

91-
private static KnnVectorsFormat getVectorsFormat(DenseVectorFieldMapper.DenseVectorIndexOptions indexOptions) {
91+
private static KnnVectorsFormat getVectorsFormat(
92+
DenseVectorFieldMapper.DenseVectorIndexOptions indexOptions,
93+
DenseVectorFieldMapper.VectorSimilarity similarity
94+
) {
9295
if (indexOptions.getType() == DenseVectorFieldMapper.VectorIndexType.HNSW) {
9396
DenseVectorFieldMapper.HnswIndexOptions hnswIndexOptions = (DenseVectorFieldMapper.HnswIndexOptions) indexOptions;
9497
int efConstruction = hnswIndexOptions.efConstruction();
@@ -97,6 +100,17 @@ private static KnnVectorsFormat getVectorsFormat(DenseVectorFieldMapper.DenseVec
97100
}
98101
return new ES92GpuHnswVectorsFormat(hnswIndexOptions.m(), efConstruction);
99102
} else if (indexOptions.getType() == DenseVectorFieldMapper.VectorIndexType.INT8_HNSW) {
103+
if (similarity == DenseVectorFieldMapper.VectorSimilarity.MAX_INNER_PRODUCT) {
104+
throw new IllegalArgumentException(
105+
"GPU vector indexing does not support ["
106+
+ similarity
107+
+ "] similarity for [int8_hnsw] index type. "
108+
+ "Instead, consider using ["
109+
+ DenseVectorFieldMapper.VectorSimilarity.COSINE
110+
+ "] or "
111+
+ " [hnsw] index type."
112+
);
113+
}
100114
DenseVectorFieldMapper.Int8HnswIndexOptions int8HnswIndexOptions = (DenseVectorFieldMapper.Int8HnswIndexOptions) indexOptions;
101115
int efConstruction = int8HnswIndexOptions.efConstruction();
102116
if (efConstruction == HnswGraphBuilder.DEFAULT_BEAM_WIDTH) {

‎x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsWriter.java‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -313,9 +313,18 @@ private CagraIndex buildGPUIndex(
313313
CuVSMatrix dataset
314314
) throws Throwable {
315315
CagraIndexParams.CuvsDistanceType distanceType = switch (similarityFunction) {
316-
case EUCLIDEAN -> CagraIndexParams.CuvsDistanceType.L2Expanded;
317-
case DOT_PRODUCT, MAXIMUM_INNER_PRODUCT -> CagraIndexParams.CuvsDistanceType.InnerProduct;
318316
case COSINE -> CagraIndexParams.CuvsDistanceType.CosineExpanded;
317+
case EUCLIDEAN -> CagraIndexParams.CuvsDistanceType.L2Expanded;
318+
case DOT_PRODUCT -> {
319+
if (dataType == CuVSMatrix.DataType.BYTE) {
320+
yield CagraIndexParams.CuvsDistanceType.CosineExpanded;
321+
}
322+
yield CagraIndexParams.CuvsDistanceType.InnerProduct;
323+
}
324+
case MAXIMUM_INNER_PRODUCT -> {
325+
assert dataType != CuVSMatrix.DataType.BYTE;
326+
yield CagraIndexParams.CuvsDistanceType.InnerProduct;
327+
}
319328
};
320329

321330
// TODO: expose cagra index params for algorithm, NNDescentNumIterations

‎x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormatTests.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,9 @@
1313
import org.apache.lucene.tests.util.LuceneTestCase;
1414
import org.apache.lucene.tests.util.TestUtil;
1515
import org.elasticsearch.common.logging.LogConfigurator;
16+
import org.elasticsearch.index.IndexVersion;
17+
import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
18+
import org.elasticsearch.index.mapper.vectors.DenseVectorFieldTypeTests;
1619
import org.elasticsearch.xpack.gpu.GPUSupport;
1720
import org.junit.BeforeClass;
1821

@@ -39,7 +42,8 @@ protected Codec getCodec() {
3942

4043
@Override
4144
protected VectorSimilarityFunction randomSimilarity() {
42-
return VectorSimilarityFunction.values()[random().nextInt(VectorSimilarityFunction.values().length)];
45+
return DenseVectorFieldTypeTests.randomGPUSupportedSimilarity(DenseVectorFieldMapper.VectorIndexType.INT8_HNSW)
46+
.vectorSimilarityFunction(IndexVersion.current(), DenseVectorFieldMapper.ElementType.FLOAT);
4347
}
4448

4549
@Override

0 commit comments

Comments
 (0)