サーバーレス推論のための大規模言語モデルの未来

GPT-4やPaLMなどの大規模言語モデルの最近の進歩により、自然言語タスクにおける変革的な能力が実現しました。LLMは、チャットボット、検索エンジン、プログラミングアシスタントなどのさまざまなアプリケーションに組み込まれています。しかし、LLMを大規模に提供することは、GPUとメモリの要件が大きいため、依然として課題です。

これを克服するためのアプローチは、一般的に2つのカテゴリに分けられます。

1. モデル圧縮技術

これらの技術は、モデルを小さくしながら精度を維持することを目的としています。一般的なアプローチには、次のものがあります。

• プルーニング – モデルから冗長または重要度の低いパラメータを削除します。これにより、パラメータが少ないスパースモデルが作成されます。
• 量子化 – fp32またはfp16ではなく、int8やbfloat16などの低精度数値を使用して重みを表します。これにより、メモリのフットプリントが削減されます。
• 知識の伝達 – 小さい「学生」モデルを大きな「先生」モデルに似せるようにトレーニングします。小さいモデルは推論に使用されます。

2. 選択的な実行

これらの技術は、圧縮されたモデルではなく、推論ごとにモデルの一部のみを選択的に実行します。

• スパースなアクティベーション – ゼロのアクティベーションに対する計算をスキップします。
• 条件付き計算 – 入力に基づいて、特定の層のみを実行します。

ソフトウェアアーキテクチャの観点から、LLMのより迅速な展開を可能にするために、研究者はサーバーレス推論システムを提案しました。サーバーレスアーキテクチャでは、LLMは共有GPUクラスタにホストされ、需要に応じて動的に割り当てられます。これにより、GPUの効率的な利用と開発者のコスト削減が可能になります。著名な実装には、Amazon (AMZN ) SageMaker、Microsoft Azure ML、およびKServeなどのオープンソースオプションがあります。

サーバーレスLLMの約束にもかかわらず、既存のシステムはインタラクティブアプリケーションでユーザーエクスペリエンスを低下させる高い待機時間オーバーヘッドを示しています。

1. コストのかかるチェックポイントのダウンロード – LLMには、大きなメモリフットプリントがあり、サイズはギガバイトからテラバイトに及ぶことがあります。最適化されたネットワークでさえも、リモートストレージからのダウンロードには20秒以上かかります。
2. 非効率的なチェックポイントのロード – 甚至ローカルSSDストレージの場合も、チェックポイントをGPUメモリにロードするには、テンソルのデシリアライゼーションや割り当てなどの要因により10秒以上かかります。これにより、コンテナの起動時間以外に大きな遅延が発生します。

これらの問題に対処するために、MIT CSAILの研究者は、ServerlessLLMという、LLMのサーバーレス推論のための革新的なシステムを提案しました。ServerlessLLMは、LLMの展開のために、多層サーバーストレージの豊富で未利用の容量と帯域幅を利用することで、ローカリティを高めます。

LLMサーバーレス推論システムの概要

ServerlessLLMの重要な革新

1. 高速なチェックポイントのロード

• ロード最適化されたチェックポイント形式 – テンソルは、高速なシーケンシャルリードと効率的なインメモリテンソルアドレッシングを可能にするように、GPUごとにバイナリファイルにグループ化されます。
• マルチ階層チェックポイントロードパイプライン – ネットワーク、SSD、DRAM、GPUメモリ全体で、ダイレクトI/O、ピンミュートランスファー、並列性などのテクニックを使用して帯域幅の利用率を最大化します。

2. ライブマイグレーションによるローカリティ駆動型推論

• トークンに基づくマイグレーション – スナップショット転送よりもはるかに少ないデータを転送するために、ネットワーク上で必要なトークンのみを転送します。
• 2段階のマイグレーション – 宛先サーバーは、トークンからキャッシュ状態を事前に計算し、ソースサーバーは最終的なトークンを転送する前にリソースを解放します。これにより、推論の停止を防ぎます。

3. 待機時間最適化されたサーバー割り当て

• 精密なモデル – 各サーバーのチェックポイントロード時間とマイグレーション時間を、キューの遅延、モデルサイズ、測定された帯域幅などのメトリックを使用して予測します。
• ローカリティに基づくスケジューラー – 各推論リクエストに対して、サーバー全体の推定ロード時間とマイグレーション時間を評価し、期待される起動待機時間を最小化するサーバーを選択します。

これらの最適化により、ServerlessLLMは、既存のシステムと比較して、LLMのロード時間を4〜8倍、エンドツーエンドの起動時間を25倍以上短縮できます。

チェックポイントのロードを高速化する

ServerlessLLMが対処する最初の大きなボトルネックは、ストレージからGPUメモリへのLLMチェックポイントのロードの高い待機時間です。

高速なチェックポイントロードを可能にするために、ServerlessLLMは次のものを導入します。

1. ロード最適化されたチェックポイント形式

PyTorchなどのフレームワークで使用される標準的なチェックポイントは、モデルトレーニングとデバッグのために設計されています。しかし、サーバーレス推論の場合、チェックポイントは読み取り専用であり、繰り返しアクセスされます。

このような読み取り集中の使用を最適化するために、ServerlessLLMはチェックポイントを、2つの重要な特性を持つ形式に変換します。

• シーケンシャルチャンクベースの読み取り – テンソルは、GPUごとにバイナリファイルにグループ化され、大きなシーケンシャルリードを可能にします。
• 効率的なテンソルアドレッシング – テンソル名をメモリオフセットにマッピングするインデックスにより、デシリアライゼーションなしに直接インメモリ復元が可能になります。

2. マルチ階層チェックポイントロードパイプライン

ServerlessLLMは、GPUサーバーの階層化アーキテクチャを利用し、ストレージメディア、ネットワーク、SSD、DRAM、GPUメモリを接続します。

システムは、すべての階層全体で帯域幅の利用率を最大化するための、マルチステージパイプラインを組み込みます。

• インメモリデータチャンクは、GPU転送のためにピンミュートランスファーを使用して割り当てられます。
• ダイレクトI/Oは、キャッシングオーバーヘッドなしで効率的なSSD読み取りを可能にします。
• 複数のスレッドが、ストレージの異なるチャンクを並列に読み取ります。
• ステージ間の調整は、非同期タスクキューを介して行われます。

これらすべてが組み合わさることで、最速の階層であるNVMe RAIDの帯域幅容量を飽和させることができます。実験により、ServerlessLLMはPyTorch/TensorFlowと比較して6〜8倍高速にロードできることがわかり、起動時間を1分以上から10秒未満に短縮します。

ライブマイグレーションによるローカリティ駆動型LLM推論

高速化されたロードとともに、ServerlessLLMは、忙しいサーバーでの推論を中断せずに、事前にロードされたチェックポイントをローカリティを利用する方法を探る新しい課題に直面しています。

ServerlessLLMは、GPUサーバー間でLLM推論をシームレスに移行できる、ライブマイグレーションという新しいテクニックを導入します。

ライブLLMマイグレーションの重要な要素:

1. トークンに基づくマイグレーション

モデル全体の状態をスナップショットするのではなく、ServerlessLLMはネットワーク上で必要なトークンのみを移行します。これにより、スナップショットよりもはるかに少ないデータが転送されます。

2. 2段階のマイグレーション

宛先サーバーは、トークンからキャッシュ状態を事前に計算します。準備が整ったら、ソースサーバーは最終的なトークンを転送する前にリソースを解放します。これにより、推論の停止を防ぎます。

実験により、トークンに基づくマイグレーションが長いシーケンスの場合でも10秒以上から1秒未満にマイグレーション時間を短縮することがわかりました。ライブマイグレーションは、ローカリティ駆動型の割り当てを実現するために不可欠です。

待機時間最適化されたモデルスケジューリング

エンドツーエンドの待機時間を最小限に抑えるために、ServerlessLLMはスケジューラーをローカリティを考慮して最適化します。これには、次のものが含まれます。

1. ロード時間の詳細な推定

モデルは、ネットワーク、SSDキャッシュ、メモリの各サーバーからのロード時間を、キューの遅延、モデルサイズ、測定された帯域幅などのメトリックを使用して予測します。

2. マイグレーション時間の正確な予測

スケジューラーは、サーバーごとに推論の進行状況を追跡することなく、プロンプトと出力トークンの数を使用してマイグレーション時間を予測します。

3. ローカリティに基づく割り当て

各推論リクエストに対して、スケジューラーはサーバー全体の推定ロード時間とマイグレーション時間を評価し、期待される起動待機時間を最小化するサーバーを選択します。

スケジューラーはまた、サーバーのタスクキューを維持し、ストロングコンシステントストアを使用して障害耐性を提供します。これらすべての革新により、スケジューリングのオーバーヘッドが削減され、ローカリティの利点が最大化されます。

ServerlessLLMのパフォーマンスの評価

包括的な実験は、OPT-175BやAzureトレースに基づくワークロードなどの実際のモデルを使用して、ServerlessLLMのエンドツーエンドの有効性を既存のシステムと比較します。

重要な結果:

• マイクロベンチマーク – ServerlessLLMは、PyTorch/TensorFlowと比較して、チェックポイントのロードを3.6〜8.2倍高速化します。ストレージの帯域幅を完全に飽和させることができます。
• スケジューリング – ServerlessLLMは、ランダムなスケジューリングと比較して、割り当ての待機時間を4〜12倍削減します。ローカリティの利点が強調されます。ライブマイグレーションにより、キューの遅延が防止されます。
• エンドツーエンドの提供 – 大きなモデルであるOPT-30Bの場合、ServerlessLLMは、KServeやRay Serveなどのシステムと比較して、99パーセンタイルの待機時間を28〜200倍改善します。また、リソースの効率も向上します。

これらの大幅な改善は、ServerlessLLMが既存のサーバーレス実装のボトルネックを克服し、インタラクティブなサービスでLLMの力を解放できることを示しています。

ServerlessLLMで導入された最適化、たとえばマルチ階層ロード、ライブマイグレーション、待機時間駆動型スケジューリングは、将来のサーバーレスアーキテクチャの設計に役立つことができます。システムのチェックポイントのロード時間と起動時間を短縮する能力は、実用的アプリケーションで大規模言語モデルのスケーラブルな展開を可能にします。

今後への展望：継続的な課題

大きな飛躍であるServerlessLLMは、サーバーレス推論のための大規模言語モデルの最適化の最初のステップを表します。まだいくつかの開かれた問題があります。

• モデル需要の予測 – プロビジョニングと事前のロードを導く
• サーバー全体でチェックポイントを配置してキャッシュヒットを最大化する
• スケジューリングアルゴリズムをより大きなクラスターに効率的にスケーリングする
• モデルと開発者全体でリソースの割り当てを公平に確保する
• ライブマイグレーションのような革新を他のサーバーレスワークロードに一般化する

これらの分野に取り組むことで、サーバーレスLLMの約束をさらに実現し、その機能をさらにアクセスしやすくすることができます。システムレベルの最適化を超えて、大規模モデルがもたらす大きな炭素足跡と潜在的な危害を減らすことも、緊急の優先事項です。

ServerlessLLMは、AIワークロードのための次世代のサーバーレスアーキテクチャにおける革新のための大きな余地があることを示しています。LLMがサイズと人気で増大し続けるにつれて、ServerlessLLMのようなスケーラビリティを解放するソリューションは、さらに重要になります。システムと機械学習の研究の融合により、AIモデルを安全に、持続可能に提供、共有、スケーリングするための新しいパラダイムが導入されます。