CTranslate2という Python と C++で書かれた高速推論用ライブラリがあり、いつか試そうと思っていたのだけど、モデルを変換する必要があったため億劫になって試していなかった。しかし hf_hub_ctranslate2 という、何もやらずにも透過的に HuggingFace のモデルを CTranslate2 で推論できる形式に変換して利用できるライブラリを知ったので試してみたところ、とても簡単に GPU で推論が 1.6 倍速に、CPU で1.9倍速になり、かつ精度もほぼ変わらなかったので、もっと早く使うべきだった、のでメモ。

CTranslate2 とは

CTranslate2(以下CT2) とは、GitHub プロジェクトページの概要に書かれている "CTranslate2 is a C++ and Python library for efficient inference with Transformer models." との通り、Transformers のモデルをさまざまな最適化によって効率的に推論するライブラリである。同様に効率推論するライブラリのllama.cppなどは、基本デコーダーモデルのみに対応しているが、CT2 はデコーダーモデルのみならず、エンコーダー・デコーダーモデルとエンコーダーモデルのある程度にも対応している。対応エンコーダーモデルではBertにも対応しているため、Bert系列モデルも効率的な推論が可能である。

Bert?そんな古いアーキテクチャを使う機会あるの？と思う方も居るかもしれないが、例えば私が日々 embeddings を生成しているモデルである、multilingual-e5-smallも御多分に洩れずBert系列のモデルであり、まだまだ使う機会は多い。

CTranslate2 を使った embeddings (SentenceTransformer) の推論

CTranslate2 の、とりわけ SentenceTransformer 互換モデルとしての利用方法はとても簡単で、例えば SentenceTransformer の以下のコード

from sentence_transformers import SentenceTransformer
model = SentenceTransformer(model_name, device=device)
embs = model.encode(texts)

を、以下へと変更するだけ。

from hf_hub_ctranslate2 import CT2SentenceTransformer
model = CT2SentenceTransformer(
    model_name, device=device, compute_type=compute_type
)
embs = model.encode(texts)

これだけで、基本1.6~2倍速で推論できるようになる。メモリも省メモリ。

また CT2SentenceTransformer は SentenceTransformer の子クラスとして実装してあるので、SentenceTransformerとほぼ同じ形で利用ができる。compute_type は後述する。

実際の推論速度と、推論結果の差異

では、どれぐらい速度に差があるのか、実測値を見てみよう。データは wikipedia 日本語の2万サンプルの先頭テキスト512トークンをmultilingual-e5-smallで類似検索タスク("query :" を prefix につける)としてembeddingsに変換したものを、元々のSentenceTransformerでの推論時と、CT2 + 各種 compute_type の推論時で比較した。ノートブックはこちら。GPUはRTX4090, CPUはRyzen9 5950。speed は SentenceTransformer を 1.0 としたときの相対比較。

今回の結果では、CT2 + int8などの量子化はむしろ速度が悪くなっていて、CT2 + float16が速度面では一番速い。GPUではcompute_type="float16" で1.66倍速、推論結果を評価する指標として mAP@100 を見てみたところの結果は1.0、つまり順位変動無しだったので、もうちょっと細かい精度の差を見るために MSE で差を計測するも、MSE もほぼ変わらない(0.0表記だが、実際は3e-09)結果であった。なお、GPU時の compute_type = "auto" は、どうやら "int8_float" のようである。

また CPU も compute_type = "auto" 時に約1.9倍速、MAP@100は引き続き1.0、MSEは0.000004とほんのわずかな差が観測された程度で、実運用ではほぼ問題ないぐらいの数値であろう。推論はCPUに、というケースも多いだろうから、CPU推論で1.9倍速になる恩恵は結構大きいと思っている。また今回は計測していないが、省メモリもうたってるし実際省メモリだったと思うので、計算機リソースが厳しい環境では嬉しさがさらに高まりそうだ。

もっと評価されて良い CTranslate2

CTranslate2 は、エンコーダーモデルにも使え、hf_hub_ctranslate2 を使うことで HuggingFace のモデルも簡単に利用できる。今回は SentenceTransformer の代わりに利用してみたが、未だBert系列のモデルもさまざまなタスクで使うことも多いと思うので、利用できる範囲が広いと感じる。

しかしながら、CTranslate2 は GitHub のスター数が現在 2.2k (llama.cpp は 44.6k !)と、昨今のLLMブームでやたらスターがつくプロジェクトも多い中、控えめな人気。CTranslate2という名前(初期は機械翻訳モデルの高速化に使っていたため?)から、どんなことができるライブラリか察せられず勿体無い感もある。ので、みんなどんどん使ってみて欲しい。