Amazon SageMaker AIとAmazon BedrockでvLLMを用いた多数の微調整済みモデルを効率的に配信

Amazonは、特に最近注目されているMixture of Experts（MoE）モデルファミリーなど、複数のカスタムAIモデルを実行している組織や個人向けに、vLLMコミュニティと提携し、Multi-Low-Rank Adaptation（Multi-LoRA）による効率的な配信ソリューションを開発しました。Multi-LoRAは、モデルの全重みを再学習する代わりに、元の重みを凍結し、モデルの層に小さな学習可能なアダプターを挿入することで、モデルを微調整する方法です。このアプローチにより、複数のカスタムモデルが同じGPUを共有して実行でき、例えば、各々が専用GPUの10%しか使用していない5人の顧客は、単一のGPUでサービスを受けることができます。 このブログ投稿では、vLLMにおけるMoEモデルのMulti-LoRA推論の実装方法、行ったカーネルレベル最適化、そしてこれらの改善を利用する方法について説明します。GPT-OSS 20Bを主要な例として使用しています。Amazon SageMaker AIまたはAmazon BedrockでLoRAカスタマイズされたモデルをホストすることで、これらの最適化を活用できます。

Multi-LoRAは、複数のLoRAアダプターが異なるユーザーやタスクのために同時に配信されるマルチ-LoRA配信設定では、各専門家ごとに4つの操作（shrinkとexpand）が必要となるため、MoEモデルの主要なパフォーマンスボトルネックとなります。これらの4つの操作には、1つの次元（LoRAランクr）が他の次元（隠れ状態と中間表現の次元）よりも100〜300倍小さいマトリックスが含まれます。標準的なGEMMカーネルは、ほぼ正方形のマトリックス用に設計されており、細長いマトリックスではパフォーマンスが低下するため、後述するカーネル最適化が必要となります。

さらに、MoE LoRAは、専門家のルーティング（異なる専門家に割り当てられたトークン）とアダプターの選択（異なるLoRAアダプターを使用するリクエスト）という2つの疎密性のソースを組み合わせるため、複合的な疎密性が必要です。これらの課題に対処するために、vLLMにfused_moe_loraカーネルを作成しました。この新しいカーネルは、gate_upとdownプロジェクションのLoRA shrinkとexpand GEMMsを実行します。fused_moe_loraカーネルは、fused_moeカーネルと同じロジックに従い、対応するアクティブ化されたLoRAアダプターのグリッドに追加の次元を追加します。

この実装がvLLMに統合されると、H200 GPU上でGPT-OSS 20Bを使用したマルチ-LoRA配信を実行し、Sonnetデータセット（詩に基づくベンチマーク）で入力長1600、出力長600、並列度16のOTPS 26とTTFT 1053 msに達することができました。これらの結果を再現するには、vLLM GitHubリポジトリのリリース0.11.1.rc3からのPRを確認してください。このブログの残りの部分では、これらのベースラインの実装からパフォーマンスを向上させるための最適化方法について説明します。

背景

MoEモデルは、複数の専門家ネットワークで構成され、各専門家は特定のタスクに特化した処理を行います。このアプローチにより、より大きなモデルを少ない計算リソースで処理できます。Multi-LoRAは、モデルの全重みを再学習する代わりに、元の重みを凍結し、小さなアダプターを挿入することで、MoEモデルを微調整する方法です。

重要用語解説

今後の影響

この技術は、大規模言語モデルの効率的な配信と実行に役立ちます。特に、MoEモデルのような複雑なモデルをより少ないリソースで運用できるようになります。これは、AI開発者や研究者がより高度なAIシステムを構築し、デプロイするのに役立つ可能性があります。