-
Notifications
You must be signed in to change notification settings - Fork 99
feat: support MRLE with split RaBitQ #2560
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 3 commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -27,16 +27,11 @@ Quantizer)**. | |
|
|
||
| ## Quick start | ||
|
|
||
| `tq` is currently exposed as a **public, externally configurable** quantization type only by | ||
| **HGraph**. HGraph maps the top-level keys `tq_chain` and `rabitq_pca_dim` into the nested | ||
| `base_codes.quantization_params` JSON via its external-parameter mapping | ||
| (`src/algorithm/hgraph.cpp:370-385`). IVF, BruteForce, Pyramid and WARP all internally render | ||
| a `tq_chain` field into their inner JSON template, but none of them expose `tq_chain` (or any | ||
| other TQ parameter) in their external mapping today. `CheckAndMappingExternalParam` rejects | ||
| unknown external keys with `invalid config param` | ||
| (`src/utils/util_functions.cpp:50-53`), so passing `tq_chain` in the `index_param` JSON of | ||
| those indexes will fail at index construction. Configuring TQ on non-HGraph indexes | ||
| therefore requires code-side changes to add the external mapping. | ||
| `tq` is exposed as a public, externally configurable quantization type by **HGraph** and | ||
| **Pyramid**. Both map `tq_chain` and dimension parameters into | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P3] Document the new parameters on the Pyramid pages This advertises Pyramid as supporting external TQ and MRLE split configuration, but the canonical English and Chinese Pyramid pages still omit
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P3] Document the new Pyramid TQ parameters on Pyramid pages This now advertises Pyramid as publicly supporting TQ and MRLE split, but the canonical English and Chinese Pyramid pages still omit |
||
| `base_codes.quantization_params`. The MRLE + RaBitQ x+y split combination is available in both | ||
| indexes and automatically reorders from the base split datacell. IVF, BruteForce and WARP do | ||
| not currently expose `tq_chain` through their external parameter mapping. | ||
|
|
||
| ```cpp | ||
| std::string params = R"({ | ||
|
|
@@ -85,6 +80,7 @@ Examples: | |
| | `"pca, rom, sq8_uniform"` | PCA reduction, random rotation, then 8-bit uniform — the example chain. | | ||
| | `"pca, rom, rabitq"` | PCA + rotation feeding the RaBitQ binary quantizer. | | ||
| | `"mrle, fp32"` | MRLE projection then store as fp32 (MRLE must be first). | | ||
| | `"mrle, rabitq"` | MRLE reduction followed by RaBitQ; with x+y split storage, filter and supplement codes are produced by the terminal RaBitQ. | | ||
|
|
||
| Constraints (`transform_quantizer_parameter.cpp:33-45`): | ||
|
|
||
|
|
@@ -98,6 +94,9 @@ Constraints (`transform_quantizer_parameter.cpp:33-45`): | |
| when `is_transform_quantizer` is true (`src/datacell/flatten_interface.cpp:166`), so using | ||
| any of those three as the terminal quantizer fails at index construction with an | ||
| "unsupported quantization type" error. | ||
| - RaBitQ x+y split storage supports TQ only for the exact chain `"mrle, rabitq"`. The split | ||
| datacell reuses the terminal RaBitQ encoder and keeps RaBitQ internal FHT/ROM rotation | ||
| behavior unchanged. Other transformed split chains are rejected. | ||
| - Any unrecognized transformer name raises `INVALID_ARGUMENT: invalid transformer name` | ||
| (`transform_quantizer.h:225-227`). | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,167 @@ | ||
| # Issue #2533:MRLE + RaBitQ Split 实施计划 | ||
|
|
||
| ## 1. 目标 | ||
|
|
||
| 在 HGraph 与 Pyramid 中支持以下组合: | ||
|
|
||
| ```json | ||
| { | ||
| "base_quantization_type": "tq", | ||
| "tq_chain": "mrle, rabitq", | ||
| "mrle_dim": 768, | ||
| "precise_quantization_type": "rabitq", | ||
| "use_reorder": true, | ||
| "rabitq_bits_per_dim_base": 3, | ||
| "rabitq_bits_per_dim_precise": 5 | ||
| } | ||
| ``` | ||
|
|
||
| 数据路径保持严格分层: | ||
|
|
||
| ```text | ||
| 原始向量 | ||
| -> TransformQuantizer 执行 MRLE 截断 | ||
| -> 末端 RaBitQuantizer 执行自身的 FHT/ROM 随机旋转与 RaBitQ 编码 | ||
| -> RaBitQSplitDataCell 将 RaBitQ code 拆成 x-bit filter 与 y-bit supplement | ||
| ``` | ||
|
|
||
| MRLE 与 split storage 解耦。MRLE 只改变送入末端量化器的向量;split 编码、lower | ||
| bound、full distance、scalar-code optimized build 仍全部由 `RaBitQuantizer` 实现。 | ||
|
|
||
| ## 2. 架构决策 | ||
|
|
||
| ### 2.1 不新增 bridge 或 codec | ||
|
|
||
| `TransformQuantizer` 已能组合任意末端量化器,因此不引入 | ||
| `RaBitQSplitBridge`、`TransformFlattenDataCell` 或另一套 transform chain。 | ||
|
|
||
| 采用 `BottomQuantizerAccessor` 静态策略: | ||
|
|
||
| ```text | ||
| RaBitQuantizer | ||
| -> bottom quantizer/computer 是自身 | ||
|
|
||
| TransformQuantizer<RaBitQuantizer> | ||
| -> bottom quantizer 是 inner RaBitQuantizer | ||
| -> bottom computer 是 inner Computer<RaBitQuantizer> | ||
| -> optimized build 输入先经过 TransformBaseVector | ||
| ``` | ||
|
|
||
| 该策略没有热路径虚调用,也不复制 RaBitQ 编码逻辑。 | ||
|
|
||
| ### 2.2 支持范围 | ||
|
|
||
| split + TQ 首版只接受精确链: | ||
|
|
||
| ```text | ||
| mrle, rabitq | ||
| ``` | ||
|
|
||
| 不接受外部 PCA、ROM 或 FHT。RaBitQ 内部已有的 FHT/ROM 随机旋转保持不变,顺序为 | ||
| MRLE 降维后再执行 RaBitQ 内部旋转。 | ||
|
|
||
| ### 2.3 code layout | ||
|
|
||
| `TransformQuantizer` 的 terminal quantizer code 位于 full code 起始位置,transform metadata | ||
| 位于其后。MRLE 当前 metadata size 为 0,因此 split datacell 合并 full code 时将外层 buffer | ||
| 清零,再让底层 RaBitQ 写入 code 起始区域即可。 | ||
|
|
||
| ## 3. 代码修改 | ||
|
|
||
| ### 3.1 BottomQuantizerAccessor | ||
|
|
||
| 新增: | ||
|
|
||
| - `src/quantization/bottom_quantizer_accessor.h` | ||
|
|
||
| 职责: | ||
|
|
||
| - direct quantizer 返回自身 quantizer/computer; | ||
| - `TransformQuantizer` 返回 inner quantizer/computer; | ||
| - optimized scalar build 时为 TQ 准备变换后的 base input。 | ||
|
|
||
| ### 3.2 TransformQuantizer | ||
|
|
||
| 修改: | ||
|
|
||
| - `src/quantization/transform_quantization/transform_quantizer.h` | ||
| - `src/quantization/transform_quantization/transform_quantizer_parameter.cpp` | ||
|
|
||
| 内容: | ||
|
|
||
| - 新增 `GetTransformedDim()` 与 `TransformBaseVector()`; | ||
| - 训练 batch 按 transformed dim 紧凑分配和寻址,修复降维后第二条向量起 stride 错误; | ||
| - 无 metadata 输出时允许 transform chain 不编码 meta; | ||
| - compatibility 同时比较 transformer 参数与末端 quantizer 完整参数。 | ||
|
|
||
| ### 3.3 RaBitQSplitDataCell | ||
|
|
||
| 修改: | ||
|
|
||
| - `src/datacell/rabitq_split_datacell.h` | ||
| - `src/datacell/rabitq_split_datacell_factory*.{h,cpp}` | ||
| - `src/datacell/flatten_interface.cpp` | ||
|
|
||
| 内容: | ||
|
|
||
| - datacell 增加 `QuantizerT` 模板参数,默认仍为 direct `RaBitQuantizer`; | ||
| - 用 `BottomQuantizerAccessor` 获取底层 RaBitQ 与 query computer; | ||
| - split、merge、lower bound、full distance、scalar code 均调用底层 RaBitQ; | ||
| - 普通 Train/Encode/FactoryComputer/Serialize 仍调用外层 quantizer; | ||
| - factory 静态实例化 direct 与 `TransformQuantizer<RaBitQuantizer>` 两条路径; | ||
| - factory 校验 TQ split 只能使用 `mrle,rabitq`。 | ||
|
|
||
| 现有 direct split 类型、序列化布局和行为保持不变。 | ||
|
|
||
| ### 3.4 共享参数映射 | ||
|
|
||
| 修改: | ||
|
|
||
| - `src/algorithm/inner_index_parameter.h/.cpp` | ||
| - `src/algorithm/hgraph/hgraph_param_mapping.cpp` | ||
| - `src/algorithm/pyramid/pyramid.cpp` | ||
| - `src/algorithm/pyramid/pyramid_zparameters.cpp` | ||
|
|
||
| 新增共享内部函数: | ||
|
|
||
| - `MapRaBitQSplitParam`:校验 exact chain、x/y bits、terminal RaBitQ,并生成 split 内部参数; | ||
| - `ValidateMRLEDim`:统一校验 `mrle_dim` 范围。 | ||
|
|
||
| HGraph 与 Pyramid 都将 split reorder source 固定为 `base`。 | ||
|
|
||
| ### 3.5 Pyramid base reorder | ||
|
|
||
| Pyramid 增加 `reorder_by_base_` 与两个小型 helper: | ||
|
|
||
| - `has_precise_codes()`:仅 precise reorder 时返回 true; | ||
| - `graph_codes()`:direct split/TQ split 返回 base,传统 reorder 返回 precise。 | ||
|
|
||
| base reorder 不创建、不训练、不插入、不序列化 precise datacell;`FlattenReorder` 直接包装 | ||
| base split datacell,从 supplement code 完成最终重排。 | ||
|
|
||
| ## 4. 测试 | ||
|
|
||
| 单元测试覆盖: | ||
|
|
||
| - MRLE + split factory 创建; | ||
| - 多向量降维训练的紧凑 batch 布局; | ||
| - optimized scalar-code build 的 base transform; | ||
| - split query 与 merged full-code distance 一致; | ||
| - 非 `mrle,rabitq` chain 被拒绝; | ||
| - HGraph/Pyramid 映射为 `rabitq_split` 与 `reorder_source=base`; | ||
| - TQ compatibility 比较 MRLE 与 RaBitQ 参数; | ||
| - direct split 回归。 | ||
|
|
||
| 功能测试覆盖: | ||
|
|
||
| - HGraph MRLE + RaBitQ split 构建与 KNN 搜索; | ||
| - Pyramid MRLE + RaBitQ split 构建与 KNN 搜索。 | ||
|
|
||
| ## 5. 验收标准 | ||
|
|
||
| - `clang-format-15` 通过; | ||
| - `git diff --check` 通过; | ||
| - debug library build 通过; | ||
| - `[MRLE]` 单元和功能测试通过; | ||
| - direct split、optimized build、TQ compute 与既有参数映射回归通过; | ||
| - 英文和中文量化文档同步说明 exact chain 与 RaBitQ 内部旋转顺序。 |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -17,78 +17,9 @@ | |
| #include "common.h" | ||
| #include "hgraph.h" // IWYU pragma: keep | ||
| #include "hgraph_parameter.h" | ||
| #include "quantization/rabitq_quantization/rabitq_quantizer_parameter.h" | ||
|
|
||
| namespace vsag { | ||
|
|
||
| namespace { | ||
|
|
||
| void | ||
| map_rabitq_split_param(const JsonType& external_json, JsonType& inner_json) { | ||
| if (not external_json.Contains(RABITQ_BITS_PER_DIM_PRECISE)) { | ||
| return; | ||
| } | ||
|
|
||
| CHECK_ARGUMENT( | ||
| external_json.Contains(RABITQ_BITS_PER_DIM_BASE), | ||
| fmt::format("{} requires {}", RABITQ_BITS_PER_DIM_PRECISE, RABITQ_BITS_PER_DIM_BASE)); | ||
| CHECK_ARGUMENT( | ||
| external_json.Contains(HGRAPH_BASE_QUANTIZATION_TYPE), | ||
| fmt::format("{} requires {}", RABITQ_BITS_PER_DIM_PRECISE, HGRAPH_BASE_QUANTIZATION_TYPE)); | ||
| CHECK_ARGUMENT( | ||
| external_json.Contains(HGRAPH_PRECISE_QUANTIZATION_TYPE), | ||
| fmt::format( | ||
| "{} requires {}", RABITQ_BITS_PER_DIM_PRECISE, HGRAPH_PRECISE_QUANTIZATION_TYPE)); | ||
|
|
||
| const auto base_quantization_type = external_json[HGRAPH_BASE_QUANTIZATION_TYPE].GetString(); | ||
| const auto precise_quantization_type = | ||
| external_json[HGRAPH_PRECISE_QUANTIZATION_TYPE].GetString(); | ||
| CHECK_ARGUMENT(base_quantization_type == QUANTIZATION_TYPE_VALUE_RABITQ, | ||
| fmt::format("{} requires {}={}", | ||
| RABITQ_BITS_PER_DIM_PRECISE, | ||
| HGRAPH_BASE_QUANTIZATION_TYPE, | ||
| QUANTIZATION_TYPE_VALUE_RABITQ)); | ||
| CHECK_ARGUMENT(precise_quantization_type == QUANTIZATION_TYPE_VALUE_RABITQ, | ||
| fmt::format("{} requires {}={}", | ||
| RABITQ_BITS_PER_DIM_PRECISE, | ||
| HGRAPH_PRECISE_QUANTIZATION_TYPE, | ||
| QUANTIZATION_TYPE_VALUE_RABITQ)); | ||
|
|
||
| const int64_t filter_bits = external_json[RABITQ_BITS_PER_DIM_BASE].GetInt(); | ||
| const int64_t supplement_bits = external_json[RABITQ_BITS_PER_DIM_PRECISE].GetInt(); | ||
| CHECK_ARGUMENT( | ||
| filter_bits >= 1, | ||
| fmt::format("{} must be in [1, 8], got {}", RABITQ_BITS_PER_DIM_BASE, filter_bits)); | ||
| CHECK_ARGUMENT( | ||
| filter_bits <= 8, | ||
| fmt::format("{} must be in [1, 8], got {}", RABITQ_BITS_PER_DIM_BASE, filter_bits)); | ||
| CHECK_ARGUMENT( | ||
| supplement_bits >= 1, | ||
| fmt::format("{} must be in [1, 8], got {}", RABITQ_BITS_PER_DIM_PRECISE, supplement_bits)); | ||
| CHECK_ARGUMENT( | ||
| supplement_bits <= 8, | ||
| fmt::format("{} must be in [1, 8], got {}", RABITQ_BITS_PER_DIM_PRECISE, supplement_bits)); | ||
| const int64_t total_bits = filter_bits + supplement_bits; | ||
| CHECK_ARGUMENT(total_bits <= 8, | ||
| fmt::format("{} + {} must be no greater than 8, got {}", | ||
| RABITQ_BITS_PER_DIM_BASE, | ||
| RABITQ_BITS_PER_DIM_PRECISE, | ||
| total_bits)); | ||
|
|
||
| inner_json[REORDER_SOURCE_KEY].SetString(HGRAPH_REORDER_SOURCE_BASE); | ||
| inner_json[BASE_CODES_KEY][CODES_TYPE_KEY].SetString(RABITQ_SPLIT_CODES); | ||
| inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_VERSION_KEY].SetString( | ||
| RaBitQuantizerParameter::RABITQ_VERSION_SPLIT); | ||
| inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_BITS_PER_DIM_QUERY_KEY] | ||
| .SetInt(32); | ||
| inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_BITS_PER_DIM_FILTER_KEY] | ||
| .SetInt(filter_bits); | ||
| inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][RABITQ_QUANTIZATION_BITS_PER_DIM_BASE_KEY] | ||
| .SetInt(total_bits); | ||
| } | ||
|
|
||
| } // namespace | ||
|
|
||
| JsonType | ||
| HGraph::map_hgraph_param(const JsonType& hgraph_json) { | ||
| static const ConstParamMap external_mapping = { | ||
|
|
@@ -613,7 +544,7 @@ HGraph::map_hgraph_param(const JsonType& hgraph_json) { | |
| std::string str = format_map(hgraph_params_template, DEFAULT_MAP); | ||
| auto inner_json = JsonType::Parse(str); | ||
| mapping_external_param_to_inner(hgraph_json, external_mapping, inner_json); | ||
| map_rabitq_split_param(hgraph_json, inner_json); | ||
| MapRaBitQSplitParam(hgraph_json, inner_json); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] HGraph MRLE split does not set |
||
|
|
||
| return inner_json; | ||
| } | ||
|
|
@@ -628,16 +559,7 @@ HGraph::CheckAndMappingExternalParam(const JsonType& external_param, | |
| inner_json[RAW_VECTOR_KEY][CODES_TYPE_KEY].SetString(SPARSE_CODES); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [suggestion] HGraph MRLE split does not set At HGraph with MRLE split cannot decode vectors to fp32 for statistics ( The |
||
| } | ||
|
|
||
| if (external_param.Contains(INDEX_MRLE_DIM)) { | ||
| CHECK_ARGUMENT(external_param[INDEX_MRLE_DIM].IsNumberInteger(), | ||
| fmt::format("mrle_dim must be an integer, got {}", | ||
| external_param[INDEX_MRLE_DIM].Dump())); | ||
| int64_t mrle_dim = external_param[INDEX_MRLE_DIM].GetInt(); | ||
| bool valid_mrle_dim = mrle_dim >= 0 and mrle_dim <= static_cast<int64_t>(common_param.dim_); | ||
| CHECK_ARGUMENT( | ||
| valid_mrle_dim, | ||
| fmt::format("mrle_dim({}) must be in range [0, {}]", mrle_dim, common_param.dim_)); | ||
| } | ||
| ValidateMRLEDim(external_param, common_param.dim_); | ||
|
|
||
| auto hgraph_parameter = std::make_shared<HGraphParameter>(); | ||
| hgraph_parameter->data_type = common_param.data_type_; | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.