「Amazon SageMakerを使用して、Llama 2モデルのスループット性能を向上させる」

Improving throughput performance of Llama 2 model using Amazon SageMaker.

機械学習（ML）の広範な採用において、私たちは興奮の点に立っており、generative AIによってほとんどの顧客体験やアプリケーションが再構築されると信じています。Generative AIは、会話、ストーリー、画像、ビデオ、音楽など、新しいコンテンツやアイデアを作成することができます。Generative AIもMLモデルによって動作しますが、これらは非常に大規模なモデルであり、膨大な量のデータでトレーニングされ、ファウンデーションモデル（FM）と呼ばれています。FMはtransformerに基づいています。transformerは、モデルの大きさのため、長いテキストシーケンスの生成において遅く、メモリを多く消費します。テキストシーケンスを生成するために使用される大規模な言語モデル（LLM）は、膨大な計算能力が必要であり、利用可能な高帯域幅メモリ（HBM）と計算能力にアクセスするのが困難です。これは、モデルのパラメータの読み込みと自己回帰デコーディングプロセスによって利用可能なメモリ帯域幅の大部分が消費されるためです。その結果、巨大な計算能力を持っていても、LLMはメモリI/Oと計算の制限によって制約され、利用可能なハードウェアリソースを最大限に活用することができません。

全体として、LLMのgenerative inferenceには3つの主な課題（Pope et al. 2022による）があります：

• 巨大なモデルパラメータとデコーディング中の一時状態による大きなメモリフットプリント。パラメータはしばしば単一のアクセラレータチップのメモリを超えます。アテンションキーバリューキャッシュも大量のメモリを必要とします。
• 低い並列性は、特に大きなメモリフットプリントの場合に遅延を増加させ、パラメータとキャッシュを各ステップでコンピュートコアにロードするために大量のデータ転送を必要とします。これにより、遅延ターゲットを達成するために高い総メモリ帯域幅が必要となります。
• シーケンスの長さに対するアテンションメカニズムの計算の二次的なスケーリングは、遅延と計算の課題を複合化します。

これらの課題に対処するためのテクニックの1つに、バッチ処理があります。バッチ処理は、複数の入力シーケンスをLLMに一緒に送信し、LLMの推論のパフォーマンスを最適化するプロセスを指します。このアプローチにより、モデルパラメータは各入力シーケンスのためにロードする必要がなくなります。パラメータは1回ロードされ、複数の入力シーケンスを処理するために使用されます。バッチ処理はアクセラレータのHBM帯域幅を効率的に利用し、コンピュートの利用率を向上させ、スループットを向上させ、費用対効果の高い推論を実現します。

この記事では、LLMの並列化されたgenerative inferenceにおいてバッチ処理技術を最大限に活用するためのテクニックを調査します。メモリフットプリントを削減し、並列性を向上させ、アテンションの二次的なスケーリングを緩和するためのさまざまなバッチ処理方法について説明します。目標は、HBMやアクセラレータなどのハードウェアを完全に活用して、メモリ、I/O、および計算のボトルネックを克服することです。そして、Amazon SageMaker large model inference（LMI）ディープラーニングコンテナ（DLC）がこれらのテクニックにどのように役立つかについても紹介します。最後に、LMI DLCを使用したSageMakerでの各バッチング戦略のスループット改善の比較分析を提供します。Llama v2などのモデルのスループットを向上させるためのサンプルノートブックは、SageMakerの例のGitHubリポジトリで入手できます。

• 「包括的な革新：Amazon SageMakerでのHack.The.Bias」
• 認知コンピューティング：定義、動作、例など
• テキストから音楽を生成するAI：Stability Audio、GoogleのMusicLMなど

大規模言語モデル（LLMs）の推論

自己回帰デコーディングは、GPTなどの言語モデルがトークンを1つずつ生成してテキストの出力を生成するプロセスです。これには次の手順が含まれます：

1. モデルは、シーケンスの前のトークンを入力として受け取ります。最初のステップでは、これはユーザーが提供した開始プロンプトです。
2. モデルは次のトークンの語彙に対する分布を予測します。
3. 最も高い予測確率を持つトークンが選択され、出力シーケンスに追加されます。ステップ2とステップ3はデコーディングの一部です。現時点では、最も一般的なデコーディング方法は、貪欲探索、ビームサーチ、相対探索、サンプリングです。
4. この新しいトークンは、次のデコーディングステップのための入力シーケンスに追加されます。
5. モデルは、終了シーケンスマーカーが生成されるか、所望の出力長に達するまで、ステップごとに1つの新しいトークンを生成することを繰り返します。

LLMsのモデルサービング

LLMsのモデルサービングは、テキスト生成のための入力リクエストを受け取り、推論を行い、結果を要求元のアプリケーションに返すプロセスを指します。モデルサービングに関与する主要な概念は次のとおりです：

• クライアントは、各リクエストがトークンのシーケンスまたは入力プロンプトのシーケンスで構成される複数の推論リクエストを生成します。
• リクエストは推論サーバー（たとえば、DJLServing、TorchServe、Triton、またはHugging Face TGI）によって受け取られます。
• 推論サーバーは推論リクエストをバッチ処理し、ジェネレーティブ言語モデル上の順方向パス（出力トークンシーケンスの予測）を実行するためのモデルパーティショニングライブラリ（Transformers-NeuronX、DeepSpeed、Accelerate、またはFasterTransformerなど）を含む実行エンジンにバッチをスケジュールします。
• 実行エンジンはレスポンストークンを生成し、応答を推論サーバーに送信します。
• 推論サーバーは生成された結果をクライアントに返信します。

推論サーバーがリクエストレベルで実行エンジンとやり取りする際には、リクエストごとにPythonプロセスを使用するため、モデルの別のコピーが必要で、メモリに制約があります。例えば、以下の図に示すように、96 GBの総アクセラレータデバイスメモリを持つ機械学習（ML）インスタンスに、80 GBのサイズのモデルの単一のコピーをロードすることしかできません。同時に追加のリクエストを処理するためには、モデル全体の追加のコピーをロードする必要があります。これはメモリとコストの面で効率的ではありません。

リクエストレベルのスケジューリングによる課題を理解したので、スループットを最適化するための異なるバッチング技術を見ていきましょう。

バッチング技術

このセクションでは、さまざまなバッチング技術を説明し、SageMaker LMIコンテナを使用してそれらを実装する方法を示します。

推論リクエストのバッチングには、主に2つのタイプがあります：

• クライアント側（静的） – 通常、クライアントがサーバーにリクエストを送信すると、サーバーはデフォルトで各リクエストを順次処理しますが、これはスループットには最適ではありません。スループットを最適化するために、クライアントは単一のペイロード内の推論リクエストをバッチ処理し、サーバーはバッチを複数のリクエストに分割してそれぞれのリクエストを個別に実行する前処理ロジックを実装します。このオプションでは、クライアントはバッチ処理のためにコードを変更する必要があり、ソリューションはバッチサイズと密接に結び付いています。
• サーバー側（動的） – バッチ処理の別のテクニックとして、推論を使用してサーバー側でバッチ処理を行う方法があります。独立した推論リクエストがサーバーに到着すると、推論サーバーはそれらをサーバー側で大きなバッチに動的にグループ化することができます。推論サーバーは指定されたレイテンシターゲットを満たすためにバッチ処理を管理し、所望のレイテンシ範囲内にとどまりながらスループットを最大化します。推論サーバーはこれを自動的に処理するため、クライアント側のコードの変更は必要ありません。サーバー側のバッチ処理には、自己回帰デコーディングに基づく生成言語モデルのスループットをさらに最適化するためのさまざまなテクニックが含まれます。これらのバッチ処理技術には、動的バッチ処理、連続バッチ処理、およびPagedAttention（vLLM）バッチ処理が含まれます。

動的バッチ処理

動的バッチ処理は、入力リクエストを組み合わせてバッチとして送信し、推論に使用します。動的バッチ処理は、コンピュータビジョン（CV）、自然言語処理（NLP）など、すべてのタスクに対して動作する一般的なサーバー側バッチ処理技術です。

LMIコンテナでは、以下の設定に基づいてリクエストのバッチ処理を設定できます（serving.properties）：

• batch_size – バッチのサイズ
• max_batch_delay – バッチの集約の最大遅延時間

これらのいずれかの閾値が満たされる（最大バッチサイズに達するか、待機期間が完了する）と、新しいバッチが準備され、推論用のモデルにプッシュされます。次の図は、異なる入力シーケンスの長さを持つリクエストの動的なバッチ処理をモデルが一緒に処理している様子を示しています。

SageMakerで動的バッチ処理を実装するには、LMIコンテナのserving.propertiesを次のように設定します：

#動的バッチ処理
engine=Python
option.entryPoint=djl_python.huggingface
batch_size=64 #例
max_batch_delay=1000 #例
option.tensor_parallel_degree=2 #例

動的バッチ処理は、バッチ処理を行わない場合と比べてスループットを最大4倍向上させることができますが、GPUの利用率が最適ではないことが観察されます。なぜなら、すべてのリクエストの処理が完了するまで、システムは別のバッチを受け入れることができないからです。

連続バッチ処理

連続バッチ処理は、テキスト生成に特化した最適化手法です。スループットを向上させ、最初のバイトの待ち時間を犠牲にしません。連続バッチ処理（またはイテレーションバッチ処理、ローリングバッチ処理とも呼ばれます）は、アイドルなGPU時間の課題に対処し、動的バッチ処理の手法をさらに改善することで、常に新しいリクエストをバッチに追加し続けます。以下の図は、リクエストの連続バッチ処理を示しています。リクエスト2と3が処理を終えると、新しい一連のリクエストがスケジュールされます。

以下のインタラクティブな図は、連続バッチ処理の仕組みについてさらに詳しく説明しています。

(提供: https://github.com/InternLM/lmdeploy)

LLMやテキスト生成を効率的に行うための強力なテクニックとして、一部のアテンション行列をキャッシュする方法があります。これにより、プロンプトの最初のパスと後続のフォワードパスが異なるということです。最初のパスでは、完全なアテンション行列を計算する必要がありますが、後続の処理では新しいトークンのアテンションのみを計算すれば済みます。このコードでは最初のパスを「prefill」と呼び、後続の処理を「decode」と呼びます。prefillはdecodeよりもはるかに負荷が高いため、常に実行する必要はありませんが、現在実行中のクエリはおそらくdecodeを行っているでしょう。前述のように連続バッチ処理を使用したい場合は、decodeグループに参加するために必要なアテンション行列を作成するために、ある時点でprefillを実行する必要があります。

このテクニックにより、アイドルなGPUを効果的に活用することで、バッチ処理を行わない場合と比べてスループットを最大20倍向上させることができます。

連続バッチ処理の使用に関して、LMIコンテナのserving.propertiesで以下のパラメータを微調整することができます：

• engine – コードのランタイムエンジンです。値にはPython、DeepSpeed、FasterTransformer、MPIなどがあります。連続バッチ処理を有効にするにはMPIを使用します。
• rolling_batch – サポートされている戦略のいずれかを使用して、イテレーションレベルのバッチ処理を有効にします。値にはauto、scheduler、lmi-distなどがあります。Llama 2の連続バッチ処理をオンにするためにはlmi-distを使用します。
• max_rolling_batch_size – 連続バッチ内の同時リクエストの最大数を制限します。デフォルトは32です。
• max_rolling_batch_prefill_tokens – キャッシュするトークンの数を制限します。これはバッチサイズと入力シーケンスの長さに基づいてチューニングする必要があります。GPUのメモリ不足を避けるためです。これはrolling_batch=lmi-distの場合にのみサポートされます。私たちの推奨値は、同時リクエストの数×リクエストごとに必要な入力トークンと出力トークンのメモリを基に設定することです。

以下は、連続バッチ処理を設定するためのserving.propertiesのサンプルコードです：

#Continuous Batching
engine=MPI
option.entryPoint=djl_python.huggingface
option.rolling_batch=auto
option.max_rolling_batch_size=64 #例
option.paged_attention=false
option.max_rolling_batch_prefill_tokens=16080 #例
option.tensor_parallel_degree=2 #例

PagedAttentionバッチ処理

自己回帰デコーディングプロセスでは、LLMにおけるすべての入力トークンがアテンションキーと値のテンソルを生成し、これらのテンソルは次のトークンを生成するためにGPUメモリに保持されます。これらのキャッシュされたキーと値のテンソルは、KVキャッシュまたはアテンションキャッシュと呼ばれることがあります。論文「vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention」によれば、Llama 13Bの単一のシーケンスでは、KVキャッシュは最大1.7GBのメモリを使用します。また、そのサイズはシーケンスの長さに依存し、非常に可変性と予測不可能性が高いです。そのため、KVキャッシュを効率的に管理することは大きな課題となっています。論文では、既存のシステムがフラグメンテーションと過剰予約によりメモリを60-80%も浪費していることがわかりました。

PagedAttentionは、UC Berkeleyが開発した新しい最適化アルゴリズムであり、アテンションキャッシュ（KVキャッシュ）を非連続にすることで、メモリを固定サイズのページまたはブロックに割り当てることにより、連続バッチングプロセスを改善します。これは、オペレーティングシステムで使用される仮想メモリとページングの概念に触発されました。

vLLM論文によると、トークンの各シーケンスのアテンションキャッシュはブロックに分割され、ブロックテーブルを介して物理ブロックにマップされます。アテンションの計算中、PagedAttentionカーネルはブロックテーブルを使用して物理メモリからブロックを効率的に取得することができます。これにより、メモリの浪費が大幅に削減され、より大きなバッチサイズ、高いGPU利用率、および高いスループットが可能になります。以下の図は、アテンションキャッシュを非連続のページに分割する様子を示しています。

以下の図は、PagedAttentionを使用した推論の例を示しています。主要なステップは次のとおりです：

1. 入力プロンプトを含む推論リクエストが受け取られます。
2. プリフィルフェーズでは、アテンションが計算され、キーと値が非連続な物理メモリに保存され、論理的なキーと値のブロックにマッピングされます。このマッピングはブロックテーブルに保存されます。
3. 入力プロンプトはモデルを通過し（順方向パス）、最初の応答トークンを生成します。応答トークンの生成中に、プリフィルフェーズのアテンションキャッシュが使用されます。
4. 後続のトークンの生成中、現在の物理ブロックがいっぱいの場合、非連続なメモリを追加で割り当てることができます。これにより、必要な時点での割り当てが可能になります。

PagedAttentionは、最適なメモリ使用とメモリの浪費の削減に役立ちます。これにより、より多くのリクエストをまとめてバッチ処理することができ、推論のスループットが大幅に向上します。

次のコードは、SageMakerのLMIコンテナでPagedAttentionバッチングを設定するためのサンプルのserving.propertiesです：

#Paged Attention Batching
engine=MPI
option.entryPoint=djl_python.huggingface
option.rolling_batch=auto
option.max_rolling_batch_size=64 #example
option.paged_attention=true
option.max_rolling_batch_prefill_tokens=16080 #example
option.tensor_parallel_degree=2 #example

どのバッチング技術を使用するか

以下の図は、LMI on SageMakerのserving.propertiesとともに、サーバーサイドのバッチング技術を要約しています。

以下の表は、異なるバッチング技術とそれらの使用例をまとめています。

	PagedAttentionバッチング	連続バッチング	ダイナミックバッチング	クライアント側バッチング	バッチなし
動作原理	新しいリクエストをトークンレベルでマージし、ページ化されたブロックとバッチ推論を行います。	新しいリクエストをトークンレベルで常にマージし、バッチ推論を行います。	新しいリクエストをリクエストレベルでマージし、バッチ形成のために数ミリ秒遅延することができます。	クライアントは、複数の推論リクエストを同じペイロードでバッチ処理してから推論サーバーに送信する責任があります。	リクエストが到着すると、即座に推論を実行します。
最適な使用場面	これは、サポートされているデコーダーのみのモデルにおすすめのアプローチです。スループット重視のワークロードに適しています。テキスト生成モデルにのみ適用されます。	同じデコーディング戦略で異なるタイミングで発生する並行リクエスト。スループット重視のワークロードに適しています。テキスト生成モデルにのみ適用されます。	同じデコーディング戦略で異なるタイミングで発生する並行リクエスト。応答時間に敏感なワークロードで、より高いスループットが必要な場合に適しています。CV、NLP、およびその他のモデルに適用されます。	レイテンシ制約がなく、スループットを最大化するためのオフライン推論ユースケースに適しています。	頻繁でない推論リクエストまたは異なるデコーディング戦略の推論リクエスト。レスポンスタイムのレイテンシニーズを持つワークロードに適しています。

SageMaker上の大規模生成モデルの異なるバッチング技術のスループット比較

この記事では、Llama v2 7Bモデルを使用してSageMaker上でのパフォーマンスベンチマークを行い、この投稿で説明されている異なるバッチング技術と並行して入力リクエストが50、総リクエスト数が5,000の条件で行いました。

パフォーマンステストでは、変数長の3つの異なる入力プロンプトを使用しました。連続およびPagedAttentionバッチングでは、3つの入力プロンプトに対して出力トークンの長さを64、128、および256に設定しました。ダイナミックバッチングでは、一貫した出力トークンの長さ128トークンを使用しました。テストでは、ml.g5.24xlargeのインスタンスタイプでSageMakerエンドポイントをデプロイしました。以下の表に、パフォーマンスベンチマークテストの結果が示されています。

モデル	バッチング戦略	ml.g5.24xlargeでの1秒あたりのリクエスト数
LLaMA2-7b	ダイナミックバッチング	3.24
LLaMA2-7b	連続バッチング	6.92
LLaMA2-7b	PagedAttentionバッチング	7.41

PagedAttentionバッチングを使用することで、Llama2-7BモデルをSageMaker上でLMIコンテナを使用してダイナミックバッチングと比較して約2.3倍のスループットの向上が見られました。

結論

この投稿では、LLMの推論における異なるバッチング技術と、スループットの向上にどのように役立つかを説明しました。連続およびPagedAttentionバッチングを使用することで、メモリ最適化技術がハードウェアの効率を向上させ、ダイナミックバッチングよりも高いスループット値を提供することを示しました。また、Llama2-7BモデルをSageMaker上でLMIコンテナを使用してダイナミックバッチングと比較して約2.3倍のスループットの向上が見られました。異なるバッチング技術をテストするために使用されたノートブックはGitHubで確認できます。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!