parallel searching in multi-gpu setting using faiss
CONTRIBUTOR
And I don't see any speed up when increasing the number of GPUs while calling `get_nearest_examples_batch`.
2022-07-28T15:07:32Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1198282457
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HbFbZ
https://api.github.com/repos/huggingface/datasets/issues/comments/1198282457/reactions
2022-07-28T15:07:32Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1198282457
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
MEMBER
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 ?
2022-07-29T09:31:20Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199069659
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HeFnb
https://api.github.com/repos/huggingface/datasets/issues/comments/1199069659/reactions
2022-07-29T09:31:20Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199069659
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
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
    )
```
2022-07-29T09:58:56Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199093310
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HeLY-
https://api.github.com/repos/huggingface/datasets/issues/comments/1199093310/reactions
2022-07-29T09:59:52Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199093310
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
CONTRIBUTOR
My version of datasets is `2.4.1.dev0`.
2022-07-29T10:00:39Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199094908
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HeLx8
https://api.github.com/repos/huggingface/datasets/issues/comments/1199094908/reactions
2022-07-29T10:00:39Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199094908
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
MEMBER
The code looks all good to me, do you see all the GPUs being utilized ? What version of faiss are you using ?
2022-07-29T13:18:37Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199276639
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5He4Jf
https://api.github.com/repos/huggingface/datasets/issues/comments/1199276639/reactions
2022-07-29T13:18:37Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199276639
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
I can see the memory usage of all the GPUs.
My version of `faiss-gpu` is `1.7.2`
2022-07-29T13:20:59Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199278970
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5He4t6
https://api.github.com/repos/huggingface/datasets/issues/comments/1199278970/reactions
2022-07-29T13:20:59Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199278970
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
MEMBER
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 ?
2022-07-29T13:30:11Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199288429
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5He7Bt
https://api.github.com/repos/huggingface/datasets/issues/comments/1199288429/reactions
2022-07-29T13:30:11Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199288429
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
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.
2022-07-29T13:38:54Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199298078
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5He9Ye
https://api.github.com/repos/huggingface/datasets/issues/comments/1199298078/reactions
2022-07-29T13:48:57Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199298078
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
MEMBER
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
2022-07-29T13:54:56Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199315104
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HfBig
https://api.github.com/repos/huggingface/datasets/issues/comments/1199315104/reactions
2022-07-29T13:54:56Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199315104
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
> 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. 😆 
2022-07-29T14:05:51Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199338664
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HfHSo
https://api.github.com/repos/huggingface/datasets/issues/comments/1199338664/reactions
2022-07-29T14:05:51Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199338664
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
MEMBER
> 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
2022-07-29T14:09:47Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199347873
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HfJih
https://api.github.com/repos/huggingface/datasets/issues/comments/1199347873/reactions
2022-07-29T14:09:47Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199347873
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
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)
```
2022-07-29T15:24:34Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1199527204
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5Hf1Uk
https://api.github.com/repos/huggingface/datasets/issues/comments/1199527204/reactions
2022-07-29T15:24:34Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1199527204
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
CONTRIBUTOR
@lhoestq Hi, could you please test the code above if you have time? 😄 
2022-07-30T06:53:26Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1200104399
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HiCPP
https://api.github.com/repos/huggingface/datasets/issues/comments/1200104399/reactions
2022-07-30T06:53:26Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1200104399
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
MEMBER
Maybe @albertvillanova you can take a look ? I won't be available in the following days
2022-07-31T20:51:12Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1200497792
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HjiSA
https://api.github.com/repos/huggingface/datasets/issues/comments/1200497792/reactions
2022-07-31T20:51:12Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1200497792
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
@albertvillanova Hi, can you help with this issue?
2022-08-02T01:36:46Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1201916028
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5Ho8h8
https://api.github.com/repos/huggingface/datasets/issues/comments/1201916028/reactions
2022-08-02T01:36:46Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1201916028
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
MEMBER
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...
2022-08-02T05:33:34Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1202036741
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HpaAF
https://api.github.com/repos/huggingface/datasets/issues/comments/1202036741/reactions
2022-08-02T05:33:34Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1202036741
https://avatars.githubusercontent.com/u/8515462?v=4
https://api.github.com/users/albertvillanova/events{/privacy}
https://api.github.com/users/albertvillanova/followers
https://api.github.com/users/albertvillanova/following{/other_user}
https://api.github.com/users/albertvillanova/gists{/gist_id}
https://github.com/albertvillanova
albertvillanova
MDQ6VXNlcjg1MTU0NjI=
https://api.github.com/users/albertvillanova/orgs
https://api.github.com/users/albertvillanova/received_events
https://api.github.com/users/albertvillanova/repos
https://api.github.com/users/albertvillanova/starred{/owner}{/repo}
https://api.github.com/users/albertvillanova/subscriptions
User
https://api.github.com/users/albertvillanova
public
MEMBER
Have you tried passing `gpu=-1` and check if there is a speedup?
2022-08-02T05:48:30Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1202045635
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HpcLD
https://api.github.com/repos/huggingface/datasets/issues/comments/1202045635/reactions
2022-08-02T05:48:30Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1202045635
https://avatars.githubusercontent.com/u/8515462?v=4
https://api.github.com/users/albertvillanova/events{/privacy}
https://api.github.com/users/albertvillanova/followers
https://api.github.com/users/albertvillanova/following{/other_user}
https://api.github.com/users/albertvillanova/gists{/gist_id}
https://github.com/albertvillanova
albertvillanova
MDQ6VXNlcjg1MTU0NjI=
https://api.github.com/users/albertvillanova/orgs
https://api.github.com/users/albertvillanova/received_events
https://api.github.com/users/albertvillanova/repos
https://api.github.com/users/albertvillanova/starred{/owner}{/repo}
https://api.github.com/users/albertvillanova/subscriptions
User
https://api.github.com/users/albertvillanova
public
CONTRIBUTOR
> Have you tried passing `gpu=-1` and check if there is a speedup?

yes, there is a speed up using GPU compared with CPU. 
2022-08-02T12:21:06Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1202415255
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5Hq2aX
https://api.github.com/repos/huggingface/datasets/issues/comments/1202415255/reactions
2022-08-02T12:21:06Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1202415255
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
MEMBER
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`
2022-08-02T16:58:26Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1202988917
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HtCd1
https://api.github.com/repos/huggingface/datasets/issues/comments/1202988917/reactions
2022-08-02T17:19:02Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1202988917
https://avatars.githubusercontent.com/u/8515462?v=4
https://api.github.com/users/albertvillanova/events{/privacy}
https://api.github.com/users/albertvillanova/followers
https://api.github.com/users/albertvillanova/following{/other_user}
https://api.github.com/users/albertvillanova/gists{/gist_id}
https://github.com/albertvillanova
albertvillanova
MDQ6VXNlcjg1MTU0NjI=
https://api.github.com/users/albertvillanova/orgs
https://api.github.com/users/albertvillanova/received_events
https://api.github.com/users/albertvillanova/repos
https://api.github.com/users/albertvillanova/starred{/owner}{/repo}
https://api.github.com/users/albertvillanova/subscriptions
User
https://api.github.com/users/albertvillanova
public
CONTRIBUTOR
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)?
2022-08-03T00:33:12Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1203350275
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5HuasD
https://api.github.com/repos/huggingface/datasets/issues/comments/1203350275/reactions
2022-08-03T00:33:12Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1203350275
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
CONTRIBUTOR
@albertvillanova @lhoestq Sorry for the bother, is there any progress on this issue? 😃 
2022-08-07T13:52:54Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1207413627
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5H96t7
https://api.github.com/repos/huggingface/datasets/issues/comments/1207413627/reactions
2022-08-07T13:52:54Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1207413627
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
MEMBER
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 ?
2022-08-18T14:59:10Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1219600695
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5IsaE3
https://api.github.com/repos/huggingface/datasets/issues/comments/1219600695/reactions
2022-08-18T14:59:10Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1219600695
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
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      |




2022-08-21T02:32:58Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1221453775
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5IzefP
https://api.github.com/repos/huggingface/datasets/issues/comments/1221453775/reactions
2022-08-21T02:32:58Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1221453775
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
MEMBER
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
2022-08-24T15:28:32Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1225880899
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5JEXVD
https://api.github.com/repos/huggingface/datasets/issues/comments/1225880899/reactions
2022-08-24T15:28:32Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1225880899
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
> 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 :)
2022-08-27T02:08:49Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1229099747
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5JQpLj
https://api.github.com/repos/huggingface/datasets/issues/comments/1229099747/reactions
2022-08-27T02:08:49Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1229099747
https://avatars.githubusercontent.com/u/48146603?v=4
https://api.github.com/users/Jiaxin-Wen/events{/privacy}
https://api.github.com/users/Jiaxin-Wen/followers
https://api.github.com/users/Jiaxin-Wen/following{/other_user}
https://api.github.com/users/Jiaxin-Wen/gists{/gist_id}
https://github.com/Jiaxin-Wen
Jiaxin-Wen
MDQ6VXNlcjQ4MTQ2NjAz
https://api.github.com/users/Jiaxin-Wen/orgs
https://api.github.com/users/Jiaxin-Wen/received_events
https://api.github.com/users/Jiaxin-Wen/repos
https://api.github.com/users/Jiaxin-Wen/starred{/owner}{/repo}
https://api.github.com/users/Jiaxin-Wen/subscriptions
User
https://api.github.com/users/Jiaxin-Wen
public
NONE
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
2023-07-21T02:07:09Z
https://github.com/huggingface/datasets/issues/4761#issuecomment-1644887900
https://api.github.com/repos/huggingface/datasets/issues/4761
IC_kwDODunzps5iCv9c
https://api.github.com/repos/huggingface/datasets/issues/comments/1644887900/reactions
2023-07-21T02:07:09Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1644887900
https://avatars.githubusercontent.com/u/40389487?v=4
https://api.github.com/users/NamburiSrinath/events{/privacy}
https://api.github.com/users/NamburiSrinath/followers
https://api.github.com/users/NamburiSrinath/following{/other_user}
https://api.github.com/users/NamburiSrinath/gists{/gist_id}
https://github.com/NamburiSrinath
NamburiSrinath
MDQ6VXNlcjQwMzg5NDg3
https://api.github.com/users/NamburiSrinath/orgs
https://api.github.com/users/NamburiSrinath/received_events
https://api.github.com/users/NamburiSrinath/repos
https://api.github.com/users/NamburiSrinath/starred{/owner}{/repo}
https://api.github.com/users/NamburiSrinath/subscriptions
User
https://api.github.com/users/NamburiSrinath
public
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