[BUG] cuVS Scaler Quantization memory issue

[singhmanas1]

Describe the bug
Running into std::bad_alloc: out_of_memory: CUDA error while training scalar embeddings with cuVS. Runs below 800,000 embeddings but throws error after.

Embedding model - https://build.nvidia.com/nvidia/llama-3_2-nv-embedqa-1b-v2/modelcard
Dimensions - 2048
Precision - FP16

Steps/Code to reproduce bug

def scalar_quantize_cuvs(embeddings):
    embeddings_gpu = cp.asarray(embeddings)     
    params = scalar.QuantizerParams(quantile=0.99)
    quantizer = scalar.train(params, embeddings_gpu)
    transformed = scalar.transform(quantizer, embeddings_gpu)
    return transformed, quantization_time

Error Statement

Error during quantization: std::bad_alloc: out_of_memory: CUDA error at: [/pyenv/versions/3.12.9/lib/python3.12/site-packages/librmm/include/rmm/mr/device/cuda_memory_resource.hpp:62](https://vscode-remote+ssh-002dremote-002b10-002e176-002e1-002e125.vscode-resource.vscode-cdn.net/pyenv/versions/3.12.9/lib/python3.12/site-packages/librmm/include/rmm/mr/device/cuda_memory_resource.hpp:62): cudaErrorMemoryAllocation out of memory
Traceback:
Traceback (most recent call last):
  File "/tmp/ipykernel_765663/1026619002.py", line 35, in scalar_quantize_cuvs
    quantizer = scalar.train(params, embeddings_gpu)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "resources.pyx", line 110, in cuvs.common.resources.auto_sync_resources.wrapper
  File "scalar.pyx", line 126, in cuvs.preprocessing.quantize.scalar.scalar.train
  File "exceptions.pyx", line 37, in cuvs.common.exceptions.check_cuvs
cuvs.common.exceptions.CuvsException: std::bad_alloc: out_of_memory: CUDA error at: [/pyenv/versions/3.12.9/lib/python3.12/site-packages/librmm/include/rmm/mr/device/cuda_memory_resource.hpp:62](https://vscode-remote+ssh-002dremote-002b10-002e176-002e1-002e125.vscode-resource.vscode-cdn.net/pyenv/versions/3.12.9/lib/python3.12/site-packages/librmm/include/rmm/mr/device/cuda_memory_resource.hpp:62): cudaErrorMemoryAllocation out of memory

Expected behavior
Successful code execution

Environment details (please complete the following information):

• Environment location: [Bare-metal, Docker, Cloud(specify cloud provider)]

• Method of RAFT install: [conda, Docker, or from source]
  pip install cuvs-cu12 --extra-index-url=https://pypi.nvidia.com

[cjnolet]

@mfoerste4 @benfred any chance you have a moment to take a look at this? Wondering if this could be a simple case where we're overflowing from the use of a 32-bit type to represent number of bytes.

[mfoerste4]

@tfeher, @cjnolet , i am able to reproduce the crash locally. The quantile selection utilizes raft::sample_without_replacement which has high memory requirements here and another worksize allocation for the subsequent cub-sortPairs call.
The internal datatypes for IdxT and WeightsT are both 8-byte wide. With 2-byte input the memory requirements sum up to about 26x of the input. We did not even hit the int32 limit with this example, although this will almost certainly raise more issues e.g. here.

The irony here is that we pick the sample in order to have fewer elements to sort for quantile selection. Now the sampling not just does a global sort, but also with pairs of an even larger data type. A quick fix could be to skip the sampling and just sort a full copy of the input.

To really solve this I believe we need to invest in dedicated quantile functionality in raft.

[cjnolet]

To really solve this I believe we need to invest in dedicated quantile functionality in raft.

This ultimately sounds like a good idea. Is there anything we can do in the mantime to reduce the memory requirement in the meantime while we ultimately work towards this? @tfeher we are seeing quite a bit of speedup in cuVS for scalar quantization over FAISS CPU version atm. it would be great if we could continue to invest in these preprocessing techniques and reduce memory requirement as much as possible.

[tfeher]

The quantile selection utilizes raft::sample_without_replacement which has high memory requirements

Déjà vu: rapidsai/raft#2155

We had the same problem earlier, and implemented a solution for that: raft::matrix::sample_rows()

[tfeher]

We might still need to revisit max_num_samples based on the available memory.