parallel searching in multi-gpu setting using faiss
{'login': 'xwwwwww', 'id': 48146603, 'node_id': 'MDQ6VXNlcjQ4MTQ2NjAz', 'avatar_url': 'https://avatars.githubusercontent.com/u/48146603?v=4', 'gravatar_id': '', 'url': 'https://api.github.com/users/xwwwwww', 'html_url': 'https://github.com/xwwwwww', 'followers_url': 'https://api.github.com/users/xwwwwww/followers', 'following_url': 'https://api.github.com/users/xwwwwww/following{/other_user}', 'gists_url': 'https://api.github.com/users/xwwwwww/gists{/gist_id}', 'starred_url': 'https://api.github.com/users/xwwwwww/starred{/owner}{/repo}', 'subscriptions_url': 'https://api.github.com/users/xwwwwww/subscriptions', 'organizations_url': 'https://api.github.com/users/xwwwwww/orgs', 'repos_url': 'https://api.github.com/users/xwwwwww/repos', 'events_url': 'https://api.github.com/users/xwwwwww/events{/privacy}', 'received_events_url': 'https://api.github.com/users/xwwwwww/received_events', 'type': 'User', 'site_admin': False}
[]
[]
["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\r\n\r\n> And I don't see any speed up when increasing the number of GPUs while calling get_nearest_examples_batch.\r\n\r\nThat's unexpected, can you share the code you're running ?"
 'here is the code snippet\r\n\r\n```python\r\n\r\n# add faiss index\r\nsource_dataset = load_dataset(source_path)\r\nqueries = load_dataset(query_path)\r\ngpu = [0,1,2,3]\r\nsource_dataset.add_faiss_index(\r\n                "embedding",\r\n                device=gpu,\r\n            )\r\n\r\n\r\n# batch query\r\nbatch_size = 32\r\nfor i in tqdm(range(0, len(queries), batch_size)):\r\n    if i + batch_size >= len(queries):\r\n        batched_queries = queries[i:]\r\n    else:\r\n        batched_queries = queries[i:i+batch_size]\r\n\r\n    batched_query_embeddings = np.stack([i for i in batched_queries[\'embedding\']], axis=0)\r\n    scores, candidates = source_dataset.get_nearest_examples_batch(\r\n        "embedding",\r\n        batched_query_embeddings,\r\n        k=5\r\n    )\r\n```'
 '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.\r\nMy 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\r\nsource dataset: 1e6\r\nembedding size: 768\r\nindex: Flat\r\ntopk: 20\r\nGPU: V100\r\n\r\nThe 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 \r\n\r\nThe memory used per GPU is unchanged while adding more GPU. Is this unexpected?\r\n\r\nI 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. 😆\r\n\r\nOh indeed that's possible, I wasn't sure. Anyway you can check that calling get_nearest_examples_batch simply calls search under the hood: \r\n\r\nhttps://github.com/huggingface/datasets/blob/f90f71fbbb33889fe75a3ffc101cdf16a88a3453/src/datasets/search.py#L375"
 'Here is a runnable script. \r\nMulti-GPU searching still does not work in my experiments.\r\n\r\n\r\n```python\r\nimport os\r\nfrom tqdm import tqdm\r\nimport numpy as np\r\nimport datasets\r\nfrom datasets import Dataset\r\n\r\nclass DPRSelector:\r\n\r\n    def __init__(self, source, target, index_name, gpu=None):\r\n        self.source = source\r\n        self.target = target\r\n        self.index_name = index_name\r\n\r\n        cache_path = \'embedding.faiss\'\r\n\r\n        if not os.path.exists(cache_path):\r\n            self.source.add_faiss_index(\r\n                column="embedding",\r\n                index_name=index_name,\r\n                device=gpu,\r\n            )\r\n            self.source.save_faiss_index(index_name, cache_path)\r\n        else:\r\n            self.source.load_faiss_index(\r\n                index_name,\r\n                cache_path,\r\n                device=gpu\r\n            )\r\n        print(\'index builded!\')\r\n\r\n    def build_dataset(self, top_k, batch_size):\r\n        print(\'start search\')\r\n\r\n        for i in tqdm(range(0, len(self.target), batch_size)):\r\n            if i + batch_size >= len(self.target):\r\n                batched_queries = self.target[i:]\r\n            else:\r\n                batched_queries = self.target[i:i+batch_size]\r\n\r\n\r\n            batched_query_embeddings = np.stack([i for i in batched_queries[\'embedding\']], axis=0)\r\n            search_res = self.source.get_nearest_examples_batch(\r\n                self.index_name,\r\n                batched_query_embeddings,\r\n                k=top_k\r\n            )\r\n      \r\n        print(\'finish search\')\r\n\r\n\r\ndef get_pseudo_dataset():\r\n    pseudo_dict = {"embedding": np.zeros((1000000, 768), dtype=np.float32)}\r\n    print(\'generate pseudo data\')\r\n\r\n    dataset = Dataset.from_dict(pseudo_dict)\r\n    def list_to_array(data):\r\n        return {"embedding": [np.array(vector, dtype=np.float32) for vector in data["embedding"]]} \r\n    dataset.set_transform(list_to_array, columns=\'embedding\', output_all_columns=True)\r\n\r\n    print(\'build dataset\')\r\n    return dataset\r\n\r\n\r\n\r\nif __name__=="__main__":\r\n\r\n    np.random.seed(42)\r\n\r\n\r\n    source_dataset = get_pseudo_dataset()\r\n    target_dataset = get_pseudo_dataset()\r\n\r\n    gpu = [0,1,2,3,4,5,6,7]\r\n    selector = DPRSelector(source_dataset, target_dataset, "embedding",  gpu=gpu)\r\n\r\n    selector.build_dataset(top_k=20, batch_size=32)\r\n```'
 '@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?\r\n\r\nyes, 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:\r\n```\r\nimport faiss\r\n\r\nngpus = faiss.get_num_gpus()\r\nprint(ngpus)\r\n```\r\n\r\nWhen passing a list of integers to `device`, then only that number of GPUs are used (multi GPU as well)\r\n- the speedup should be proportional (more or less) to the ratio of the number of elements passed to `device` over `ngpus`\r\n- 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)\r\n\r\nWhen passing a positive integer to `device`, then only a single GPU is used.\r\n- this time should be more or less proportional to the time when passing `device=-1` over `ngpus`"
 'Thanks for your help!\r\nHave 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))`.\r\n\r\nCould 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 ?']
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
{'url': 'https://api.github.com/repos/huggingface/datasets/issues/4761/reactions', 'total_count': 0, '+1': 0, '-1': 0, 'laugh': 0, 'hooray': 0, 'confused': 0, 'heart': 0, 'rocket': 0, 'eyes': 0}