[diskann-garnet] Fix Q8 quantizer state loading#1264
Open
metajack wants to merge 1 commit into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes restart correctness for the Q8 (MinMax 8-bit) quantizer in diskann-garnet by persisting and restoring the quantizer’s randomized DoubleHadamard transform state via FlatBuffers, instead of assuming reconstruction yields an identical transform.
Changes:
- Added FlatBuffers serialization/deserialization support for
MinMaxQuantizer(including transform state) indiskann-quantization, plus round-trip tests. - Updated
diskann-garnetQ8 initialization to load persisted MinMax quantizer state when present and persist it when first created; added a restart test for Q8. - Added a MinMax FlatBuffers schema and build integration; bumped
diskann-garnetpackage versions to4.0.2.
Reviewed changes
Copilot reviewed 9 out of 10 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| diskann-quantization/src/minmax/quantizer.rs | Adds MinMax quantizer FlatBuffers serialize/deserialize + tests to ensure transform state round-trips. |
| diskann-quantization/src/flatbuffers/minmax/quantizer_generated.rs | New generated FlatBuffers bindings for MinMax quantizer state. |
| diskann-quantization/src/flatbuffers.rs | Wires MinMax FlatBuffers module into the crate import tree. |
| diskann-quantization/schemas/minmax_v001.fbs | Introduces the MinMax quantizer schema with identifier mq00. |
| diskann-quantization/build.rs | Includes MinMax schema in FlatBuffers build generation. |
| diskann-garnet/src/quantization.rs | Uses MinMaxQuantizer FlatBuffers serialization and adds construction-from-bytes for restart. |
| diskann-garnet/src/provider.rs | Persists/loads Q8 quantizer state in metadata; adds a Q8 restart test. |
| diskann-garnet/diskann-garnet.nuspec | Version bump to 4.0.2. |
| diskann-garnet/Cargo.toml | Version bump to 4.0.2. |
| Cargo.lock | Updates locked version for diskann-garnet. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+215
to
+233
| let quantizer = if let Some(quant_state) = | ||
| callbacks.read_varsize_iid::<u8>(&context.term(Term::Metadata), QUANT_STATE_KEY) | ||
| { | ||
| quantization::MinMax8Bit::new_from_bytes(metric_type, &quant_state)? | ||
| } else { | ||
| let quantizer = quantization::MinMax8Bit::new(dim, metric_type)?; | ||
|
|
||
| if !callbacks.write_iid( | ||
| &context.term(Term::Metadata), | ||
| QUANT_STATE_KEY, | ||
| &quantizer.serialize()?, | ||
| ) { | ||
| return Err(GarnetError::Write.into()); | ||
| } | ||
|
|
||
| quantizer | ||
| }; | ||
|
|
||
| let quantizer = Box::new(quantizer) as Box<dyn GarnetQuantizer>; |
Comment on lines
+2166
to
+2178
| assert!(provider.callbacks.read_single_iid( | ||
| &ctx.term(Term::Vector), | ||
| 1, | ||
| bytemuck::cast_slice_mut::<f32, u8>(&mut fv), | ||
| )); | ||
|
|
||
| quantizer.compress(&fv, &mut qv).unwrap(); | ||
|
|
||
| assert!( | ||
| provider | ||
| .callbacks | ||
| .read_single_iid(&ctx.term(Term::Quantized), 1, &mut orig_qv) | ||
| ); |
Codecov Report❌ Patch coverage is ❌ Your patch status has failed because the patch coverage (82.00%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #1264 +/- ##
==========================================
- Coverage 90.16% 90.14% -0.03%
==========================================
Files 509 510 +1
Lines 97238 97535 +297
==========================================
+ Hits 87678 87921 +243
- Misses 9560 9614 +54
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The previous work on saving/restoring quantizer state in diskann-garnet did not save/restore state for the Q8 quantizer assuming that constructing it was sufficient. However, the DoubleHadamard transform generates a random transform on construction, so a fresh quantizer will not work to decode quant vectors.
This PR saves and restores the Q8 quantizer state, which is essentially just the DoubleHadamard matrix.
It adds similar serialization/deserialization to MinMaxQuantizer as the spherical quantizer which is then used by diskann-garnet.