[Fix] multi-GPU rerank retrievalタスクにおいてOOM防止のための修正 - Githubissues

sbintuitions / JMTEB

The evaluation scripts of JMTEB (Japanese Massive Text Embedding Benchmark)

Creative Commons Attribution Share Alike 4.0 International

24 stars 4 forks source link

[Fix] multi-GPU rerank retrievalタスクにおいてOOM防止のための修正 #51

Closed akiFQC closed 1 month ago

akiFQC commented 1 month ago

関連する Issue / PR

https://github.com/sbintuitions/JMTEB/issues/50

PR をマージした後の挙動の変化

緊急のdebugとして、TransfomersEmbedderで複数GPUを使った場合のOOMを防ぐ

挙動の変更を達成するために行ったこと

(緊急のdebugとして) torchrunのそれぞれのプロセスで、異なるGPUで類似度行列の計算を行う
類似度行列のchunk sizeを1000000→10000に変更

動作確認

[x] テストが通ることを確認した
[x] マージ先がdevブランチであることを確認した