Cannot search multi-GPU index with GPU tensor and faiss.index_cpu_to_all_gpus

[greeneggsandyaml]

Summary

I am trying to use a multi-GPU index (with the standard IndexReplicas multi-GPU mode).

I have tensors that are already on the GPU, and I am trying to use them to train/add to/search the index.

I am doing import faiss.contrib.torch_utils to get GPU support, and this works for a single-GPU setup. Unfortunately, it does not seem to work for a multi-GPU setup. I have included more details below.

Platform

OS: Ubuntu 20.04.1

Faiss version: 1.7.1

Installed from: anaconda

Running on:

• CPU
• GPU

Interface:

• C++
• Python

Reproduction instructions

Here is the setup:

import torch
import faiss
import faiss.contrib.torch_utils

# Device
device = 'cuda'

# Load features
features = torch.randn(10000, 768, device=device)
queries = torch.randn(100, 768, device=device)

# Params
N, DIM = features.shape
K = 10
NLIST = 200
M = 16
BITS = 8

# Create index
quantizer = faiss.IndexFlatIP(DIM)
index = faiss.IndexIVFPQ(quantizer, DIM, NLIST, M, BITS)

Now, the following code works (single GPU):

# Single GPU
res = faiss.StandardGpuResources()  # use a single GPU
index = faiss.index_cpu_to_gpu(res, 0, index)

# Train
index.train(features)

I would expect the following code to work (multi-GPU):

# Multiple GPUs
print(f"Using {faiss.get_num_gpus()} GPUs")
index = faiss.index_cpu_to_all_gpus(index)

# Train
index.train(features)

However, I get the following error:

Traceback (most recent call last):
  File "helpers/create-faiss-database.py", line 43, in <module>
    index.train(features)
  File "/home/luke/anaconda3/envs/faiss/lib/python3.8/site-packages/faiss/contrib/torch_utils.py", line 195, in torch_replacement_train
    assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
AssertionError: GPU tensor on CPU index not allowed

Honestly, this is not a big deal for the training and adding vectors to the index. However, it is a big deal for searching the index, because my data is already on GPU. I could transfer it to CPU, but that would be slow. If I train/add to the index using CPU numpy arrays and then search with a torch tensor, I get the same error.

As a note, I can see with nvidia-smi that the index is correctly using the GPUs, even when I add them as CPU bumpy arrays.

I feel like there's something obvious here that I'm missing. Could someone provide a bit of help?

[KevinMusgrave]

This is a pretty big issue

[SirRob1997]

Is there a workaround for this? I've encountered something similar described here #2074. This is a pretty big issue!!

[phanxuanduc1996]

I used Anaconda, I installed bellow, it worked. If you use Docker, please install Miniconda.
conda install faiss-gpu cudatoolkit=11.1 -c pytorch-gpu
conda install -c anaconda pytorch-gpu

[mdouze]

The IndexReplicas is a CPU index. There is no native GPU-only index support.
Nb that for most use cases, CPU->GPU transfer or even a roundtrip is not the bottleneck. The search time is the bottleneck.