Amazon SagemakerがLarge Model Inference (LMI) DLCの新モデルにてTensorRT-LLMをサポートしました。 #AWSreinvent

AWS re:Invent 2023

#AWS

せーの

2023.12.01

せーのでございます。

Amazon Sagemakerの大規模モデル推論（LMI）ディープラーニングコンテナ（DLC）の新バージョン（0.25.0）が発表されましたのでご紹介します。

LMI DLCとは

SageMakerには、AWSインフラストラクチャ上でGPT、T5、OPT、BLOOM、Stable Diffusionなどの大規模モデルをホストするための一般的なオープンソースライブラリを備えたディープラーニングコンテナ（DLC）があります。
DLCを使うと、DLCに対応したサードパーティライブラリを使って、モデル並列化技術を使用してモデルパラメータを分割し、推論に複数のGPUのメモリを活用することができます。

サードパーティライブラリにはDeepSpeed Inference、Hugging Face Accelerate、FasterTransformerなどがありましたが、今回これにTensorRT-LLMが加わりました。

Amazon SageMaker LMI TensorRT-LLM DLCは、旧バージョンと比較して、Llama2-70B、Falcon-40B、CodeLlama-34Bモデルのレイテンシを平均33％削減し、スループットを平均60％向上させます。

TensorRT-LLM ライブラリ

ではこのTensorRT-LLM ライブラリはどんな機能があるでしょう。

最新のLMI DLCは、スループットを改善するための推論リクエストの継続的なバッチ処理のサポート、レイテンシを改善するための効率的な推論集団操作、GPUでのパフォーマンスを最大化するためのNVIDIAの最新のTensorRT-LLMライブラリを提供します。

LMI TensorRT-LLM DLCは、モデルIDとオプションのモデルパラメータを入力するだけで、TensorRT-LLMのコンパイルを簡素化するローコードインターフェイスを提供します。また、GPTQ、AWQ、SmoothQuantといった最新の量子化技術もLMI DLCで利用できます。

TensorRT-LLM ライブラリはNVIDIA GPUを使用する際に、SmoothQuant、FP8、大規模言語モデルの連続バッチ処理など、最先端の最適化を可能にします。

TensorRTはよくJetsonなどのエッジコンピュータにて機械学習モデルを動かしたいときに、限られたGPUリソースでも快適に動くように量子化技術などを使ってモデルを軽くするNVIDIAの目玉技術ですね。今回はこの技術をLLMに応用して、推論の効率を上げているものが、AWSのコンテナや並列化などのインフラ技術と手を組んだ、というなかなか驚くニュースです。

専用ツールキット

推論を高速化するためにTensorRT-LLMライブラリを最適化し、ジャスト・イン・タイムのモデル変換をサポートするツールキットもできました。
このツールキットを使用して、Hugging FaceモデルIDを提供し、モデルをエンドツーエンドでデプロイすることができます。このツールキットは、ストリーミングによる連続バッチ処理もサポートしています。Llama-2 7Bと13Bモデルは約1-2分、70Bモデルは約7分でコンパイルできます。

また、SageMakerのエンドポイントを設定し、インスタンスをスケールする際のコンパイルのオーバーヘッドを回避するために、Ahead of Time (AOT)コンパイルを使用することができます。

Triton Server用に構築されたTensorRT LLMモデルであれば、LMI DLCで使用することができます。これにより、例えばLlama-2 70BなどをTensorRT LLM　LMIコンテナを使って効率的に学習させたりすることができます。

最新のLMI DLCは全リージョンで利用可能です。

サンプルのノートブックも用意されているので、こちらを参考にまずは触ってみてください。

https://github.com/aws/amazon-sagemaker-examples/blob/main/inference/generativeai/llm-workshop/deploy-V7-lmi/llama2_70b-lmi-trtllm.ipynb