NVIDIA TensorRT-LLM で大規模言語モデルの推論を最適化

NVIDIA は、NVIDIA GPU 上の最新の LLMの推論性能を高速化および最適化する TensorRT-LLM の一般提供を発表しました。このオープンソース ライブラリは、現在、/NVIDIA/TensorRT-LLM GitHub レポジトリおよび NVIDIA NeMo フレームワークの一部として無料で提供されています。

大規模言語モデル (LLM) は人工知能の分野に革命をもたらし、デジタルの世界と対話するまったく新しい方法を生み出しました。しかし、世界中の組織やアプリケーション開発者が LLM を業務に取り入れようとする中で、これらのモデルを実行する上でのいくつかの課題が明らかになってきています。

簡単に言えば、LLM は巨大です。そのため、適切な技術がなければ、LLM の実行が遅くなり、運用費用がかかる可能性があります。

カーネル融合や量子化のようなモデル最適化から、C++ 実装、KV キャッシュ、インフライト バッチング、Paged Attention などのランタイム最適化まで、多くの最適化技術がこの問題に対処するために台頭してきました。どれが自分のユース ケースに適しているかを判断し、これらのテクニックと時には互換性のない実装との相互作用をナビゲートするのは難しいかもしれません。

そのため NVIDIA は、推論のために LLM をコンパイルし最適化するための包括的なライブラリである TensorRT-LLM を発表しました。TensorRT-LLM は、新しいモデルを定義して構築するための直感的な Python API を提供しながら、これらの最適化すべてを組み込んでいます。

TensorRT-LLM オープンソース ライブラリは、NVIDIA GPU 上の最新の LLM の推論パフォーマンスを加速します。このライブラリは、生成 AI アプリケーションを構築、カスタマイズし、実運用のためにデプロイする際の最適化バックボーンとして、エンドツーエンドのフレームワークである NVIDIA NeMo のLLM 推論機能を 使用しています。NeMo は、TensorRT-LLM と NVIDIA Triton を含む、生成 AI デプロイ向けの完全なコンテナーを提供します。

TensorRT-LLM は現在、ベータ リリースとしてネイティブ Windows でも利用可能です。アプリケーション開発者や AI に熱心な方は、NVIDIA RTX および NVIDIA GeForce RTX GPU を搭載した PC やワークステーション上でローカルに実行される高速化された LLM からメリットを得ることができます。

TensorRT-LLM は、TensorRT のディープラーニング コンパイラを備え、LLM 実行のための最先端の実装である FlashAttention と Masked Multi-head Attention (MHA) ののために作られた最新の最適化されたカーネルを含んでいます。

また、TensorRT-LLM は、前処理と後処理のステップ、そしてマルチ GPU/マルチノード通信プリミティブを備えた、シンプルでオープンソースの Python API から成り立っています。これにより、GPU 上で画期的な LLM 推論のパフォーマンスが実現されます。

TensorRT-LLM のハイライトは以下の通りです。

• Llama 1 および 2、ChatGLM、Falcon、MPT、Baichuan、Starcoder などの LLM をサポート
• インフライト バッチングと Paged Attention
• マルチ GPU マルチノード (MGMN: Multi-GPU multi-node) 推論
• FP8 搭載 NVIDIA Hopper Transformer Engine
• NVIDIA Ampere アーキテクチャ、NVIDIA Ada Lovelace アーキテクチャ、および NVIDIA Hopper GPU のサポート
• ネイティブ Windows サポート (ベータ)

過去 2 年間、NVIDIA は、Anyscale、Baichuan、Cohere、Deli、Grammarly、Meta、Mistral AI、MosaicML (現在は Databricks の一部)、OctoML、Perplexity AI、Tabnine、Together.ai、Zhipu など、LLM を牽引する企業と緊密に協力し、LLM 推論の高速化と最適化に取り組んできました。

ライブラリとその使い方の感覚をつかむために、一般に公開されている LLM である Llama 2 を、Linux 上で TensorRT-LLM と NVIDIA Triton を使用してデプロイする方法の例をご紹介します。ベータリ リースを使い始めるには、TensorRT-LLM for native Windows GitHub リポジトリをご覧ください。

異なるモデル、異なる最適化、マルチ GPU 実行などを含む、より詳細な情報については、TensorRT-LLM サンプルの一覧を参照してください。

インストールの開始

TensorRT-LLM ライブラリのクローンとビルドから始めます。TensorRT-LLM をビルドし、すべての依存関係を取得する最も簡単な方法は、同梱の Dockerfile を使用することです。

git lfs install
git clone -b release/0.5.0 https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git submodule update --init --recursive
make -C docker release_build

これらのコマンドは、ベース コンテナーを取得し、TensorRT-LLM に必要なすべての依存関係をコンテナー内にインストールします。その後、TensorRT-LLM 自体をビルドしてコンテナーにインストールします。

モデルの重みの取得

TensorRT-LLM は、LLM 推論のためのライブラリであるため、これを使用するには、学習済みの重みのセットを提供する必要があります。NVIDIA NeMo のようなフレームワークで学習した独自のモデルの重みを使用するか、HuggingFace Hub のようなリポジトリから学習済みの重みのセットを引っぱって来ることができます。

この記事におけるコマンドは、HuggingFace Hub から 7B パラメーターの Llama 2 モデルのチャットチューニングされたバリアント用の重みとトークナイザー ファイルを自動的に取り出します。また、以下のコマンドを使えば、重みを自分でダウンロードしてオフラインで使うこともできます。このディレクトリを指すように、後のコマンドのパスを更新するだけです。

git lfs install
git clone https://huggingface.co/meta-llama/Llama-2-7b-chat-hf

このモデルの使用には特定のライセンスが必要です。必要なファイルをダウンロードするには、条件に同意し、Hugging Face で認証してください。

モデルのコンパイル

プロセスの次のステップは、モデルを TensorRT エンジンにコンパイルすることです。そのためには、モデルの重みと TensorRT-LLM Python API で書かれたモデル定義が必要です。

TensorRT-LLM リポジトリには、定義済みの様々なモデル アーキテクチャが含まれています。この記事では、自分で書く代わりに、含まれている Llama モデル定義を使用します。これは、TensorRT-LLM で利用可能な最適化の最小限の例です。

利用可能なプラグインと量子化についての詳細は、Llama の完全な例と Numerical Precision を参照してください。

# Launch the Tensorrt-LLM container
make -C docker release_run LOCAL_USER=1
 
# Log in to huggingface-cli
# You can get your token from huggingface.co/settings/token
huggingface-cli login --token *****
 
# Compile model
python3 examples/llama/build.py \
    --model_dir meta-llama/Llama-2-7b-chat-hf \
    --dtype float16 \
    --use_gpt_attention_plugin float16 \
    --use_gemm_plugin float16 \
    --remove_input_padding \
    --use_inflight_batching \
    --paged_kv_cache \
    --output_dir examples/llama/out

TensorRT-LLM API でモデル定義を作成すると、ニューラル ネットワークの層を形成する NVIDIA TensorRT プリミティブから操作のグラフが構築されます。これらの操作は、特定のカーネル、つまり GPU 用にあらかじめ書き込まれたプログラムにマッピングされます。

TensorRT コンパイラは、グラフを走査して、各処理と利用可能な GPU に最適なカーネルを選択することができます。重要なのは、複数の演算を 1 つのカーネルに統合するのに適した候補となるグラフのパターンを特定できることです。これにより、必要なメモリ移動量と、複数の GPU カーネルを起動のオーバーヘッドが削減されます。

TensorRT はまた、操作のグラフを単一の CUDA グラフにコンパイルし、1 度にすべてを起動できるようにすることで、カーネル起動のオーバーヘッドをさらに削減します。

TensorRT コンパイラは、レイヤーを融合して実行速度を向上させるのに非常に強力ですが、FlashAttention のように、多くの操作を相互にインターリーブする必要があり、自動的に検出できない複雑なレイヤーの融合もあります。そのような場合は、コンパイル時にグラフの一部を明示的にプラグインに置き換えることができます。

この例では、FlashAttention のような fused attention kernel を実装する gpt_attention プラグインと、FP32 アキュムレーションで行列乗算を実行する gemm プラグインが含まれています。また、HuggingFace からダウンロードした重みのデフォルトの精度と一致するように、フルモデルに必要な精度を FP16 として呼び出します。

このスクリプトを実行し終わると、以下のようになります。examples/llama/out フォルダーには、以下のファイルがあります。

• Llama_float16_tp1_rank0.engine: ビルド スクリプトのメイン出力で、モデルの重みを埋め込んだ実行可能な操作グラフが含まれています。
• config.json: 一般的な構造や精度などのモデルの詳細情報や、どのプラグインがエンジンに組み込まれているかなどの情報が含まれています。
• model.cache: モデル コンパイルで得られたタイミングと最適化情報の一部をキャッシュし、連続するビルドを高速化します。

モデルの実行

それでは、モデル エンジンを手に入れたところで、それで何ができるのでしょうか?

エンジン ファイルにはモデルを実行するために必要な情報が含まれていますが、LLM を実際に使用するには、モデルを 1 回フォワード パスするだけでは不十分です。TensorRT-LLM には、ビルドされた LLM エンジンを実行し、モデル出力からのトークンのサンプリング、KV キャッシュの管理、リクエストのバッチ処理などのプロセスを管理するために、高度に最適化された C++ ランタイムが含まれています。

このランタイムを直接使用してローカルでモデルを実行することもできますし、NVIDIA Triton Inference Server 用の TensorRT-LLM ランタイム バックエンドを使用して、複数のユーザーにモデルを提供することもできます。

モデルをローカルで実行するには、以下のコマンドを実行します。

python3 examples/llama/run.py --engine_dir=examples/llama/out --max_output_len 100 --tokenizer_dir meta-llama/Llama-2-7b-chat-hf --input_text "How do I count to nine in French?"

Triton Inference Server でデプロイ

ローカルで実行するだけでなく、NVIDIA Triton Inference Server を使用して、LLM の本番環境でのデプロイを作成することもできます。

NVIDIA は、新しい TensorRT-LLM 用の Triton Inference Server バックエンドをリリースします。このバックエンドは、TensorRT-LLM C++ ランタイムを活用して推論を高速に実行し、インフライト バッチングや Paged KV Cache などの技術を含みます。TensorRT-LLM バックエンドを備えた Triton Inference Server は、NGC を通じてビルド済みコンテナーとして利用することができます。

まず、Triton Inference Server がモデルと関連するメタ データを読み込めるように、モデル リポジトリを作成します。tensorrtllm_backend リポジトリには、 all_models/inflight_batcher_llm/ 配下に適切なモデル リポジトリのスケルトンが含まれています。

• /preprocessing および /postprocessing: テキスト入力をトークン化し、モデル出力をデトークン化して、文字列とモデルを操作するトークン ID を変換するための、Triton Inference Server Python バックエンドのスクリプトが含まれています。
• /tensorrt_llm: 以前コンパイルしたモデル エンジンを配置する場所。
• /ensemble: 前の 3 つのコンポーネントを一緒にリンクするモデル アンサンブルを定義し、Triton Inference Server にデータをどのように流すかを指示します。

サンプル モデル リポジトリをプルダウンし、前のステップでコンパイルしたモデルをそこにコピーします。

# After exiting the TensorRT-LLM docker container
cd ..
git clone -b release/0.5.0 \ 
https://github.com/triton-inference-server/tensorrtllm_backend.git
cd tensorrtllm_backend
cp ../TensorRT-LLM/examples/llama/out/*   all_models/inflight_batcher_llm/tensorrt_llm/1/

次に、リポジトリのスケルトンにある設定ファイルのいくつかを、以下のような情報で修正します。

• コンパイルされたモデル エンジンの場所
• どのトークナイザーを使用するか
• 推論をバッチで実行する際の KV キャッシュのメモリ割り当ての処理方法

python3 tools/fill_template.py --in_place \
      all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt \
      decoupled_mode:true,engine_dir:/all_models/inflight_batcher_llm/tensorrt_llm/1,\
max_tokens_in_paged_kv_cache:,batch_scheduler_policy:guaranteed_completion,kv_cache_free_gpu_mem_fraction:0.2,\
max_num_sequences:4
 
python tools/fill_template.py --in_place \
    all_models/inflight_batcher_llm/preprocessing/config.pbtxt \
    tokenizer_type:llama,tokenizer_dir:meta-llama/Llama-2-7b-chat-hf
 
python tools/fill_template.py --in_place \
    all_models/inflight_batcher_llm/postprocessing/config.pbtxt \
    tokenizer_type:llama,tokenizer_dir:meta-llama/Llama-2-7b-chat-hf

Docker コンテナーを起動し、Triton サーバーを起動します。モデルが構築された GPU の数である「world_size」を指定し、先ほど設定した model_repo を指します。

docker run -it --rm --gpus all --network host --shm-size=1g \
-v $(pwd)/all_models:/all_models \
-v $(pwd)/scripts:/opt/scripts \
nvcr.io/nvidia/tritonserver:23.10-trtllm-python-py3
 
# Log in to huggingface-cli to get tokenizer
huggingface-cli login --token *****
 
# Install python dependencies
pip install sentencepiece protobuf
 
# Launch Server
python /opt/scripts/launch_triton_server.py --model_repo /all_models/inflight_batcher_llm --world_size 1

リクエストの送信

実行中のサーバーにリクエストを送信してやりとりするには、Triton Inference Server クライアント ライブラリのいずれかを使用するか、エンドポイント作成に HTTP リクエストを送信します。まずは、より高機能なクライアント スクリプトか、以下の curl コマンドを使用してください。

curl -X POST localhost:8000/v2/models/ensemble/generate -d \
'{
"text_input": "How do I count to nine in French?",
"parameters": {
"max_tokens": 100,
"bad_words":[""],
"stop_words":[""]
}
}'

まとめ

TensorRT-LLM と Triton Inference Server を組み合わせることで、LLM を効率的に最適化、デプロイ、実行するための不可欠なツールキットが提供されます。TensorRT-LLM が GitHub でオープンソース ライブラリとしてリリースされたことで、組織やアプリケーション開発者は、これらのモデルの可能性を活用することがこれまで以上に容易になりました。

LLM の世界に飛び込むのをご希望でしたら、今こそ TensorRT-LLM を使い始める時です。TensorRT-LLM の機能を探求し、さまざまなモデルや最適化を試し、AI による言語モデルの驚異的なパワーを解き放つ旅に出ましょう。私たちは、皆さんが開発する素晴らしいものを目にする日が待ちきれません。

TensorRT-LLM を使い始めるのに役立つ詳細情報については、以下を参照してください。

• NVIDIA/TensorRT-LLM GitHub レポジトリにあるオープンソース ライブラリにアクセスしてください。
• 生成 AI デプロイ用の完全なコンテナー (TensorRT-LLM と NVIDIA Triton を含む) を提供している NVIDIA NeMo で詳細をご覧ください。
• サンプル コード、ベンチマーク、TensorRT-LLM のドキュメントは、GitHub で提供しています。
• エンタープライズグレードのセキュリティ、安定性、管理性、そしてサポートを備えたミッションクリティカルな AI 推論を実現するためには、TensorRT を含むエンドツーエンドの AI ソフトウェアプラットフォームである NVIDIA AI Enterprise をご購入ください。TensorRT-LLM もまもなく含まれます。詳細については、営業チームまでお問い合わせください。
• TensorRT と TensorRT-LLM を使い始めるには、こちらの関連情報をご覧ください。

翻訳に関する免責事項

この記事は、「Optimizing Inference on Large Language Models with NVIDIA TensorRT-LLM, Now Publicly Available」の抄訳で、お客様の利便性のために機械翻訳によって翻訳されたものです。NVIDIA では、翻訳の正確さを期すために注意を払っておりますが、翻訳の正確性については保証いたしません。翻訳された記事の内容の正確性に関して疑問が生じた場合は、原典である英語の記事を参照してください。