Integrating NVIDIA TensorRT-LLM with the Databricks Inference Stack | Databricks Blogの翻訳です。
本書は抄訳であり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
Databricks / MosaicのR&Dチームは7ヶ月前から推論サービスのアーキテクチャに対する最初のイテレーションをスタートしました。それ以来、急速に成長を続ける生成AIランドスケープにおけるすべての新たな進歩を取り込むことができるスケーラブルでモジュール化され、パフォーマンスのあるプラットフォームを提供するために数多くのストライドを行ってきました。2024年1月には、NVIDIA TensorRT-LLM上に構築された大規模言語モデル(LLM)をサービングするための新たな推論エンジンを使い始めます。
NVIDIA TensorRT-LLMのご紹介
TensorRT-LLMは最先端のLLM推論のためのオープンソースライブラリです。これはいくつかのコンポーネントから構成されます: NVIDIAのTensorRT deep learning compilerを用いたファーストクラスのインテグレーション、言語モデルにおけるキーオペレーションに最適化されたカーネル、効率的なマルチGPUサービングを可能とするコミュニケーションプリミティブです。これらの最適化は、NVIDIAのTensor Core GPUで強化された推論サービスでシームレスに動作し、最先端のパフォーマンスを提供する際のキーパーツとなります。
図1: TensorRT-LLM推論サーバーによって複数クライアントからの推論リクエストが集約されます。推論サーバーでは複雑な多対多の最適化問題を解決しなくてはなりません: 投入されるリクエストはバッチ化されたテンソルを動的にグルーピングされ、これらのテンソルは大量のGPUに分散させる必要があります。
過去6ヶ月において、われわれの推論サービスとTensorRT-LLMをインテグレーションするためのコラボレーションを継続しており、達成したことについて興奮しています。TensorRT-LLMを用いることで、最初のトークンに至る時間と、出力トークンごとの時間の両方において大きな改善を行うことができます。過去の記事で議論したように、これらのメトリクスはLLMを取り扱う際のユーザー体験の品質に対するキーとなる見積もり要素となります。
NVIDIAとのコラボレーションは相互に有益なものです。TensorRT-LLMプロジェクトの最初期フェーズにおいて、我々のチームはMPTモデルの変換スクリプトに貢献し、Hugging FaceからのMPTモデルやMPTアーキテクチャを用いてご自身で事前トレーニング、ファインチューニングしたモデルサービングすることを迅速かつ容易にしました。また、NVIDIAチームはインストール手順を追加し、H100 Tensor Core GPUにおける量子化とFP8のサポートを導入することで、MPTモデルのサポートを拡張しました。このコラボレーションは我々のチームやお客様のだけのメリットを提供するだけでなく、最先端の推論パフォーマンスを用いて、自分たち固有の要件のために自由にMPTモデルを導入することで、より広範なコミュニティを支援できるので、我々はこのようなTensorRT-LLMでMPTアーキテクチャのファーストクラスサポートを提供できることを嬉しく思っています。
プラグインによる柔軟性
最新のモデルアーキテクチャでTensorRT-LLMを拡張することはスムーズなプロセスとなっています。TensorRT-LLMの持つ柔軟性とプラグインを通じて更なる最適化を追加する能力によって、我々のエンジニアは我々固有のモデリング要件をサポートするためにクイックに修正を行うことができます。この柔軟性は、開発プロセスを加速しただけではなく、NVIDIAチームはすべてのユーザー要件に対するサポートを統合したいと言う要件を緩和しました。
LLM推論における重要なコンポーネント
それぞれのクラウドプラットフォームにおいて、すべてのGPUモデル(A10G, A100, H100)に対する包括的なベンチマークを実施しました。最小のコストで最適なレーテンシーを達成するために、以前カバーしたように連続的なバッチサイズ、テンソルのシャーディング、モデルのパイプライン作成のようなTensorRT-LLMの設定を最適化しました。それぞれのクラウドにおいてLLAMA-2やMixtralを含むトップのLLMに対する最適な設定やインスタンス設定をデプロイし、新たなモデルやハードウェアがリリースされるたびに同様のことを行っていきます。Databricksモデルサービングを用いることで常にベストなLLMパフォーマンスを手に入れることができるのです！
インテグレーションを容易にするPython API
ネイティブのインフライト(連続)バッチサポートと組み合わせて活用することで、TensorRT-LLMのオフライン推論パフォーマンスはさらにパワフルなものとなります。我々は、大量のトラフィックがある状況で高いリクエストのスループットを維持するためには、インフライトのバッチ処理が重要なコンポーネントになることを知りました。最近では、NVIDIAチームはC++で記述されたバッチマネージャに対するPythonサポートに取り組んでおり、これによって、TensorRT-LLMはシームレスに我々のバックエンドウェブサーバーと連携できるようになります。
図2. インフライト(連続)バッチ処理の図。Seq 2の長さによって全てのスロットがアイドルになるのを待つのではなく、バッチマネージャは他のスロットで次のシーケンス(Seq 4と5)を処理し始めることができます。(Source: NVIDIA.com)
実験を始める準備はできましたか？
Databricksを利用しているのであれば、すぐにでもAI Playground(現在パブリックプレビュー)を通じて推論サーバーを利用することができます。ログインし、ナビゲーションバーの機械学習の配下にあるPlaygroundをクリックするだけです。
我々は、LLMをホスティングする推論エンジンとしてTensorRT-LLMをインテグレーションというジャーニーを通じて取り組みを行っていますので、非常に重要なコラボレーターとしてNVIDIAチームに感謝の意を表します。 Databricksの推論エンジンにおける今後のリリースにおけるイノベーションの基礎としてTensorRT-LLMを活用する予定です。過去の実装に対する我々のプラットフォームのパフォーマンス改善を共有できることを楽しみにしています。(効率的なLLM推論のための多大なるオープンソースコミュニティであるvLLMが優れた選択肢を提供し、モーメントを生み出していることにも触れさせてください)来月に公開予定のブログ記事では、パフォーマンスの詳細にディープダイブしますので楽しみにしていてください。