🌠

TensorRT-LLM を動かしてみる

かわら

2023/12/03に公開

こんにちは、初めましての方は初めまして。株式会社 Fusic の瓦です。ビールのアドベントカレンダーを買ったので、12 月は毎日ビールが飲めます。冬ってとてもいい季節ですね。

この記事は Fusic Advent Calendar 2023 の三日目の記事となります。昨日は @hkusaba さんの rectorとrector-laravelでLaravelをアップデートしてみるという記事でした。明日は @gorogoroyasu さんの予定です。

昨年の 12 月に ChatGPT が公開されてからこの一年間 LLM に関する話を見ない日はなく、様々な企業が LLM を開発しては公開しています。LLM は便利に扱える一方で、一般的にモデルのサイズがかなり大きくなり、推論に時間がかかってしまうといった問題があります。その問題を解決するために、量子化やキャッシュの使用などの推論を高速化するための様々な手法が提案されてきました。この記事ではそんな高速化を行うためのライブラリである、TensorRT-LLM という LLM に特化した TensorRT を触って試してみたいと思います。

TensorRT-LLM とは

LLM を TensorRT で動かすためのライブラリです。TensorRT のページには

NVIDIA® TensorRT™, an SDK for high-performance deep learning inference, includes a deep learning inference optimizer and runtime that delivers low latency and high throughput for inference applications. ^[1]

と書いてあるように、TensorRT は推論を高速で実行するための SDK です。TensorRT-LLM は LLM でよく使われる計算に特化した TensorRT といったところでしょうか。TensorRT-LLM を使うことで LLM での言語生成が高速化できることが期待できます。

試してみる

試した環境の情報は以下になります。

Ubuntu 22.04 (WSL)
GeForce RTX 4070 Ti
TensorRT-LLM release/0.5.0

環境構築

Docker イメージが用意されているので、概ね公式ドキュメント^[2]に従えば TensorRT-LLM が使用できる環境が出来上がります。ただ私の環境では mpi4py というモジュールのインストールでエラーが出てイメージの構築が出なかったので、同じようなエラーが出るかもしれない方のために何をやったかを書いておきます。

ログを見ると pip 経由で mpi4py のインストールが失敗しているので、まず初めに requirements.txt から mpi4py をコメントアウトします。（requirements-dev.txt というファイルもありますが、開発用イメージでも requirements.txt が使われていそうなので間違えないようにしましょう）

コメントアウトするとイメージの構築が出来るようになります。ただ、mpi4py がないと TensorRT-LLM が動かないため、コンテナを立ち上げた後に手動で入れる必要があります。私は mpi4pyをOpenMPIでインストールした時の話を参考に mpi4py のソースをダウンロードして解凍し、python setup.py build && python setup.py install を実行してインストールしました。

コンテナが立ち上がり中に入ったら TensorRT-LLM のビルド、インストールを行います。これはドキュメント通りにコマンドを打てば大丈夫だと思います。これで TensorRT-LLM を使う準備が出来ました。何回やっても Docker を使った環境構築は楽でいいなと実感します。

実際に試してみる

stabilityai/japanese-stablelm-3b-4e1t-instruct を対象に、TensorRT-LLM を実際に動かしてみます。このモデルは Llama と同様のアーキテクチャとなっているらしい^[3]ので、Llama 用の README を読みながら動くまでを試してみます。

まずはモデルのダウンロードを行います。以下のように、HuggingFace で記載されている通りに重みをコピーします。実行したディレクトリ以下に japanese-stablelm-3b-4e1t-instruct というディレクトリが作成され、その中にモデルの設定や重みが格納されていると思います。

git lfs install
git clone https://huggingface.co/stabilityai/japanese-stablelm-3b-4e1t-instruct

次に TensorRT 用にモデルの変換を行います。GPU が一つしかないので、並列処理などはしない設定で変換します。これは README に書いてあるコマンドをそのまま打てばいいだけです。

python build.py --model_dir /path/to/japanese-stablelm-3b-4e1t-instruct \
                --dtype float16 \
                --remove_input_padding \
                --use_gpt_attention_plugin float16 \
                --enable_context_fmha \
                --use_gemm_plugin float16 \
                --output_dir /path/to/stablelm-3b-instruct-trt-engines

/path/to/japanese-stablelm-3b-4e1t-instruct や /path/to/stablelm-3b-instruct-trt-engines は環境に合わせて適当に変えてください。これで推論にモデルを使用する準備は出来ました。

モデルを使ってテキストを生成させてみます。要約を行うためのコードが提供されていますが、そのままではこのモデルは動かないため一部修正します。トークナイザーを読み込む部分で LlamaTokenizer が指定されていますが、使っているモデルと異なるため、AutoTokenizer をインポートして AutoTokenizer 経由でトークナイザーを読み込むようにします。

-from transformers import AutoModelForCausalLM, LlamaTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer

-    tokenizer = LlamaTokenizer.from_pretrained(hf_model_location,
-                                               legacy=False,
-                                               padding_side='left')
+    tokenizer = AutoTokenizer.from_pretrained(
+        hf_model_location, legacy=False, padding_side="left"
+    )

これでモデルが動くようになると思います。それでは実際に動かしてみましょう。

python summarize.py --test_trt_llm \
                    --hf_model_location /path/to/japanese-stablelm-3b-4e1t-instruct \
                    --data_type fp16 \
                    --engine_dir ./path/to/stablelm-3b-instruct-trt-engines

# > Output
# [12/02/2023-18:04:53] [TRT-LLM] [I]  Article : ...
# [12/02/2023-18:04:53] [TRT-LLM] [I]
#  Highlights : ...
# [12/02/2023-18:04:53] [TRT-LLM] [I]
#  Summary : ...
# ...
# [12/02/2023-18:05:23] [TRT-LLM] [I] TensorRT-LLM (total latency: 27.308032512664795 sec)
# [12/02/2023-18:05:23] [TRT-LLM] [I] TensorRT-LLM beam 0 result
# [12/02/2023-18:05:24] [TRT-LLM] [I]   rouge1 : 0.0
# [12/02/2023-18:05:24] [TRT-LLM] [I]   rouge2 : 0.0
# [12/02/2023-18:05:24] [TRT-LLM] [I]   rougeL : 0.0
# [12/02/2023-18:05:24] [TRT-LLM] [I]   rougeLsum : 0.0

色々省略していますが、出力はこんな感じで出ると思います。元々の HuggingFace のモデルでの性能も試せるので、その出力も見てみます（--test_hf と --test_trt_llm を引数として与えれば両方のモデルの評価を行えるのですが、GPU のメモリが大変なことになったので一回ずつコマンドで試しています）

python summarize.py --test_hf \
                    --hf_model_location /path/to/japanese-stablelm-3b-4e1t-instruct \
                    --data_type fp16

# > Output
# [12/02/2023-20:38:01] [TRT-LLM] [I] Hugging Face (total latency: 228.3670153617859 sec)
# [12/02/2023-20:38:01] [TRT-LLM] [I] HF beam 0 result
# [12/02/2023-20:38:01] [TRT-LLM] [I]   rouge1 : 20.84750521729989
# [12/02/2023-20:38:01] [TRT-LLM] [I]   rouge2 : 5.330592119018671
# [12/02/2023-20:38:01] [TRT-LLM] [I]   rougeL : 15.368175809931033
# [12/02/2023-20:38:01] [TRT-LLM] [I]   rougeLsum : 17.979822200167813

推論時間が 228.37 秒から 56.32 秒へと縮まり、かなり高速に推論が出来ていそうです。一方で ROUGE スコアが 0 になっており、なにかおかしな部分がありそうです…
少し調べたのですが、時間が足りず原因が突き止められなかったので解決編は次回に持ち越したいと思います（ちゃんと動いたら推論時間とか精度の比較もしたいな…）

まとめ

今回は TensorRT-LLM の環境を構築して実際に動かすところまで試してみました。本当は推論時間や精度の比較も行いたかったのですが、ちょっとうまくいっていない部分があるので今回は動かしてみるところまでになりましたが、うまく動いたら次回は複数の GPU を使った場合やより大きなモデルの場合も試してみたいです。

最後に宣伝になりますが、機械学習でビジネスの成長を加速するために、Fusicの機械学習チームがお手伝いたします。機械学習のPoCから運用まで、すべての場面でサポートした実績があります。もし、困っている方がいましたら、ぜひFusicにご相談ください。お問い合わせからでも気軽にご連絡いただけます。またTwitterのDMからでも大歓迎です！

脚注

GitHubで編集を提案

Fusic 技術ブログPublication

さまざまな個性を受け入れて有機的につなぐ社内環境を整える。あらゆる事業機会の創出と実現を繰り返し、世の中に対する視点を絶えず増やして成長していく。あっと驚くような角度から発展できるポイントを見つけ、そこにいい感じにフィットする形でテクノロジーを組み込んで、世の中をちょっとずつ、時には大胆にアップデートしつづけていく。

TensorRT-LLM とは

試してみる

環境構築

実際に試してみる

まとめ

Discussion