HUOXIU

KubeAI 大規模モデル推論加速実践 | Dewu Technology

出典:デウーテクノロジー

目次

I. 背景

II. 大規模モデル開発における課題

III. 大規模モデル推論を加速するための最適化の方向性

1. Llama2のモデル構造

2. 大規模モデル推論を加速するための方向性は何ですか?

IV. FlashAttention-Attention演算子計算の最適化

V. PageAttention - メモリ管理の最適化 VI. MOE - 推論中のモデルパラメータの削減

VII. テンソル並列化

VIII. 推奨される推論フレームワーク

IX. 要約と展望

1つ

背景

私たちは最近、大規模モデルに特化した推論クラスターを本番環境に導入しました。その結果、70Bを含む大規模モデルの推論速度が50%向上し、導入コストを大幅に削減し、本番環境での安定した運用を実現できました。この記事では、大規模モデル推論クラスターの導入経験に基づき、大規模モデルの推論速度を向上させる効果的な手法をいくつかご紹介します。最後に、私たちが評価し、優れたパフォーマンスを発揮することが確認された大規模モデル推論フレームワークをいくつかご紹介します。これらの提案が、読者の皆様がプロジェクトに最適な推論フレームワークを選択する一助となれば幸いです。
OpenAIの科学者であるヒョン・ウォン・チョン氏は、2023年の公開プレゼンテーション「大規模言語モデル」[8]において、大規模モデルの一部の機能は、ある程度の規模に達したときに初めて発揮されると指摘しました。大規模モデルのパラメータ数は今後確実に増加していくことは明らかであり、これは大規模モデルの開発動向でもあります。パラメータ数が増加するにつれて、大規模モデルの推論速度に対する要求はますます高くなります。大規模モデルの推論速度やスループットを向上させるには、どのような方法があるでしょうか?
まず、大規模モデルの高速化最適化の方向性を探ります。次に、タイムラインに基づいて、業界で古典的かつ実用的な大規模モデル高速化手法をいくつか紹介します。これには「FlashAttention[1]」や「PageAttention[3]」などが含まれますが、これらに限定されません。
以下は、業界における大規模モデル推論の加速化手法の代表的な例を時系列で示したものです。この記事では、大規模モデル加速化手法の概要を時系列順に紹介します。
上記の手法に加えて、量子化手法など、大規模モデルの推論速度を向上させる方法は他にもありますが、ここでは説明しません。これらについては、機会があれば別の記事で紹介する予定です。

大規模モデル開発における課題

今後の大規模モデルのパラメータ数は確実に増加しますが、これは大規模モデルの開発傾向でもあり、推論加速に対する要件はますます高くなります。
OpenAIは論文「ニューラル言語モデルのスケーリング則」[7]において、大規模モデルのスケーリング則を提示し、モデルの能力とサイズの関係を説明しています。具体的には、モデルの能力は、モデルパラメータの数、データセットのサイズ、学習に必要な計算量など、そのサイズに大きく依存します。さらに、OpenAIの科学者であるHyung Won Chung氏は、2023年の公開講演「大規模言語モデル」[8]において、大規模モデルの一部の能力は、ある一定の規模に達したときに初めて明らかになると指摘しました。
上の図は、ヒョン・ウォン・チョン氏のプレゼンテーションスライド[8]から引用したものです。この図は主に、GPT3からGPT4のようにモデルサイズが大きくなるにつれて、モデルの能力がますます強くなり、さらには新たな能力が出現する可能性があるという見解を表しています。
しかし、モデルサイズが大きくなるにつれて、パラメータ数の増加に伴うGPU計算量の増加により、大規模モデルの推論速度は徐々に低下します。この推論速度の低下はユーザーエクスペリエンスの悪化につながるため、大規模モデルの推論高速化がますます重要になっています。

三つ

大規模モデル推論を加速するための最適化の方向性

Llama2モデル構造

まず、Llama 2の論文[9]を参考に、Llama 2モデルシリーズの構造を簡単に理解しましょう。現在、Llamaシリーズのような生成言語モデルのほとんどは、TransformerアーキテクチャのDecoderモジュールを主に利用しています。Huggingfaceプラットフォームでは、この種のモデル構造は通常、CausalLM(Causal Language Model)と呼ばれます。
上の画像はLlama2モデルの構造を示しており、最も重要なコンポーネントはAttention計算(Llama Attention)です。これは推論プロセス全体の中で最も時間のかかるモジュールでもあり、その後の最適化のほとんどはAttentionに基づいています。Llama2モデルの構造をより深く理解するために、まずLlama2モデルの推論プロセス全体を簡単に説明します。興味のない方はこのステップを飛ばしてください。
  1. ユーザーがモデルにプロンプ​​トを送信すると、モデルの最初の操作は次の文字(トークン)を予測し、予測された文字を入力に追加して予測を続行します。このプロセスは、モデルがストップトークンを出力するまで継続され、ストップトークンが出力されると予測は停止し、最終結果を出力します。
  2. 次のトークンを生成するために、モデルはN回のLlamaデコーダーレイヤー計算を実行する必要があります。具体的には、Llama-2-7Bモデルは32回の計算を実行し、Llama-2-13Bモデルは40回の計算を実行します。
  3. ラマデコーダー層における最も重要な計算ステップは、アテンション(Llama Attention)の計算です。推論時間の大部分はアテンションの計算に費やされるため、様々な最適化手法がアテンション計算の効率向上を目指しています。

大規模モデル推論を加速するための方向性は何ですか?

Llama 2 モデルの構造分析から、推論計算中の大規模モデルの次の特性をまとめることができます。
  1. 推論プロセス全体の中で最も時間のかかる部分は、注目度の計算です。注目度の計算速度を最適化することで、全体的な推論性能を大幅に向上させることができます。
  2. アテンション計算中、キーバリューキャッシュ(KVキャッシュ)は大量のGPUメモリを消費します。13Bモデルを例にとると、1つのプロンプトシーケンスを処理するには約3GBの追加GPUメモリが必要となり、このメモリは頻繁に割り当てと解放が行われるため、大きな断片化が発生します。GPUメモリの断片化を軽減することで、大規模モデルのスループットを向上させることができます。
  3. 推論プロセスでは、膨大な数のパラメータをGPUで処理・計算する必要があります。7Bモデルは70億、13Bモデルは130億のパラメータを持ち、最新かつ世界最強のモデルであるDBRXは驚異的な1300億のパラメータを誇ります。そのため、これらのパラメータを効率的に処理することが不可欠です。しかしながら、最適化の余地は残されています。
上記の 3 つの特性に対処するために、業界ではいくつかの効果的な最適化手法が提案されています。その一部を以下に示します。
1. FlashAttention-Attention計算速度の最適化
FlashAttention[1]は、Attention演算子の計算結果を変えることなく、Attention演算子の計算速度を向上させます。FlashAttentionは、様々なモデルやタスクにおいて大幅な性能向上を示してきました。例えば、BERT-largeやGPT-2などのモデルにおいて、FlashAttentionはベースライン実装と比較して、エンドツーエンドで15%から3倍の高速化を達成できます。
2. PageAttention-KV キャッシュメモリ管理の最適化
PageAttention[3]の目的は、メモリの断片化を減らすことです。Pag​​eAttentionに基づくVLLMシステムは、安定した時間消費分布を維持しながら、一般的な大規模言語モデル(LLM)のスループットを10倍以上向上させることができます。
3. MOE - 推論中のモデルパラメータの削減
MOE(Mixture of Experts)[4]の目的は、モデル推論中の計算に関係するパラメータの数を減らすことです。
実験結果:Mixtralモデルは、ほとんどのベンチマークテストでLlama 2 70Bモデルを上回り、推論速度は6倍高速です。このモデルは複数の言語をサポートし、強力なコード生成機能を備え、特定の指示に従うように細かく設定できるため、MT-Benchベンチマークテストで高得点を達成しました。
次に、上記の各方向について詳しく説明します。

4つ

FlashAttention - Attention演算子の計算最適化

FlashAttentionは、Attention演算子の最適化を説明する2つの論文、FlashAttention-1[1]とFlashAttention-2[2]を発表しています。ここでは、FlashAttention-1[1]を例に挙げ、その最適化原理を理解していきます。
まず、GPUのメモリ階層を理解しましょう。FlashAttention-1[1]の論文から引用した次の図を参照してください。
GPUのメモリ階層は、SRAM、HBM、DRAMという3つの主要な部分で構成されています。以下はA100 GPUのリファレンス構成です。
SRAM (Static Random Access Memory) はアクセス速度が最も速い (19TB/s) ですが、容量は比較的小さい (わずか 20MB) です。
HBM (High Bandwidth Memory) は、大容量ストレージ (40 GB) と高速データ アクセス (1.5 TB/秒) を提供します。
DRAM (ダイナミック ランダム アクセス メモリ) は、具体的には GPU 外部のメイン メモリを指し、容量は最大 (1 TB 以上) ですが、アクセス速度は最も遅い (12.8 GB/秒) です。
上記の構成からわかるように、メモリ容量が小さいほど、処理速度は速くなります。
従来のAttention計算では、HBMへのアクセスによって多数の入出力操作が実行されます。FlashAttentionアルゴリズムは、Attention計算プロセスを最適化し、HBMへのアクセス回数を削減することで計算効率を向上させるため、I/Oを考慮した最適化アルゴリズムとなっています。
下の図は、論文FlashAttention-1[1]に掲載されているFlashAttentionの加速方法を示しています。
FlashAttentionは、高速かつメモリ効率の高いアテンション計算を実現する巧妙な手法を採用しています。入力データをタイリングすることで、巨大なアテンション行列全体を一度に処理することを回避します。タイリングは通常、膨大なメモリと計算リソースを必要とします。巨大な図書館(行列)があると想像してみてください。FlashAttentionのアプローチは、図書館の本を小さな山に分割し、一度に1つの山を処理するようなものです。こうすることで、すべての本を一度にテーブルに取り出す必要がなくなります(大きなテーブルと膨大な時間が必要になります)。
具体的には、FlashAttentionは行列計算を行う際に、データをブロックに分割し、GPU上の高速だが小容量のSRAMを計算に利用することで、低速だが大容量のメモリ(HBM)へのアクセスを効果的に削減します。これにより、計算速度が向上するだけでなく、GPUメモリの需要も大幅に削減されます。
FlashAttentionは、低速ストレージへの依存を減らすことで、モデルのパフォーマンスを維持、あるいは向上させながら、モデルのトレーニング速度を大幅に向上させます。例えば、BERT-largeのトレーニングはMLPerf 1.1の記録よりも15%高速化し、GPT-2のトレーニングはHuggingFaceおよびMegatron-LMのベースラインよりも3倍高速化し、長シーケンスドメインのトレーニング速度は2.4倍向上しました。
以下の画像はhuggingfaceのブログ[14]から引用したもので、フラッシュアテンションを紹介しており、フラッシュアテンションがどのようにマトリックスを分割するかを理解するのに役立ちます。
Flash Attentionは計算を高速化できますが、どのようなフレームワークがFlash Attention計算をサポートしているのでしょうか?この記事の後半では、優れた推論フレームワークをいくつかご紹介します。

PageAttention - メモリ管理の最適化

PageAttention[3]の概念は、VLLMの著者であるWoosuk Kwon氏に​​よって初めて提案され、VLLM推論フレームワークにおける最も重要な最適化戦略でもあります。Woosuk Kwon氏は論文の中で、PageAttentionを用いて大規模言語モデル(LLM)サービスにおける重要な問題を解決する方法、すなわち、メモリを効果的に管理することでレイテンシを増加させずにスループットを向上させる方法を紹介しました。
まず、推論における大規模モデルのメモリ構造の分布を理解しましょう。次の図は論文[3]からの抜粋です。
これは、13Bのパラメータを持つ大規模言語モデルをNVIDIA A100で提供するためのメモリレイアウトです。13BのLLM推論メモリ割り当ては26GBで、KVキャッシュはリクエストごとに12GBのメモリを占有します。QPSが増加すると、KVキャッシュは急速に増加し、頻繁に割り当てと解放が行われるため、メモリの断片化が大量に発生します。放置すると、システムは徐々にクラッシュします。
では、VLLMはPageAttentionを通じてビデオメモリの断片化の問題をどのように解決するのでしょうか?次の図は論文[14]からの引用で、VLLMのビデオメモリ管理技術を示しています。
PageAttentionは、キーバリューキャッシュ(KVキャッシュ)を固定サイズのブロック(または「ページ」)に分割し、これらのブロックをメモリ内に非連続に格納できるようにすることで機能します。このアプローチは、オペレーティングシステムの仮想メモリとページング技術に着想を得たもので、メモリリソースをより柔軟かつ効率的に管理することを目的としています。
従来のアテンションメカニズムでは、リクエストのキーバリューキャッシュをメモリ内に連続して保存する必要があり、メモリの断片化と非効率的なメモリ共有という2つの主な問題が発生します。メモリの断片化はバッチサイズを制限し、メモリを共有できないことでデータの重複が発生し、貴重なメモリリソースが無駄になります。
PageAttention は次の手順でこれらの問題に対処します。
  1. セグメント化されたKVキャッシュ:各リクエストのKVキャッシュを複数の小さなブロックに分割します。これらのブロックのサイズは固定されており、モデルやハードウェアの特定のニーズに応じて調整できます。
  2. 非連続ストレージ:従来のキーバリューキャッシュブロックはメモリ内に連続して保存されますが、PageAttentionではこれらのブロックを物理メモリ内に非連続に分散させることができます。これにより、メモリブロックを実際のニーズに応じて動的に割り当て、再利用できるため、メモリの無駄を削減できます。
  3. 動的管理:オペレーティングシステムの仮想メモリ管理と同様に、PageAttentionはこれらのメモリブロックを動的に管理します。システムは、現在のメモリ使用量に基づいて、必要に応じてキーバリューキャッシュブロックを割り当てたり解放したりすることで、メモリ使用率を最適化します。
  4. メモリ共有:PageAttentionは、異なるリクエスト間、または同一リクエスト内の異なるシーケンス間でのキーバリューキャッシュブロックの共有もサポートしています。この共有は柔軟で、ブロックレベルで実行できるため、メモリ使用量をさらに削減し、効率性を向上させます。
このように、PageAttention を使用すると、LLM サービス システムは、同じレイテンシを維持しながらメモリの無駄を減らし、メモリ共有を増やすことで、要求のスループットを大幅に向上させることができます。
PageAttentionの最適化により、VLLMはLLaMA 7Bと13Bのスループットを10倍以上向上させます(下の図は論文[11]より)。

MOE - 推論中にモデルパラメータを削減する

最近リリースされた、1300億パラメータを有する世界最強のオープンソース大規模モデルであるDBRXと、Mistralの8x7Bオープンソース大規模モデルは、どちらもMOEアーキテクチャに基づいています。なぜ、より大きなパラメータ数を持つモデルにMOEアーキテクチャを採用すべきなのでしょうか?Mistralの8x7Bオープンソース大規模モデルを例に、MOEアーキテクチャのパフォーマンス上の利点をご紹介します。
MOE大規模モデルについて言えば、まずは上図に示すように、通常の大規模モデルとMOE大規模モデルの構造上の違いを比較してみましょう。MOE大規模モデルでは、大規模モデルのパラメータは8つのグループとルーターに分かれており、各グループはエキスパートグループと呼ばれます。リクエストが来ると、MOE大規模モデルはまずルーターに8つのエキスパートグループから2つを選択させ、選択した2つのエキスパートグループだけが計算に参加します。一方、通常の大規模モデルでは、すべてのパラメータがGPU計算に参加する必要があります。
そのため、MOE 大規模モデルは、同レベルの通常の大規模モデルよりも推論速度が約 4 倍高速になります。
Mistral MOEの実装を見てみましょう。Mistral MOEは、mistral.ai [12]によって公開された8*7Bの大規模モデルです。次の図は論文[12]からの抜粋で、8*7Bの大規模モデルのエキスパート層の構造を示しています。
Mixtral 8x7Bは、Sparse Mixture of Experts (SMoE)言語モデルです。Mistral 7Bアーキテクチャに基づいていますが、各層は8つのフィードフォワードブロック(つまりエキスパート)で構成されています。各トークンを処理する際、各層のルーティングネットワークは、現在の状態を処理する2つのエキスパートを選択し、それらの出力を結合します。各トークンは2つのエキスパートとのみ対話しますが、各タイムステップで選択されるエキスパートは異なる可能性があります。そのため、各トークンは47バイトのパラメータにアクセスできますが、推論中に使用されるアクティブパラメータは13バイトのみです。
Mixtralは複数のベンチマークにおいて、特に数学、コード生成、多言語理解において優れた性能を示しました。Llama 2 70BおよびGPT-3.5と比較すると、Mixtralはほとんどの評価指標において同等以上の性能を示しました。具体的には、MixtralはLlama 2 70B(70B)と比較して5分の1のアクティブパラメータ(13B)を使用するにもかかわらず、ほぼすべてのカテゴリにおいてLlama 2 70Bよりも優れた、あるいは同等の性能を示しました。
MOE大規模モデルは、推論速度を低下させることなくパラメータ数を増やすという目標を達成しており、これは大規模モデル開発の今後の動向です。

セブン

テンソル並列化

複数の GPU がある場合は、テンソル並列処理を使用して、大規模モデルの推論速度をさらに高速化できます。
非常に分厚い本を一度に全部コピーしたいのに、コピー機では一度に数ページしかコピーできないとします。このような場合、本をいくつかの部分に分割し、それぞれを個別にコピーし、コピーした部分を順番につなぎ合わせることで、本全体のコピーを完了できます。
テンソル並列処理では、扱う大規模モデルは分厚い本のようなもので、GPUはコピー機のようなものです。単一のGPUでは大規模モデル全体を一度に処理できないため、モデル(この場合は重みテンソル)を複数の部分に分割し、異なるGPUで個別に処理できるようにする必要があります(本の異なる部分をコピーするのと同じです)。入力データの処理は、本の各ページを個別にコピーし、それらをつなぎ合わせて完全な出力を形成するようなものです。
このように、ワークロードを分散することで、複数のGPUが連携し、単一のGPUでは処理できなかった大規模なタスクを完了します。これがテンソル並列処理の仕組みであり、非常に大規模なモデルの処理を可能にします。

記事[13]からの画像

テンソル並列処理は、大規模モデルを複数のGPUに分散させるために使用されます。行列乗算を例に挙げると、入力テンソルに最初の重みテンソルを乗算する場合、この操作は、まず重みテンソルを列ごとに分割し、次に各列に入力テンソルを乗算し、最後にこれらの積を結合すると考えることができます。これらの結合された出力はGPUからエクスポートされ、上図に示すように最終出力を形成するために集約されます(論文[13]を参照)。

推奨される推論フレームワーク

これまでのセクションでは、Flash Attention、Page Attention、MOE、テンソル並列処理といった、いくつかの高速化および最適化手法について検討しました。次に、私たち自身の実践的な経験と評価に基づき、現在高性能な推論フレームワークをいくつか推奨します。

要約と展望

本稿では、大規模モデルの推論速度向上を目的とした、Flash Attention、Page Attention、MOE、テンソル並列処理など、様々な技術と手法について詳細に考察します。大規模モデル専用の推論クラスターを本番環境に導入することで、最大70バイトのモデルを含む推論速度を50%削減することに成功し、安定した本番環境への適用を通じてこれらの最適化手法の有効性と実用性を実証しました。
大規模モデルが様々な分野で利用されるようになるにつれ、推論速度の向上と推論コストの削減を効果的に図ることが課題となっています。私たちの実践では、現在利用可能な加速技術を紹介するだけでなく、経験に基づき、高性能な大規模モデル推論フレームワークをいくつか推奨します。これらの推奨は、読者が数多くの選択肢に直面した際に、ニーズに最適な推論フレームワークを選択できるよう支援することを目的としています。
今後、継続的な技術進歩と新たなアルゴリズムの登場により、さらなる高速化・最適化技術が開発され、大規模モデル推論の効率がさらに向上すると考えています。また、大規模モデル推論の速度を向上させるための新たな技術や手法を今後も探求し、導入していく機会を楽しみにしています。

参考文献
[1] FlashAttention: IOを考慮した高速かつメモリ効率の高い正確なアテンション(https://arxiv.org/abs/2205.14135)
[2] FlashAttention-2: より高速なアテンションとより優れた並列処理および作業分割(https://arxiv.org/abs/2307.08691)
[3] PagedAttentionを用いた大規模言語モデルサービングのための効率的なメモリ管理(https://arxiv.org/abs/2309.06180)
[4] ミストラル・オブ・エキスパート(https://mistral.ai/news/mixtral-of-experts/)
[5] 専門家のミックストラル(https://arxiv.org/abs/2401.04088)
[6] MEDUSA: 複数のデコードヘッドを備えたシンプルなLLM推論加速フレームワーク(https://arxiv.org/pdf/2401.10774.pdf)
[7] ニューラル言語モデルのスケーリング則(https://arxiv.org/pdf/2001.08361.pdf)
[8] Hyung Won Chung(OpenAI)、大規模言語モデル(2023年)、ソウル国立大学での講演
[9] ラマ2:オープンファウンデーションと微調整されたチャットモデル(https://arxiv.org/abs/2307.09288)
[10] 必要なのは注目だけ(https://arxiv.org/pdf/1706.03762.pdf)
[11] https://blog.vllm.ai/2023/06/20/vllm.html
[12] https://arxiv.org/pdf/2401.04088.pdf
[13] https://huggingface.co/docs/text-generation-inference/en/conceptual/tensor_parallelism
[14] https://huggingface.co/docs/text-generation-inference/en/conceptual/flash_attention
[15] https://blog.vllm.ai/2023/06/20/vllm.html