Open gabuzi opened 7 months ago
Thanks for the precise report. I don't know if we can do something about it @wickedfoo ?
Would there be an alternative approach that is slower but more parsimonious in memory?
Thanks for getting back to me and sorry for the delay.
I may have been a bit hasty in saying "we need such high k". The reality was more like that we were looking into it, but we have abandoned this this high-k approach due to increased temporary memory requirements during search. This is no longer an issue for us, but the problem as described in the initial post still stands.
Our main rationale for using such quantization was to save GPU memory. Naturally, the temporary memory on the GPU grows with k (pre-refinement), which we have naively ignored. Memory savings by storing the index data compressed can thus easily be annihilated by the temporary allocations required for storing the distances and indices to the temporary results on the GPU before refinement. The solution is then to reduce the query batchsize. This has been exacerbated a bit by the recent move to 64-bit indexing on the GPU.
Summary
GpuIndexIVFScalarQuantizer with scalar quantizers that require shared memory on the GPU don't seem to work for k >= 1024 in Faiss 1.7.4. See the small reproduction script at the bottom.
I need such high k in order to make up for quantization errors via refinement on CPU.
My investigation points towards exhausting the shared memory available on GPU without opt-in (see https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#fn33).
GPUs by default can only address 48kiB shared memory per block without special opt-in (see link above), but devices since compute capability 7 can potentially address more, but do require dynamic allocation and the opt-in.
For k = 1024, faiss uses 128 threads (faiss/gpu/impl/scan/IVFInterleaved1024.cu). Since the move to 64-bit indexing, this configuration fully exhausts the 48kiB shared memory already as shown in the following two lines (kNumWarps = 4 for 128 threads, and NumWarpQ is 1024) https://github.com/facebookresearch/faiss/blob/0fc8456e1db0846759bd8d76a03ebe497c511a9f/faiss/gpu/impl/IVFInterleaved.cuh#L83-L84
This leaves no shared memory to be allocated dynamically in the kernel launch resulting in the cuda error.
For k = 2048, the problem is identical, as this gets launched with half the number of threads, but double the NumWarpQ, again filling the 48k of statically allocatable shared memory.
Quick workarounds that I have successfully tested (equal results vs CPU indexes, no performance tests, and both seem suboptimal / hacky):
code for option 2: add to IVFInterleaved.cuh in the IVFINT_RUN macro for both the
useResidual
cases. Edited to add: This is based on 1.7.4. I now saw that this area has seen some recent changes, but apparently only a refactor to reorganize code files, and the solution could still hold.Related: The check in https://github.com/facebookresearch/faiss/blob/7dd06dd188643523d1a0faab511ee8fc0c5b0e3b/faiss/gpu/GpuIndexIVFScalarQuantizer.cu#L92-L110 doesn't seem to consider that the maxK selection also requires shared memory.
Platform
x86-64, intel RTX 4090
OS: Ubuntu 20.04 LTS
Faiss version: 1.7.4 Installed from: anaconda in clean docker image from nvidia/cuda:11.4.3-runtime-ubi8, but also built myself install command:
conda create -y -n faiss-issue -c pytorch -c nvidia python=3.10 faiss-gpu=1.7.4 mkl=2021 blas=1.0=mkl
Faiss compilation options:
Running on:
Interface:
Reproduction instructions
Output:
Edited to remove a leftover from a modification that I made at some point and a comment.