Improve: batch size setting and multi GPU inference with SentenceTransformers+DP - 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

Improve: batch size setting and multi GPU inference with SentenceTransformers+DP #53

Closed akiFQC closed 1 month ago

akiFQC commented 1 month ago

関連する Issue / PR

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

バッチサイズを最適化して高速に埋め込みを計算したい
安定性のために、torchrun を使わずにGPU並列化を行いたい

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

src/jmteb/embedders/base.pyのTextEmbedder に _chunk_size という変数を追加し、_chunk_size ごとに TextEmbedder.encodeを呼ぶように変更
src/jmteb/embedders/data_parallel_sbert_embedder.py追加し、sentence-transfermersでpytorchのDataParallelによる複数GPU推論が可能になった。
テストを追加

動作確認

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

lsz05 commented 1 month ago

起動時torchrunではなくpythonで起動すると理解合っていますでしょうか?

akiFQC commented 1 month ago

起動時torchrunではなくpythonで起動すると理解合っていますでしょうか?

はい。pythonで起動してください。