Repository navigation
Description
Describe the bug
The binary quantizer's sampling_median threshold is wrong when it is trained on a device dataset. In the device train() overload (cpp/src/preprocessing/quantize/detail/binary.cuh:240-246 on main 2450691), the sampled rows are transposed into transposed_sampled_dataset_chunk and each row of that buffer is sorted, but raft::copy_matrix then reads the medians from sampled_dataset_chunk, the unsorted and untransposed buffer:
raft::copy_matrix(threshold_ptr + dim_offset,
size_t(1),
sampled_dataset_chunk.data_handle() + (num_samples - 1) / 2, // should be transposed_sampled_dataset_chunk
num_samples, ...So the threshold for dimension j is an arbitrary sampled value, usually from a different column. Nothing fails, the thresholds are just wrong. The host train() overload takes the median from the sorted column and is correct. mean (the default) and the zero threshold are not affected. The existing test compares GPU and CPU transforms using the same quantizer, so it does not check the thresholds.
Steps/Code to reproduce bug
Through the C API: 10001 rows and 12 dimensions, where column j is 1000*j + 0.05*perm_j for a random permutation perm_j of 0..10000. Train one quantizer with threshold = SAMPLING_MEDIAN and sampling_ratio = 1.0 on a device copy of the data and one on the host copy, transform the same data with each, and count the set bits per dimension. With sampling_ratio = 1.0 the sample is all rows, so the median of each column is its exact median and exactly 5000 of 10001 bits should be set in every dimension.
Output on a GH200 with the libcuvs-cu13 26.08.1 wheel:
dim | bits set: device-trained / host-trained (expect 5000 / 5000)
0 | 0 / 5000
1 | 5632 / 5000
2 | 0 / 5000
3 | 0 / 5000
4 | 451 / 5000
5 | 0 / 5000
6 | 10001 / 5000
7 | 1453 / 5000
8 | 10001 / 5000
9 | 10001 / 5000
10 | 2113 / 5000
11 | 10001 / 5000
All 12 device-trained thresholds are wrong; all host-trained ones are right.
Expected behavior
Device-trained and host-trained sampling_median thresholds agree. The fix is to read from transposed_sampled_dataset_chunk in that copy_matrix call. A regression test could train both ways on data with a different offset per column and compare the thresholds. I can open a PR with the fix and that test.
Environment details (please complete the following information):
- Environment location: Bare-metal, GH200 (aarch64), driver 595.58.03
- Method of cuVS install: pip wheel (
libcuvs-cu1326.08.1), C API
Additional context
The code has read from sampled_dataset_chunk since sampling_median was added in #716.
Metadata
Metadata
Assignees
Labels
Type
Projects
- StatusShow more project fieldsNeeds Triage
Activity