Imbalanced GPU Load when generating the embeddings.

[bhattg]

Hi!

I've been trying to run the following command --

python bergen.py retriever=splade-v3 reranker=debertav3 dataset=popqa

To generate the embeddings in a reasonable time, I modified the retriever config batch size with batch_size: 512. However, as I am looking at my nvidia-smi board, it seems to be very heavily imbalanced.

========================================+======================+======================|
|   0  NVIDIA A100 80GB PCIe          On  | 00000000:18:00.0 Off |                    0 |
| N/A   39C    P0              70W / 300W |  27951MiB / 81920MiB |     17%      Default |
|                                         |                      |             Disabled |
+-----------------------------------------+----------------------+----------------------+
|   1  NVIDIA A100 80GB PCIe          On  | 00000000:3B:00.0 Off |                    0 |
| N/A   38C    P0             264W / 300W |   3163MiB / 81920MiB |      0%      Default |
|                                         |                      |             Disabled |
+-----------------------------------------+----------------------+----------------------+
|   2  NVIDIA A100 80GB PCIe          On  | 00000000:86:00.0 Off |                    0 |
| N/A   39C    P0             249W / 300W |   3163MiB / 81920MiB |      0%      Default |
|                                         |                      |             Disabled |
+-----------------------------------------+----------------------+----------------------+
|   3  NVIDIA A100 80GB PCIe          On  | 00000000:AF:00.0 Off |                    0 |
| N/A   38C    P0             264W / 300W |   3163MiB / 81920MiB |      0%      Default |
|                                         |                      |             Disabled |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      1850      G   /usr/lib/xorg/Xorg                            4MiB |
|    0   N/A  N/A    739593      C   python                                    27924MiB |
|    1   N/A  N/A      1850      G   /usr/lib/xorg/Xorg                            4MiB |
|    1   N/A  N/A    739593      C   python                                     3136MiB |
|    2   N/A  N/A      1850      G   /usr/lib/xorg/Xorg                            4MiB |
|    2   N/A  N/A    739593      C   python                                     3136MiB |
|    3   N/A  N/A      1850      G   /usr/lib/xorg/Xorg                            4MiB |
|    3   N/A  N/A    739593      C   python                                     3136MiB |
+---------------------------------------------------------------------------------------+

Can anyone please help me with this? I am using transformers v4.46.3 and torch v2.3.0+cu121

[sclincha]

Hi @bhattg !
Sorry for the late reply. I think that the multi-gpu parrellization is not yet supported for the retriever and the reranker. This is one of the feature we wish to improve very soon. In case you need the search results in 'retriever=splade-v3 reranker=debertav3 dataset=popqa', the trec runs is included in the run folder (ie https://github.com/naver/bergen/blob/main/runs/run.rerank.retriever.top_50.naver_splade-v3.rerank.top_50.popqa.kilt-100w.dev.naver_trecdl22-crossencoder-debertav3.trec)