[diskann-garnet] Fix Q8 quantizer state loading - #1264
Conversation
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.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1264 +/- ##
==========================================
+ Coverage 90.16% 90.44% +0.27%
==========================================
Files 509 498 -11
Lines 97238 95832 -1406
==========================================
- Hits 87678 86676 -1002
+ Misses 9560 9156 -404
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
hildebrandmw
left a comment
There was a problem hiding this comment.
I looked over the diskann-quantization code. One small suggestion for testing, otherwise looks fine.
f38b782 to
096f80e
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 10 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
diskann-quantization/src/minmax/quantizer.rs:821
- The serialization roundtrip test only checks
PartialEqonMinMaxQuantizer, so it can pass even if the deserialized quantizer produces different codes (e.g., ifTransform’sPartialEqmisses some internal state). Since the goal here is to ensure persisted state restores the exact quantization behavior, add a check thatcompress_intoproduces identical canonical bytes for a fixed input vector (and this will also make the NBITS-specific test variants meaningfully exercise per-bitrate codegen).
let serialized = quantizer.serialize(GlobalAllocator).unwrap();
let deserialized = MinMaxQuantizer::try_deserialize(&serialized).unwrap();
assert_eq!(deserialized, quantizer);
}
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.