Fix bug where quantization framework does not work with training #2100

ryanbogan · 2024-09-12T06:33:40Z

Description

Currently, the quantization framework is broken and does not work when training. This PR fixes this bug, so the vectors are properly quantized when the framework is enabled

Check List

New functionality includes testing.
New functionality has been documented.
API changes companion pull request created.
Commits are signed per the DCO using --signoff.
Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

ryanbogan · 2024-09-12T16:13:34Z

Rolling upgrade failure not related to this PR

release-notes/opensearch-knn.release-notes-2.17.0.0.md

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

jmazanec15

This looks good to me thank!

shatejas

verify null checks, you can adress other comments in a follow up PR

shatejas · 2024-09-12T17:35:22Z

jni/src/faiss_wrapper.cpp

+            searchParameters = &ivfParams;
+        } else {
+            auto hnswReader = dynamic_cast<const faiss::IndexBinaryHNSW*>(indexReader->index);
+            if(hnswReader != nullptr && (methodParamsJ != nullptr || parentIdsJ != nullptr)) {


nit: do you need the condition after &&

That was already there, I'm not sure if it's necessary or not

shatejas · 2024-09-12T17:36:00Z

jni/src/faiss_wrapper.cpp

-            if (parentIdsJ != nullptr) {
-                idGrouper = buildIDGrouperBitmap(jniUtil, env, parentIdsJ, &idGrouperBitmap);
-                hnswParams.grp = idGrouper.get();
+        if (ivfReader) {


might be a good idea to throw in a unit test in this or as an immediate follow up

Sure, will do in follow-up

src/main/java/org/opensearch/knn/training/FloatTrainingDataConsumer.java

shatejas · 2024-09-12T17:46:16Z

src/main/java/org/opensearch/knn/training/FloatTrainingDataConsumer.java

+                return ArrayUtils.toPrimitive((Float[]) vectors.get(position));
+            }
+        };
+        QuantizationState quantizationState = quantizer.train(trainingRequest);


nit: I am assuming QuantizationServic doesn't fit here?

shatejas · 2024-09-12T17:51:29Z

src/main/java/org/opensearch/knn/training/FloatTrainingDataConsumer.java

+            try {
+                List<byte[]> byteVectors = quantizeVectors(floats);
+                long memoryAddress = trainingDataAllocation.getMemoryAddress();
+                memoryAddress = JNICommons.storeBinaryVectorData(memoryAddress, byteVectors.toArray(new byte[0][0]), byteVectors.size());


Do an argument capture for byte array and make sure they aren't the same in unit test

shatejas · 2024-09-12T17:53:43Z

src/test/java/org/opensearch/knn/training/FloatTrainingDataConsumerTests.java

+        verify(trainingDataAllocation).setMemoryAddress(valueCapture.capture());
+        when(trainingDataAllocation.getMemoryAddress()).thenReturn(valueCapture.getValue());


nit: verify the capture

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

* Initial implementation Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Modify integration test and fix bugs in jni Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Fix unit test Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Fix integration test after merge Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Add changelog (release notes) Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Add unit test Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Remove entry for release notes Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Add null checks Signed-off-by: Ryan Bogan <rbogan@amazon.com> --------- Signed-off-by: Ryan Bogan <rbogan@amazon.com> (cherry picked from commit 5d10d64)

…) (#2102) * Initial implementation Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Modify integration test and fix bugs in jni Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Fix unit test Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Fix integration test after merge Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Add changelog (release notes) Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Add unit test Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Remove entry for release notes Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Add null checks Signed-off-by: Ryan Bogan <rbogan@amazon.com> --------- Signed-off-by: Ryan Bogan <rbogan@amazon.com> (cherry picked from commit 5d10d64) Co-authored-by: Ryan Bogan <rbogan@amazon.com>

…) (#2101) * Initial implementation Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Modify integration test and fix bugs in jni Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Fix unit test Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Fix integration test after merge Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Add changelog (release notes) Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Add unit test Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Remove entry for release notes Signed-off-by: Ryan Bogan <rbogan@amazon.com> * Add null checks Signed-off-by: Ryan Bogan <rbogan@amazon.com> --------- Signed-off-by: Ryan Bogan <rbogan@amazon.com> (cherry picked from commit 5d10d64) Co-authored-by: Ryan Bogan <rbogan@amazon.com>

ryanbogan added 5 commits September 10, 2024 13:24

Initial implementation

31a9646

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

Modify integration test and fix bugs in jni

d18c5e4

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

Fix unit test

670dee2

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

Merge branch 'main' into fix_training_bug

7ef98bf

Fix integration test after merge

9abcbde

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

ryanbogan requested review from heemin32, navneet1v, VijayanB, vamshin, jmazanec15, naveentatikonda, junqiu-lei, martin-gaievski and luyuncheng as code owners September 12, 2024 06:33

ryanbogan added skip-changelog and removed skip-changelog labels Sep 12, 2024

Add changelog (release notes)

0cd7ea3

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

ryanbogan added the skip-changelog label Sep 12, 2024

jmazanec15 reviewed Sep 12, 2024

View reviewed changes

release-notes/opensearch-knn.release-notes-2.17.0.0.md Outdated Show resolved Hide resolved

ryanbogan added 2 commits September 12, 2024 10:22

Add unit test

e0f10d1

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

Remove entry for release notes

09690c2

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

jmazanec15 previously approved these changes Sep 12, 2024

View reviewed changes

shatejas reviewed Sep 12, 2024

View reviewed changes

Add null checks

49806dc

Signed-off-by: Ryan Bogan <rbogan@amazon.com>

ryanbogan dismissed jmazanec15’s stale review via 49806dc September 12, 2024 18:23

ryanbogan requested a review from shatejas September 12, 2024 18:31

ryanbogan changed the title ~~Enable training with quantization framework~~ Fix bug where quantization framework does not work with training Sep 12, 2024

shatejas approved these changes Sep 12, 2024

View reviewed changes

heemin32 approved these changes Sep 12, 2024

View reviewed changes

ryanbogan requested a review from jmazanec15 September 12, 2024 20:17

ryanbogan added backport 2.x backport 2.17 labels Sep 12, 2024

jmazanec15 approved these changes Sep 12, 2024

View reviewed changes

ryanbogan merged commit 5d10d64 into opensearch-project:main Sep 12, 2024
38 of 39 checks passed

opensearch-trigger-bot bot mentioned this pull request Sep 12, 2024

[Backport 2.x] Fix bug where quantization framework does not work with training #2101

Merged

opensearch-trigger-bot bot mentioned this pull request Sep 12, 2024

[Backport 2.17] Fix bug where quantization framework does not work with training #2102

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix bug where quantization framework does not work with training #2100

Fix bug where quantization framework does not work with training #2100

ryanbogan commented Sep 12, 2024 •

edited

Loading

ryanbogan commented Sep 12, 2024

jmazanec15 left a comment

shatejas left a comment

shatejas Sep 12, 2024

ryanbogan Sep 12, 2024

shatejas Sep 12, 2024

ryanbogan Sep 12, 2024

shatejas Sep 12, 2024

shatejas Sep 12, 2024

shatejas Sep 12, 2024 •

edited

Loading

		verify(trainingDataAllocation).setMemoryAddress(valueCapture.capture());
		when(trainingDataAllocation.getMemoryAddress()).thenReturn(valueCapture.getValue());

Fix bug where quantization framework does not work with training #2100

Fix bug where quantization framework does not work with training #2100

Conversation

ryanbogan commented Sep 12, 2024 • edited Loading

Description

Check List

ryanbogan commented Sep 12, 2024

jmazanec15 left a comment

Choose a reason for hiding this comment

shatejas left a comment

Choose a reason for hiding this comment

shatejas Sep 12, 2024

Choose a reason for hiding this comment

ryanbogan Sep 12, 2024

Choose a reason for hiding this comment

shatejas Sep 12, 2024

Choose a reason for hiding this comment

ryanbogan Sep 12, 2024

Choose a reason for hiding this comment

shatejas Sep 12, 2024

Choose a reason for hiding this comment

shatejas Sep 12, 2024

Choose a reason for hiding this comment

shatejas Sep 12, 2024 • edited Loading

Choose a reason for hiding this comment

ryanbogan commented Sep 12, 2024 •

edited

Loading

shatejas Sep 12, 2024 •

edited

Loading