GPU memory usage of IndexIVFFlat

[jiangtann]

Question 1: options.useFloat16 cannot reduce GPU memory usage

embeddings = np.random.rand(5000000, 2048).astype(np.float32)

index = faiss.IndexFlatIP(2048)
index = faiss.IndexIVFFlat(index, 2048, 32768)

options = faiss.GpuMultipleClonerOptions()
options.shard = True
options.useFloat16 = True
gpu_index = faiss.index_cpu_to_all_gpus(index, co=options, ngpu=1)

gpu_index.train(embeddings[:1000000])

gpu_index.add(embeddings)

Whether I set options.useFloat16 = True or not, the GPU memory usage remains consistent. Does IndexIVF (e.g. IndexIVFFlat, IndexIVFPQ) require all vectors to be stored in fp32?

Question 2: How to further reduce GPU memory usage

I have a huge (56000000, 2048) embedding, which is similar to #4502. I use 8 A100 80G GPUs to perform index.add. Even though this embedding alone consumes 56,000,000 * 2048 * 4 bytes, approximately 427GB, due to useFloat16 not working and extra GPU memory usage, my total of 640GB across 8 A100 GPUs still couldn't smoothly complete the index.add operation.

Are there any methods to reduce GPU memory usage? Or is it impossible for me to use IndexIVFFlat in this situation?

[limqiying]

Hey @jiangtann, what are you using to measure memory? Is it nvidia-smi? I am trying to reproduce your issue.

[jiangtann]

Yes.

I ran two different tests, one on GPU 0 and another on GPU 1. GPU 0 was not configured with options.useFloat16 = True, while GPU 1 was. After gpu_index.add(embeddings), use nvidia-smi:

use nvitop(need pip install nvitop):

[jiangtann]

If set options.shard = True and use multiple GPUs, the total GPU memory usage is also about 51G.

[mnorris11]

@jiangtann thanks for the code snippet, it makes it a lot easier. Is quantization an option for you? IVFSQ will maintain higher accuracy, but be higher memory usage than IVFPQ. If you need PQ quantization at high dimensions like this, you can try cuVS IVFPQ through Faiss, which has no limit on subquantizer, as long as it is divisible by dimensions (so, you could do a PQ with M=1024 with cuVS IVFPQ).

I did some GPU benchmarking some time back. There is some memory pooling that happens as well for faster processing, so memory observed on nvidia-smi won't be linear. For example, I just tried your snippet with 1M and 2.5M. Peak usages were 23GB and 35GB, which is not what a straight calculation would give (1M * 2048 * 4 bytes = 8GB), and also 35 is not 2.5 times 23.

[jiangtann]

Thanks for your advise, I will try IVFSQ and cuVS IVFPQ.

I still have questions about question 1. Does options.useFloat16 have no impact on GPU memory usage because there's an issue with my code? Or is this phenomenon normal?