研究から本番環境へ:EAGLE-3がVertex AIでオープンソースLLMの推論を2-3倍高速化

EAGLE-3 Vertexカバー画像

LLM開発に携わる実務者にとって、一度に1つのトークンしか生成できないボトルネックは共通の悩みです。標準的な自己回帰生成は本質的に逐次的なプロセスであり、これによりメモリ帯域幅に制限された計算モードが生じます:GPUコアはアイドル状態になり、時間は主に各ステップでメモリから巨大なモデル重みを読み取ることに費やされます。

解決策はspeculative decoding(推測デコーディング)です。この最適化技術は、ドラフト機構を導入することで、大型LLM(ターゲットモデル)の遅いトークンごとの生成プロセスを高速化します。ドラフト機構は複数の後続トークンを素早く提案し、ターゲットモデルはこれらの提案を並列バッチで検証し、最長一致プレフィックスを受け入れて、そこから生成を続けます。

しかし、すべてのドラフト機構が効率的というわけではありません。古典的なドラフト-ターゲット手法では、小型LLMをdrafterとして別途デプロイする必要があり、ホスティングリソースとコストが増加します。

従来のドラフトモデルアーキテクチャ

ここで、EAGLE-3(Extrapolative Attention Guided LEarning)の登場です。完全に独立したモデルは不要で、ターゲットモデルの内部レイヤーに極めて軽量なdraft head(ターゲットモデルのわずか2-5%のサイズ)を直接追加します。このヘッドは特徴とトークンレベルで動作し、隠れ状態を利用して将来のトークンのツリー構造を推測します。

結果は?推測デコーディングのすべての利点を保持しながら、第2のモデルの訓練と実行のオーバーヘッドを排除します。複雑な数億パラメータのドラフトモデルと比較して、EAGLE-3は軽量ヘッドのみを訓練し、Llama 70Bなどのモデルに対して2x-3xのデコーディング高速化を実現します(マルチターン会話、コード、長いコンテキストなどのワークロードによって異なります)。

EAGLE-3ターゲットモデルアーキテクチャ

EAGLE-3を論文から大規模本番クラウドサービスへと移行するには、多くのエンジニアリング課題を克服する必要があります。本稿では技術パイプライン、主要な問題点、貴重な経験を共有します。

課題1:データ準備

EAGLE-3ヘッドには訓練が必要ですが、公開データセットには以下の問題がよくあります:

  • 厳格な利用規約:生成モデルは競合開発に使用禁止。
  • PII汚染:名前、場所、金融識別子などを含む。
  • 品質保証なし:デモには適しているが、顧客の本番ワークロードには不適。

経験1:合成データ生成パイプラインの構築

解決策:顧客シナリオに基づいて優良データセットを選択し、ユーザープロンプトを抽出し、DLPとPIIフィルタリングを経て、チャットテンプレートを適用し、トークン化後にターゲットモデル(例:Llama 3.3 70B)に入力してレスポンスを生成します。この方法により、コンプライアンスに準拠し、クリーンで、モデル分布に適合したデータが生成され、draft head訓練に完璧に適合します。

データ生成パイプライン

課題2:訓練パイプラインのエンジニアリング

データ入力方法はonline training(リアルタイム埋め込み生成)とoffline training(事前計算埋め込み)に分かれます。ハードウェア要件が低いため、我々はofflineを選択しました:事前計算された特徴/埋め込みをGCSに保存し、訓練に使用します。EAGLE-3ヘッドが小さいため、初期訓練は単一ホストで1日のみ必要で、データセット拡張後は数日に増加します。

訓練パイプライン

経験2:チャットテンプレートは不可欠

指示調整モデルの訓練で、ターゲットモデルのチャットテンプレート(例:Llama 3)を使用しない場合、埋め込みにエラーが発生し、ヘッドが偏った分布を学習します。

経験3:マスクへの注意

訓練入力にはプロンプトとレスポンスが含まれますが、EAGLE-3はレスポンスのみを予測します。損失関数でプロンプトを手動でマスクする必要があり、そうしないとヘッドは既知のプロンプトの予測に容量を無駄にし、パフォーマンスが低下します。

マスク処理の図解

課題3:サービングとスケーリング

経験4:サービングフレームワークが重要

SGLangチームと協力し、EAGLE-3を効率的に本番環境にデプロイしました。SGLangのtree attention kernelは、EAGLE-3のドラフトツリー(分岐パス)の並列検証専用に設計されており、パフォーマンス損失を回避します。

経験5:CPUにGPUの足を引っ張らせない

EAGLE-3による高速化後、CPU オーバーヘッド(カーネル起動、メタデータ管理など)が新たなボトルネックになります。同期スケジューラーでは、GPU DraftはCPU Verify準備を待つ間アイドル状態になります。

通常のスケジューラー

SGLangのZero-Overhead Overlap Schedulerがこの問題を解決します:FutureMapを利用して、GPUがVerifyを実行している間、CPUは次のDraft/Extendを並列準備し、アイドル時間を排除し、追加の10%-20%高速化を実現します。

オーバーラップスケジューラー

ベンチマーク結果

成果は顕著です:SGLang上で、Llama 4 Scout 17B Instructのベンチマークを使用して、EAGLE-3は2x-3xのデコーディングレイテンシー削減とスループット向上を実現しました。

指標1:中央値出力トークン時間(TPOT)

TPOTベンチマークグラフ

緑のEAGLE-3ラインは、すべての同時実行度で青のベースラインよりもTPOTが低く、より良いレイテンシーを示しています。

指標2:出力スループット

スループットベンチマークグラフ

緑のラインは同時実行度全体でベースラインを大幅に上回り続け、EAGLE-3の優位性を示しています。完全なnotebookを確認して独自のベンチマークを実行できます。