vllmで埋め込みモデルの推論を高速化する

小ネタです。

transformersの埋め込みモデルを使ってembeddingを作りたいことは割といろんなところであるかと思いますが、このembedding作成もそこそこ時間を食ったりします。

可能であれば埋め込みを計算するときも高速化したいので、今回はvllmで高速化してみたいと思います。

参考

ドキュメントはこちらです。実は埋め込みモデルに対しても利用できるようなページがありますね。

やるだけなら結構簡単のようなので、どれくらい高速化できるかやってみたいと思います。

まずは参考実装を作ってみたいと思います。

次にvllmを使った実装を作ってみます。

embeddingの中身を参照したい場合は

output[0].outputs.embedding

というところに入っているのでそちらを使えば大丈夫そうです。微妙に数字は変わってますが、まあ大体同じembeddingが得られていると考えて良いでしょう。(多分)

速度測定の結果はこんな感じになりました。

baseline	vllm
2min 45s	2min

劇的に速くなったというわけではありませんが、多少速くなりましたね。(今回の実行時間的には20%減くらいでしょうか)

データセットの条件が変わればまた違うのかもしれませんが、今回はこんな感じになりました。

vllmを使って推論の高速化ができることは知っていましたが、embeddingに使えるか試してみた次第です。一応vllmを使って高速化できたので満足です。覚えてたらまたどこかで使用してみようと思います。