[Fix] handle prompt length for multi-GPU #87
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
[dev to main] v1.3.1
[dev to main] v1.3.1 hot fixes
[dev to main] v1.3.2
…ngth_for_multi_gpus
lsz05
approved these changes
Dec 9, 2024
| @@ -7,7 +7,7 @@ | |||
| import torch | |||
| from accelerate.utils import find_executable_batch_size | |||
| from loguru import logger | |||
| from sentence_transformers import SentenceTransformer | |||
| from sentence_transformers import SentenceTransformer, models | |||
There was a problem hiding this comment.
nit: modelsの中のclassは一個しか使われなかったので，from sentence_transformers.models import Poolingと書くのがどうでしょう（self.modelと混同するのを防ぐため）
There was a problem hiding this comment.
確かにそちらの方が良さそうですね、ありがとうございます、変更します！(modelsはちょっと一般的すぎるモジュール名で嫌ですよね)
lsz05
changed the title
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
関連する Issue / PR
埋め込み作成時にpromptを利用するモデルについて、DPSentenceTransformerを用いてmulti GPUで
encode関数を用いるとエラーが発生する。PR をマージした後の挙動の変化
encode関数では内部でprompt_lengthと呼ばれる値を作成している。埋め込み作成時にこの値がdict型の
feature変数に格納される。この時、prompt_lengthはint型である。encode関数はDPでの推論時にfeatureの中身をgatherするが、このgatherはtorch.Tensor型以外を受け取るとエラーになる。現状は
prompt_lengthがfeature変数を介してgather関数に渡っているため、promptとともにencode関数を使用するとエラーが発生する。この問題を解決する。
挙動の変更を達成するために行ったこと
対策として、
feature中にprompt_lengthがkeyとして含まれる場合、その値を入力文数と同じ数・同じ値を持つtorch.Tensor型に変換する。また、そもそも
prompt_lengthはinclude_prompt=Falseなモデルのためのパラメータであるが、include_prompt=Trueの場合はprompt_length自体が不要なので、その場合は事前にkeyを消しておく。動作確認として、埋め込み作成時にprefixを用いる例をテストに追加した。
動作確認