LLM推論最適化テクニック

LLM建築を中心に

大規模言語モデル（LLM）はすべて、2017年にVaswaniらによって発明されたトランスフォーマーアーキテクチャに基づいている。トランスフォーマーアーキテクチャは、多様な言語タスクにおいて、優れた精度、少数ショット学習、人間に近い能力を達成している。しかし、このような基礎モデルは、数百から数千億のパラメータから構成されることが多く、学習コストが高く、推論時にリソースを消費する。推論にかかるコストは、入力データが大きいために大きな処理能力を必要とする長い入力コンテキストでエスカレートする。このため、効率的な推論、特にメモリと計算リソースの管理が重要な課題となっている。

具体的には、GPT-3、GPT-4、LLaMA、Mistral、DeepSeekなど、ほとんどの有名なLLMはデコーダのみのLLMである。これらのモデルは、因果モデリングタスクで事前学習され、次単語予測器として機能する。トークンのシーケンスを入力として処理し、停止条件に達するまで自己回帰的に次のトークンを生成する。

デコーダのみのモデルにおけるLLM推論には、プレフィルフェーズとデコードフェーズの2つの重要なフェーズがある。プレフィルフェーズでは、モデルは入力トークンを処理して、最初の新しいトークンを生成するための中間状態（キーと値）を計算する。このフェーズは行列と行列の演算に似ており、高度に並列化され、GPUの能力を効率的に利用します。逆に、デコード・フェーズでは、前のトークンの状態を頼りに、トークンを1つずつ生成する。この行列-ベクトル演算は、計算速度よりもGPUへのデータ転送が主に待ち時間を決定するため、メモリに縛られ、GPUの計算能力が十分に活用されません。

デコード段階を最適化することは、推論の課題を解決するための焦点である。解決策としては、効率的なアテンション・メカニズムの開発や、メモリ・ボトルネックを減らすためのキーと値のより良い管理などがある。この記事では、読者がトランスフォーマー・アーキテクチャーとアテンション・メカニズムを基本的に理解していることを前提に、推論性能を向上させるための実践的なアプローチを紹介する。これらの最適化は、実際のLLM展開においてスループットを向上させ、レイテンシーを削減するために極めて重要である。

さらに、LLM間でトークナイザーが異なるため、トークンの比較可能性に影響するという複雑な問題が生じる。トークンは英語の4文字にほぼ等しいが、トークナイザーによって表現が異なるため、推論のスループット（1秒あたりのトークン数など）を直接比較すると誤解を招く。このばらつきは、推論中のLLMの性能を正確に評価・比較するための標準化された評価指標の必要性を強調している。

バッチ処理

バッチ処理は、大規模言語モデル（LLM）のGPU利用率とスループットを向上させるための重要な戦略です。同じモデルを使用して複数のリクエストを同時に処理することで、バッチングはモデルの重みのメモリ・コストをリクエスト全体に分散し、より大きなバッチでより多くのGPU計算能力を活用できるようにします。しかし、バッチ・サイズには限界があります。過度に大きなバッチは、LLMのメモリ要求、特にキー・バリュー（KV）キャッシュに関連するメモリ・オーバーフローを引き起こす可能性があるからです（これについては後で詳しく説明します）。

従来の静的バッチングでは、バッチ内のリクエストが生成する 完了トークンの数が異なることが多く、実行時間にばらつきが出るため、 制限があった。このため、すべてのリクエストは最も遅いリクエストの完了を待つことになり、世代の長さが大きく異なる場合には問題となる。これに対処するために、パフォーマンスを最適化するインフライトバッチングのような高度な技術が開発されてきた。

インフライトバッチング（連続バッチングとも呼ばれる）は、LLMワークロードの動的な性質がもたらす課題に取り組むもので、単純なチャットボットの応答から複雑なドキュメントの要約やコード生成まで、さまざまなものがある。これらのタスクは、大きく異なるサイズのアウトプットを生成するため、バッチ処理で効率的にリクエストを並列実行することは困難です。静的バッチングとは異なり、インフライトバッチングでは、サーバーは完了したシーケンスを即座にバッチから退避させ、他のリクエストがまだ処理中である間に新しいリクエストの処理を開始することができます。このアプローチは、実世界のシナリオで変化するリクエストの実行時間に適応することで、GPU使用率を大幅に向上させます。

モデル並列化によるマルチGPU展開

モデルの並列化は、大規模な機械学習モデルを複数のGPUに分散させることで、メモリと計算量の要求を管理するための重要な戦略です。このアプローチにより、1つのデバイスのメモリ容量を超えるような大規模なモデルや入力バッチを扱うことができるため、メモリ制約が厳しい場合の学習と推論の両方に不可欠です。モデルの重みを分割する手法には、パイプライン並列、テンソル並列、シーケンス並列などさまざまなものがあり、それぞれモデル分散の異なる側面に対応しています。学習時に大きな入力バッチを処理するためにデバイス間でモデルの重みを複製することに重点を置くデータ並列性とは異なり、これらの手法は、学習時と推論時の両方でメモリフットプリントを削減することに大きく関係しています。

パイプライン並列では、モデルを垂直方向に順次チャンクに分割し、各チャンクには、別々のデバイスに割り当てられたレイヤーのサブセットが含まれる。たとえば、4ウェイ・パイプラインのセットアップでは、各デバイスがモデルのレイヤーの4分の1を処理し、出力を次のデバイスに順番に渡します。これにより、デバイスごとのメモリ要件が大幅に削減される一方で、「パイプライン・バブル」と呼ばれる非効率性が生じます。パイプライン・バブルでは、デバイスが前のレイヤーからの出力を待つ間にアイドル状態になることがあります。入力バッチをより小さなサブバッチに分割して順次処理するマイクロバッチは、このようなバブルを減らすことはできるが、フォワードパスやバックワードパスの間にもアイドル時間が続くため、完全になくすことはできない。

これに対してテンソル並列は、個々のレイヤーを水平方向に分割し、デバイス間で独立して実行できる小さな計算ブロックにする。これはアテンションブロックや多層パーセプトロン（MLP）のようなトランスフォーマーコンポーネントには特に効果的で、例えば異なるアテンションヘッドを別々のデバイスに割り当てて並列計算することができます。しかし、テンソル並列は、LayerNormやDropoutのような、簡単に分割することができず、デバイス間で複製されなければならず、活性化を保存するために冗長なメモリを使用することになる演算にはあまり効果的ではありません。この限界は、メモリ効率を最適化するための補完的アプローチの必要性を浮き彫りにしている。

シーケンス並列性は、シーケンス要素間の独立性を利用して、入力シーケンス次元に沿ってそれらを分割することにより、LayerNormやDropoutのような演算のメモリ非効率性に対処する。この方法は、冗長なアクティベーションのメモリフットプリントを削減し、テンソル並列処理を補完します。これらの並列化技術は相互に排他的ではなく、大規模言語モデル（LLM）をさらに最適化するために組み合わせることができる。さらに、注意モジュールのための特定の最適化戦略は、スケーラビリティを強化し、GPUあたりのメモリ需要を削減し、大規模モデルのより効率的な訓練と推論を可能にすることができます。

アテンション最適化

Vaswaniらによる2017年の論文*Attention Is All You Need*は、自己注意を礎とするTransformerモデルを紹介した。自己注意は、文中の異なる単語の関連性を相対的に評価することを可能にし、自然言語処理のようなタスクの文脈理解を強化する。この論文では、特にクエリーとキー・バリューのペアを出力にマッピングするスケールド・ドット・プロダクト・アテンション（SDPA）メカニズムを通して、自己注意を形式化し、最新のニューラルネットワークにおける極めて重要な構成要素にしている。ここでは、アテンション計算を最適化するための最も重要なテクニックをいくつか紹介する：

マルチヘッドアテンション（MHA）は、クエリ、キー、および値の行列をそれぞれ別個に投影して、複数のアテンション操作を並列に実行することによって、SDPAを基礎とする。これらの並列演算（「ヘッド」）は、異なる表現部分空間に焦点を当て、入力に対するモデルの理解を深める。これらのヘッドからの出力は連結され、線形に投影される。各ヘッドの次元を小さくする（例えば、モデルの次元をヘッドの数で割る、例えば8）ことで、単一ヘッドの注意に匹敵する計算効率を維持する。

マルチクエリーアテンション（MQA）は、複数のクエリープロジェクションを保持しながら、複数のアテンションヘッド間でキーと値のプロジェクションを共有することで、推論のためにMHAを最適化する。これにより、メモリ帯域幅の需要とキー・バリュー（KV）キャッシュのサイズが削減され、より大きなバッチサイズとより優れた計算利用が可能になる。しかし、MQAは精度を若干低下させる可能性があり、MQAを活用したモデルでは、性能を維持するためにMQAを有効にしてトレーニングや微調整を行う必要がある。

Grouped-query attention (GQA)は、クエリーヘッドをグループ化し、各グループ内でキー値の投影を共有することで、MHAとMQAのバランスをとり、MQAに近い計算効率でMHAに近い品質を達成する。Llama 2 70BのようなモデルはGQAを使用しており、MHAで訓練されたモデルは最小限の追加訓練でGQAに適応させることができる。MQAもGQAもKVキャッシュメモリの需要を削減するが、キャッシュ管理のさらなる最適化が必要である。

FlashAttentionは、GPUのメモリ階層をより効果的に活用するために計算の順序を変更することで、アテンションメカニズムを強化します。従来のレイヤーごとの処理とは異なり、FlashAttentionは演算を融合し、「タイリング」を使用して出力行列の小さな部分を一度に計算し、メモリの読み書きを最小限に抑えます。このI/Oを意識した正確なアテンション・アルゴリズムは、既存のモデルに変更を加えることなくシームレスに統合でき、データ移動を最適化することで大幅なスピードアップを実現します。

キー・バリュー・キャッシュ

KVキャッシングは、大規模言語モデル(LLM)のデコード段階において、自己アテンション計算の効率を向上させるために用いられる重要な最適化手法である。このフェーズでは、生成された各トークンは、プレフィルステージとそれに続くデコードステップで計算されたものを含む、すべての前のトークンのキー（K）と値（V）テンソルに依存する。KVキャッシュは、時間ステップごとにトークンごとにこれらのテンソルを再計算する代わりに、それらをGPUメモリに保存し、計算されるたびに新しいテンソルをキャッシュに追加します。通常、モデルの各レイヤーに対して個別のKVキャッシュが維持されるため、冗長な計算が大幅に削減され、デコード処理が高速化されます。

GPU上のLLMに必要なメモリは、主にモデル重みとKVキャッシュの2つです。例えば、Llama 2 7Bのような70億パラメータのモデルを16ビット精度で計算する場合、約14GBが必要です。一方KVキャッシュは、再計算を避けるために自己アテンソルを保存し、そのサイズはレイヤー数、アテンション・ヘッド数、ヘッド寸法、精度などの要素によって決定される。各トークンについて、キャッシュサイズは2 * num_layers * (num_heads * dim_head) * precision_in_bytesとして計算される。バッチ入力の場合、KVキャッシュ・サイズの合計は、バッチ・サイズとシーケンス長によってスケールし、シーケンス長4,096、バッチ・サイズ1のLlama 2 7Bモデルで〜2GBのような大きなサイズに達する可能性がある。

KVキャッシュの効率的な管理は、バッチサイズとシーケンス長に伴う線形成長により、スループットが制限され、ロングコンテキストの入力の取り扱いが複雑になるという課題をもたらす。一般的な非効率性は、静的な過剰プロビジョニングから生じる。この場合、実際の入力サイズに関係なく、サポートされる最大シーケンス長（例えば、2048トークン）分のメモリが予約される。これは、予約されたスペースの多くがリクエストのライフタイム中未使用のままであることが多く、貴重なGPUメモリリソースを占有するため、大幅なメモリの浪費や断片化につながります。

これらの非効率性に対処するため、PagedAttentionアルゴリズムは、オペレーティングシステムのページングにヒントを得た新しいアプローチを導入している。KVキャッシュを固定サイズのブロックに分割し、それぞれがトークンのセット数を表し、メモリに非連続に格納することができる。ブロックテーブルがこれらのブロックを追跡し、アテンション計算中に必要に応じてフェッチする。新しいトークンが生成されると、追加のブロックが動的に割り当てられる。この方法は、連続的な割り当てやオーバープロビジョニングの必要性を排除することでメモリの浪費を最小化し、より大きなバッチサイズを可能にし、スループットを向上させるため、LLMのKVキャッシュメモリ管理における重要な進歩となっている。

モデルの最適化

このセクションでは、大規模言語モデル（LLM）を最適化し、メモリ消費量を削減し、GPU上でのパフォーマンスを向上させるための様々なテクニックについて説明します。主な手法には、量子化、疎分散、蒸留があり、それぞれモデル効率の異なる側面を対象としています。これらの手法は、モデルの重みを変更し、GPUハードウェアアクセラレーションを活用し、知識をより小さなモデルに転送することで、性能を維持しながらより大きなモデルを限られたハードウェアで実行できるようにします。これらの手法はモデルの精度を低下させる可能性があるため、使用には注意が必要です。

量子化はモデルの重みと活性度の精度を下げ、通常32ビットまたは16ビットから8ビット以下にすることで、モデルのメモリ占有量を減らし、より効率的にデータを転送できるようにする。重みの量子化は学習後に固定されるため簡単ですが、活性度の量子化はダイナミックレンジを拡大する外れ値のため複雑です。LLM.int8()のようなテクニックは、特定の活性化に対して選択的に高精度を適用したり、量子化された重みのダイナミックレンジを活性化に再利用することでこれに対処しますが、GPUでは演算のために重みを高精度に変換し直す必要があるかもしれません。

スパース性は、ゼロに近いモデル値を刈り込み、より少ないメモリしか必要としないスパース行列を作成します。GPUは、4つの値のうち2つをゼロとして表現するなど、構造化されたスパース性をサポートし、計算を高速化します。スパース性を量子化と組み合わせることで、実行速度をさらに向上させることができます。LLMに最適なスパース表現の研究が続けられており、推論速度の向上が期待されている。

ディスティレーションは、より大きな「教師」モデルからより小さな「生徒」モデルへと知識を転送し、性能を維持しながらサイズを圧縮します。例えば、DistilBERTはBERTと比較して、97%の能力を保持したまま、40%のサイズ縮小と60%の速度向上を達成している。Distillationには、教師の出力を模倣したり、教師が生成したデータをトレーニングに使用したりすることが含まれ、「Distilling Step by Step！」のような手法には、効率的な学習のための根拠が組み込まれています。しかし、多くの先進的なLLMの制限されたライセンスは、蒸留に適した教師モデルの利用を制限している。

投機的推論

投機的推論（投機的サンプリングやアシスト生成としても知られる）は、GPTスタイルのような自己回帰型大規模言語モデル（LLM）の実行を並列化する手法であり、通常、トークンごとにテキストを生成する。標準的な実行では、各トークンはコンテキストのために先行するすべてのトークンに依存するため、n番目のトークンは(n+1)番目のトークンの前に生成されなければならず、並列生成は不可能である。投機的推論では、「より安価な」ドラフトモデルを使用して複数の将来のトークンを同時に予測し、それをメインモデルで並列に検証または拒否することで、より高速なテキスト生成を可能にします。

このプロセスでは、リソースをあまり消費しない方法で複数のトークンからなるドラフトを生成し、その後、ドラフトを投機的コンテキストとして使用してメインモデルによる検証を並行して行う。検証モデルがドラフトトークンと一致した場合、そのトークンは受理され、一致しな いトークンや後続のトークンは破棄され、新しいドラフトでプロセスが繰り返される。ドラフト・トークンは、複数のモデルを訓練する、事前に訓練したモデルで複数のヘッ ドを微調整して将来のトークンを予測する、あるいは、より大規模で高性能な検証モデルとともに小規模な ドラフト・モデルを採用するなど、さまざまなアプローチで生成することができます。

分解推論

分散推論とは、パフォーマンス、コスト、リソースの使用を最適化するために、計算タスクを異なるハードウェアに分割する手法である。具体的には、プリフィリングとデコーディングのフェーズを分離する。これらのフェーズを分割することで、それぞれの計算需要に最適なハードウェアに割り当てることができ、効率とスケーラビリティが向上する。

プリフィリングは計算集約的で、入力プロンプト全体を処理しKVキャッシュを生成するために、かなりの行列乗算を必要とする。このフェーズは、並列計算を得意とするGPUやTPUのような高性能ハードウェアの恩恵を受ける。プリフィリングは推論要求ごとに1回だけのタスクであるため、このようなワークロードに最適化された集中型の強力な計算ノードにオフロードすることができる。このセットアップにより、大規模なプロンプトの高速処理が可能になり、性能の低いデバイスの負担が軽減されるため、高スループットのハードウェアが利用可能なクラウドベースやデータセンター環境に最適です。

これとは対照的に、デコーディングはメモリに縛られ、トークン生成を繰り返し、KVキャッシュへのアクセスに大きく依存する。そのため、CPUやエッジ・デバイスのような、メモリが最適化されていないハードウェアに適しています。デコーディングを別のハードウェア（オンプレミスのサーバーやエッジデバイスなど、エンドユーザーに近い可能性がある）に移行することで、分離推論はレイテンシーとネットワーク帯域幅の需要を削減する。この分離により、ハイエンドのクラウドサーバーでプリフィルを実行し、ローカルまたはエッジデバイスでデコードを行う柔軟な展開が可能になり、リソースの割り当てを最適化し、リアルタイムのチャットボットや対話型AIシステムのようなアプリケーションの効率的なスケーリングが可能になります。

結論

近年、LLMの性能を向上させるために、多くの推論最適化技術が考案されている。

これらのテクニックを実装するには、LLMアーキテクチャと使用しているハードウェアを深く理解する必要があるため、一般的には、vLLM、TensorRT-LLM、LMDeployなど、これらのテクニックをすでに実装している既存の推論エンジンを使用する方が簡単です。NLP Cloudでは、これらのテクニックを独自の推論エンジンに実装しており、独自のモデルをデプロイしたい場合は、推論エンジンに関するブログ記事を書いています： こちらでお読みいただけます。.

独自のLLMを自分でデプロイできない、またはデプロイしたくない場合は、NLPクラウドを使用し、プロダクションでスケールの大きな高速生成AIモデルを活用することができます。 今すぐNLPクラウドで高速推論をお試しください！

推論エンジン全般についてご質問があれば、遠慮なくお尋ねください！

Julien
NLPクラウドCTO