WEDでML/DSエンジニアをしている園芸係です。こちらは家で育てているシダとなります。
前回までは、レシートからLLMを使って住所や電話番号・商品名やその値段を抽出する話、そのシステムを構築した話でした。
課題
弊社のレシート買取サービスONEでは、VertexAIのOnline predictionにLLMのモデルを用いてレシート情報を抽出しています。3ヶ月ほどLLMの運用と調整も行ってきましたが、より高速に推論できれば必要なGPUの数も減らせる、結果的にコスト削減できる、との観点から高速化についても検討を始めました。やはり、GPUインスタンスを使った大規模な運用ではどうしてもコストが大きくかかってしまうのが悩み所でした。
検討手法
推論の高速化にも色々な手法があります。
- 蒸留
- 元の大きいモデルを教師モデルとして、その知識を同機能の小さいモデルの学習に利用
- 変換難易度は高い
- 枝刈り
- 不要なネットワークを削除
- 変換難易度は高い
- 低精度化、量子化
- 浮動小数点float32の計算をfloat16やintで行う
- 変換難易度は低い
まず、蒸留についてです。初期ではT5-baseを学習して使っていたのですが、試しにT5-smallを同じ教師データで学習した結果、同等の精度に収束してしまいました。最初から小さいモデルを使えばそもそも蒸留はしなくてもいいかな、ということでそのようにしています。
次に、枝刈りです。不要なネットワークを削除し、計算量を減らすことができますが、一方で、ネットネットワークの形が変わってしまい、その他の高速化手法が使いにくくなるな、と考え今回は行っていません。
低精度化についてです。通常だと浮動小数点の計算は、通常だとfloat32で行われていますが、NVIDIAより提供されているTensorRTというSDKを利用することで、float16で計算するように変換しました。
なお、TensorRTは、float16やint8の低精度化だけでなく、効率良く計算が行われるように、複数の計算をまとめたり、不要な計算の削除をしたりなどの計算グラフ全体の最適化も行ってくれます。
モデル性能評価
まず、推論が最も高速に行えるよう、バッチ数固定、推論が終わればすぐに次の推論を行うという、理想的な状態で性能を調べました。
T5-baseからT5-smallへ、同等の小さいモデルへ変えることで、1枚あたりの推論時間は0.742秒から0.628秒へと、0.846倍に短縮しました。
さらに、TesorRTで最適化すると0.131秒、つまり最終的に0.176倍へと大きく短縮することができました！！！
T5-baseからT5-smallへ、float32からfloat16へ、変更することにより、推論精度の低下を心配していたのですが、トークン数が長くなった時にわずかに推論精度が下がる以外は、推論精度は劣化せず、自身も驚きました。
運用評価
モデルの性能は問題ないということで、実際にデプロイしました。その結果は、というと・・・
TensorRT高速化モデルのデプロイ前のレイテンシは次のようでしたが、
デプロイ後は、次のようにレイテンシは、半分ほどに改善しています。
また、コストもこのように大きく下がったことが確認できました。
推論の量が多い弊社では、とても効果のある施策でした。
運用上の注意
TensorRTによる高速化の1つ目の注意点として、最適化したエンジンファイル（モデルファイル）は、そのビルド時と同じ環境、GPUでないと動作しない、という制限があります。
つまり、GPU環境がないCloud Build上などで、エンジンファイルは作れないのです。
推論時と同じハードウェアの構成を準備し、さらに同じTensorRTのバージョンをインストールし、その上でビルドしエンジンファイルを作る必要があります。
（ちょっとでもバージョンが食い違うと、segmentation fault (core dumped)が表示され無情にも落ちてしまいました・・・）
2つ目の注意点として、TensorRTは内部に大量のバイナリファイルを含んでおり、これをインストールしたDockerイメージは膨大な大きさになってしまう、という点があります。
最初、Dockerイメージが30GB近くなってしまい、VertexAIへのデプロイ、スケーリングができない状態になってしまいました。その原因は、上のリンクのように、VertexAIの仕様、Dockerイメージは10-15Gくらいにする必要、とのことでした。ということで、次のようなマルチ・ステージング・ビルドを使うことで、中間生成物を含まない、binとライブラリだけを入れることで、なんとか10GB以下のDockerイメージを作りました。
FROM nvidia/cuda:12.5.1-runtime-ubuntu22.04 as builder
RUN pip3 install tensorrt
FROM nvidia/cuda:12.5.1-runtime-ubuntu22.04 as runner
COPY --from=builder /usr/local/lib/python3.10/dist-packages /usr/local/lib/python3.10/dist-packages
COPY --from=builder /usr/local/bin /usr/local/bin
TensorRTによる高速化は、ちょっとのバージョンの差異でも動かなくなってしまい、運用が難しくなる欠点を体感しましたが、一方で、その欠点を上回る速度向上、コスト削減を体感することもできました。
さらに、この高速化は、多くのモデルに対して適応でき、1日の推論量が多い、推論のコストが高いのであれば、ぜひ行っていく方がいいように思います。
Discussion