Issue: parallel searching in multi-gpu setting using faiss

While I notice that  `add_faiss_index` has supported assigning multiple GPUs, I am still confused about how it works. 

Does the `search-batch` function automatically parallelizes the input queries to different gpus?https://github.com/huggingface/datasets/blob/d76599bdd4d186b2e7c4f468b05766016055a0a5/src/datasets/search.py#L360

Comments:
And I don't see any speed up when increasing the number of GPUs while calling `get_nearest_examples_batch`.

Hi ! Yes search_batch uses FAISS search which happens in parallel across the GPUs

> And I don't see any speed up when increasing the number of GPUs while calling get_nearest_examples_batch.

That's unexpected, can you share the code you're running ?

here is the code snippet

```python

# add faiss index
source_dataset = load_dataset(source_path)
queries = load_dataset(query_path)
gpu = [0,1,2,3]
source_dataset.add_faiss_index(
                "embedding",
                device=gpu,
            )


# batch query
batch_size = 32
for i in tqdm(range(0, len(queries), batch_size)):
    if i + batch_size >= len(queries):
        batched_queries = queries[i:]
    else:
        batched_queries = queries[i:i+batch_size]

    batched_query_embeddings = np.stack([i for i in batched_queries['embedding']], axis=0)
    scores, candidates = source_dataset.get_nearest_examples_batch(
        "embedding",
        batched_query_embeddings,
        k=5
    )
```

My version of datasets is `2.4.1.dev0`.

The code looks all good to me, do you see all the GPUs being utilized ? What version of faiss are you using ?

I can see the memory usage of all the GPUs.
My version of `faiss-gpu` is `1.7.2`

It looks all good to me then ^^ though you said you didn't experienced speed improvements by adding more GPUs ? What size is your source dataset and what time differences did you experience ?

query set: 1e6
source dataset: 1e6
embedding size: 768
index: Flat
topk: 20
GPU: V100

The time taken to traverse the query set once is about 1.5h, which is almost not influenced by the value of query batch size or the number of GPUs according to my experiments.

Hmmm the number of GPUs should divide the time, something is going wrong. Can you check that adding more GPU does divide the memory used per GPU ? Maybe it can be worth looking at similar issues in the FAISS repository or create a noew issue over there to understand what's going on

> Can you check that adding more GPU does divide the memory used per GPU 

The memory used per GPU is unchanged while adding more GPU. Is this unexpected?

I used to think that every GPU loads all the source vectors and the data parallelism is at the query level. 😆 

> I used to think that every GPU loads all the source vectors and the data parallelism is at the query level. 😆

Oh indeed that's possible, I wasn't sure. Anyway you can check that calling get_nearest_examples_batch simply calls search under the hood: 

https://github.com/huggingface/datasets/blob/f90f71fbbb33889fe75a3ffc101cdf16a88a3453/src/datasets/search.py#L375

Here is a runnable script. 
Multi-GPU searching still does not work in my experiments.


```python
import os
from tqdm import tqdm
import numpy as np
import datasets
from datasets import Dataset

class DPRSelector:

    def __init__(self, source, target, index_name, gpu=None):
        self.source = source
        self.target = target
        self.index_name = index_name

        cache_path = 'embedding.faiss'

        if not os.path.exists(cache_path):
            self.source.add_faiss_index(
                column="embedding",
                index_name=index_name,
                device=gpu,
            )
            self.source.save_faiss_index(index_name, cache_path)
        else:
            self.source.load_faiss_index(
                index_name,
                cache_path,
                device=gpu
            )
        print('index builded!')

    def build_dataset(self, top_k, batch_size):
        print('start search')

        for i in tqdm(range(0, len(self.target), batch_size)):
            if i + batch_size >= len(self.target):
                batched_queries = self.target[i:]
            else:
                batched_queries = self.target[i:i+batch_size]


            batched_query_embeddings = np.stack([i for i in batched_queries['embedding']], axis=0)
            search_res = self.source.get_nearest_examples_batch(
                self.index_name,
                batched_query_embeddings,
                k=top_k
            )
      
        print('finish search')


def get_pseudo_dataset():
    pseudo_dict = {"embedding": np.zeros((1000000, 768), dtype=np.float32)}
    print('generate pseudo data')

    dataset = Dataset.from_dict(pseudo_dict)
    def list_to_array(data):
        return {"embedding": [np.array(vector, dtype=np.float32) for vector in data["embedding"]]} 
    dataset.set_transform(list_to_array, columns='embedding', output_all_columns=True)

    print('build dataset')
    return dataset



if __name__=="__main__":

    np.random.seed(42)


    source_dataset = get_pseudo_dataset()
    target_dataset = get_pseudo_dataset()

    gpu = [0,1,2,3,4,5,6,7]
    selector = DPRSelector(source_dataset, target_dataset, "embedding",  gpu=gpu)

    selector.build_dataset(top_k=20, batch_size=32)
```

@lhoestq Hi, could you please test the code above if you have time? 😄 

Maybe @albertvillanova you can take a look ? I won't be available in the following days

@albertvillanova Hi, can you help with this issue?

Hi @xwwwwww I'm investigating it, but I'm not an expert in Faiss. In principle, it is weird that your code does not work properly because it seems right...

Have you tried passing `gpu=-1` and check if there is a speedup?

> Have you tried passing `gpu=-1` and check if there is a speedup?

yes, there is a speed up using GPU compared with CPU. 

When passing `device=-1`, ALL existing GPUs are used (multi GPU): this is the maximum speedup you can get. To know the number of total GPUs:
```
import faiss

ngpus = faiss.get_num_gpus()
print(ngpus)
```

When passing a list of integers to `device`, then only that number of GPUs are used (multi GPU as well)
- the speedup should be proportional (more or less) to the ratio of the number of elements passed to `device` over `ngpus`
- if this is not the case, then there is an issue in the implementation of this use case (however, I have reviewed the code and in principle I can't find any evident bug)

When passing a positive integer to `device`, then only a single GPU is used.
- this time should be more or less proportional to the time when passing `device=-1` over `ngpus`

Thanks for your help!
Have you run the code and replicated the same experimental results (i.e., no speedup while increasing the number of GPUs)?

@albertvillanova @lhoestq Sorry for the bother, is there any progress on this issue? 😃 

I can confirm `add_faiss_index` calls `index = faiss.index_cpu_to_gpus_list(index, gpus=list(device))`.

Could this be an issue with your environment ? Could you try running with 1 and 8 GPUs with a code similar to[ this one from the FAISS examples](https://github.com/facebookresearch/faiss/blob/main/tutorial/python/5-Multiple-GPUs.py) but using `gpu_index = faiss.index_cpu_to_gpus_list(cpu_index, gpus=list(device))`, and see if the speed changes ?

Hi, I test the FAISS example and the speed indeed changes.  I set `nb=1000000`, `nq=1000000` and `d=64`

| num GPUS | time cost |
| -------- | --------- |
| 1        | 28.53     |
| 5        | 7.16      |





Ok the benchmark is great, not sure why it doesn't speed up the index in your case though. You can try running the benchmark with the same settings as your actual dataset
```
query set: 1e6
source dataset: 1e6
embedding size: 768
index: Flat
topk: 20
GPU: V100
```

Note that you can still pass a FAISS index you built yourself to a dataset using https://huggingface.co/docs/datasets/v2.4.0/en/package_reference/main_classes#datasets.Dataset.add_faiss_index_from_external_arrays

> Here is a runnable script. Multi-GPU searching still does not work in my experiments.
> 
> ```python
> import os
> from tqdm import tqdm
> import numpy as np
> import datasets
> from datasets import Dataset
> 
> class DPRSelector:
> 
>     def __init__(self, source, target, index_name, gpu=None):
>         self.source = source
>         self.target = target
>         self.index_name = index_name
> 
>         cache_path = 'embedding.faiss'
> 
>         if not os.path.exists(cache_path):
>             self.source.add_faiss_index(
>                 column="embedding",
>                 index_name=index_name,
>                 device=gpu,
>             )
>             self.source.save_faiss_index(index_name, cache_path)
>         else:
>             self.source.load_faiss_index(
>                 index_name,
>                 cache_path,
>                 device=gpu
>             )
>         print('index builded!')
> 
>     def build_dataset(self, top_k, batch_size):
>         print('start search')
> 
>         for i in tqdm(range(0, len(self.target), batch_size)):
>             if i + batch_size >= len(self.target):
>                 batched_queries = self.target[i:]
>             else:
>                 batched_queries = self.target[i:i+batch_size]
> 
> 
>             batched_query_embeddings = np.stack([i for i in batched_queries['embedding']], axis=0)
>             search_res = self.source.get_nearest_examples_batch(
>                 self.index_name,
>                 batched_query_embeddings,
>                 k=top_k
>             )
>       
>         print('finish search')
> 
> 
> def get_pseudo_dataset():
>     pseudo_dict = {"embedding": np.zeros((1000000, 768), dtype=np.float32)}
>     print('generate pseudo data')
> 
>     dataset = Dataset.from_dict(pseudo_dict)
>     def list_to_array(data):
>         return {"embedding": [np.array(vector, dtype=np.float32) for vector in data["embedding"]]} 
>     dataset.set_transform(list_to_array, columns='embedding', output_all_columns=True)
> 
>     print('build dataset')
>     return dataset
> 
> 
> 
> if __name__=="__main__":
> 
>     np.random.seed(42)
> 
> 
>     source_dataset = get_pseudo_dataset()
>     target_dataset = get_pseudo_dataset()
> 
>     gpu = [0,1,2,3,4,5,6,7]
>     selector = DPRSelector(source_dataset, target_dataset, "embedding",  gpu=gpu)
> 
>     selector.build_dataset(top_k=20, batch_size=32)
> ```

By the way,  have you run this toy example and replicated my experiment results? I think it is a more direct way to figure this out :)

Hi,

I have a similar question and would like to know if there's any progress in this issue. 

`dataset.add_faiss_index(column="embedding")`, this takes around 5minutes to add the index.

`dataset.add_faiss_index(column="embedding", device=-1)`, this ran for more than 10minutes and still didn't complete execution. 

Now, I don't understand why that's the case as I expected for GPU the indexing should be faster