HAQM の生成 AI 搭載ショッピングアシスタント Rufus を、80,000 以上の AWS AI チップを活用して Prime Day 向けにスケーリング

この記事は Scaling Rufus, the HAQM generative AI-powered conversational shopping assistant with over 80,000 AWS Inferentia and AWS Trainium chips, for Prime Day (記事公開日 : 2024 年 10 月 10 日) の翻訳です。翻訳は Annapurna Labs の常世が担当しました。

HAQM Rufus は、生成AIを活用したショッピングアシスタントです。HAQM の商品情報やウェブ上の様々な情報を活用して回答を作成し、お客様のよりスマートなお買い物をサポートします。Rufus を使えば、HAQM の商品に精通した AI アシスタントと一緒にショッピングを楽しむことができます。また、ウェブ上の幅広い情報も組み合わせることで、より確かな商品選びができるようになります。

HAQM が世界中の幅広いお客様にサービスを提供するにあたり、Rufus には数十億のパラメータを持つ大規模言語モデル（LLM）を、低コストかつ低レイテンシーで処理できる安定性の高い推論基盤が必要でした。低レイテンシーにより、ユーザーは Rufus とスムーズにやり取りでき、1 秒以内に回答が表示され始めます。これを実現するために、Rufus チームは AWS の様々なサービスや、AWS Trainium と AWS Inferentia といった専用 AI チップを活用しています。

Inferentia と Trainium は、AWS が開発した専用 AI チップで、深層学習ワークロードを高性能かつ低コストで実現します。これらのチップを活用することで、Rufus は顧客への低レイテンシーを維持しながら、他の検討していたソリューションと比べてコストを約 4.5 分の 1 に抑えることができました。この記事では、AWS AI チップを使った Rufus の推論システムの導入について、そして年間で最も負荷の高いイベントの 1 つである HAQM Prime Day をどのように実現できたのかについて詳しく紹介します。

ソリューション概要

Rufus の基盤となっているのは、HAQM の商品カタログとウェブ全体の情報を学習した LLM です。LLM を実用化する際には、モデルの規模や精度、処理性能などのバランスを取る必要があり、これが大きな課題となっています。一般的に、モデルが大きいほど知識や推論の能力は向上しますが、必要な計算量も増え、レイテンシーも増加するため、コストが上昇してしまいます。Rufus は HAQM Prime Day のようなアクセスが集中する時期にも対応できるよう、デプロイとスケーリングを行う必要がありました。この規模でのサービス提供には、必要な性能の確保はもちろん、消費電力上昇による環境への影響やシステム運用のコストなども考慮しなければなりません。これらの課題を解決するため、Rufus は以下の AWS ソリューションを組み合わせて活用しました：Inferentia2 と Trainium、HAQM Elastic Container Service（HAQM ECS）、Application Load Balancer（ALB）です。また、Rufus チームは NVIDIA とも協力し、NVIDIA の Triton Inference Server（推論を効率化するために NVIDIA が開発したオープンソースの推論サーバー）を使って AWS チップを活用したソリューションを構築しました。

Rufus は、検索拡張生成（RAG : Retrieval Augmented Generation）システムを採用しており、お客様の質問に応じて HAQM の検索結果から商品情報などを取得し、より充実した回答を提供します。この仕組みにより、LLM は信頼性が高く、質の高い正確な回答を確実に生成できるようになっています。

Prime Day を万全の体制で迎えるため、Rufus チームは Inferentia2 と Trainium を活用し、複数の AWS リージョンをまたいだハイブリッドな推論システムを構築しました。複数のリージョンでシステムを展開することで、Rufus は2つの重要なメリットを得ることができました。1 つ目は、アクセスが集中する時期に備えた追加の処理能力の確保、2 つ目は、システム全体の安定性の向上です。

Rufus チームは Inferentia2 を搭載した HAQM EC2 Inf2 インスタンス（以下、Inf2 ）と Trainium を搭載した HAQM EC2 Trn1 インスタンス（以下、Trn1） という 2 種類のインスタンスタイプを利用しました。両者は同じ AWS Neuron SDK を使用しているため、どちらのインスタンスでも同じ Rufus モデルを動かすことができました。設定で調整が必要だったのは、テンソル並列度（Inf2 が 24、Trn1 が 32）だけでした。より高性能な Trn1 インスタンスを使用することで、Inf2 と比べてレイテンシーが 20% 削減され、スループットも向上しました。

次の図はソリューションアーキテクチャを示しています。

複数のリージョンをまたぐリアルタイムのトラフィックルーティングを実現するため、Rufus は革新的なトラフィックオーケストレーターを構築しました。HAQM CloudWatch を使ったモニタリングを基盤とし、トラフィックパターンの変化に応じて 15 分以内にリージョン間のトラフィック配分を調整できるようにしました。このタイプのオーケストレーションにより、必要に応じて他のリージョンへリクエストを振り分けることが可能になりました。その際、最初のトークンまでにわずかなレイテンシーが生じることはありますが、Rufus のストリーミングアーキテクチャとリージョン間の高速な AWS ネットワークにより、エンドユーザーが体感するレイテンシーは最小限に抑えられています。

これらの取り組みにより、Rufus は 3 つのリージョンで 80,000 以上の Trainium と Inferentia2 チップまでスケールしました。その結果、Prime Day 中も顧客への初回応答までの P99 レイテンシーを 1 秒未満に保ちながら、1 分間に平均 300 万トークンを処理することができました。さらに、これらの専用チップの採用により、他の検討済みのソリューションと比べてワットあたりの性能が 54% 向上し、チームが目標としていた省エネルギー基準もクリアすることができました。

推論性能とホスト利用率の最適化

Rufus では HAQM ECS を使って推論システムを運用し、Inferentia2 と Trainium チップを搭載したインスタンスを管理しています。ECS がインフラストラクチャの管理を担当することで、Rufus チームは必要なコンテナと設定を ECSタスクとして定義するだけで済むようになりました。各コンテナでは、Python バックエンドを持つ NVIDIA Triton Inference Server を採用し、Neuron SDK を利用して vLLM を実行しています。vLLM は、高スループットのために最適化されたメモリ効率の高い推論・サービングエンジンです。また、Neuron SDK の採用により、チームは簡単に AWS チップを利用でき、PyTorch Lightning をはじめとする様々なライブラリやフレームワークも使用できるようになりました。

Neuron SDK は、Trainium と Inferentia チップ向けに最適化された、深層学習推論および学習用の SDK で、使いやすい LLM 向け分散推論および分散学習ライブラリを提供し、様々な種類の Transformer ベースの LLM アーキテクチャに対応しています。レイテンシーを削減するため、Rufus は AWS Annapurna チームと協力して、INT8（重みのみ）量子化、vLLM を使用した連続バッチング (continuous batching)、Neuron コンパイラとランタイムでのリソース、計算、メモリ帯域幅の最適化など、さまざまな最適化を開発しました。これらの最適化は現在 Rufus の本番環境にデプロイされており、Neuron SDK 2.18 以降で利用可能です。

お客様が Rufus の応答をより早く見られるようにするため、チームは推論のストリーミングアーキテクチャを開発しました。LLM 推論では大量の計算処理とメモリを必要とするため、質問に対する完全な回答の生成には数秒かかることがあります。ストリーミングアーキテクチャを採用することで、Rufus は生成した内容をリアルタイムで返せるようになりました。この改善により、お客様は 1 秒以内に応答を見始めることができます。さらに、複数のサービスが gRPC 接続を通じて連携し、ストリーミング形式の応答をリアルタイムでスマートに統合・強化しています。

以下の図に示すように、Rufus の応答には画像とリンクが組み込まれており、お客様はそれらを使ってさらに詳しい情報を探ることができます。

スケールアップ

最高の顧客体験のために低レイテンシーを維持する必要がありますが、同時にハードウェアリソースを効率的に使用してサービスのスループットを向上させることも重要です。ハードウェアの稼働率を上げることで、高性能な機器が無駄にアイドル状態になることを防ぎ、コスト増加を抑えられます。システム全体のスループットを最適化するため、チームは個々のサーバーのスループットと、複数のサーバー間での負荷分散の効率の両方を改善しました。

LLM 推論の負荷分散は、以下の課題があるため複雑です。まず、1 台のホストで同時に処理できるリクエスト数に限りがあります。次に、1 つのリクエストの処理完了までにかかる時間（エンドツーエンドのレイテンシー）は、LLMの生成する応答の長さによって数秒単位で変動する可能性があります。

これらの課題に対処するため、チームは個々のサーバーのスループットと、負荷分散を使った複数サーバー全体でのスループットの両方を考慮して最適化を行いました。

チームは ALB の最小未処理リクエスト（LOR : least outstanding requests）ルーティングアルゴリズムを採用し、以前の基準と比べてスループットを 5 倍に向上させました。これにより、各サーバーは同時に多数のリクエストを受け取っても処理しきれなくなることなく、現在処理中のリクエストを適切に扱い、gRPC 接続を使って応答をストリーミング形式で返す十分な時間を確保できます。さらに Rufus は AWS と vLLM チームと協力し、Neuron SDK と NVIDIA Triton Inference Server を使用した vLLM 統合により、一台のサーバーでの同時実行性を改善しました。

図 1. ECS タスクは水平方向にスケールし、Triton Inference Server と 依存関係をホスティングします

この統合により、Rufus は重要な最適化である継続的バッチング（continuous batching）を活用できるようになりました。継続的バッチングにより、1 台のホストのスループットを大幅に向上させることができます。また、継続的バッチングは静的バッチング（static batching）などの従来のバッチ処理技術とは異なる特長があります。たとえば、静的バッチングでは、最初のトークンが生成されるまでの時間（TTFT : time to first token）がバッチ内のリクエスト数に比例して長くなります。一方、継続的バッチングでは LLM 推論のプリフィルステージ (Prefill stage : 入力されたプロンプト全体を初期処理) を優先することで、同時実行されるリクエスト数が増えても TTFT を一定に保つことができます。この仕組みにより、Rufus は最初の応答を低レイテンシーで返せるようになり、快適な体験を提供しながら、1 台のホストのスループットも改善し、サービスのコストも抑制できるようになりました。

結論

この記事では、Rufus が Neuron SDK や Inferentia2、Trainium チップ、そして AWS の各種サービスを活用して、数十億のパラメータを持つ LLM を安定的にデプロイし、運用する方法をご紹介しました。Rufus は、生成 AI の技術進歩とお客様からのフィードバックを取り入れながら進化を続けています。皆様もぜひ Inferentia と Trainium をご活用ください。

HAQM の生成 AI 分野でのイノベーションについて、さらに詳しくはこちらをご覧ください。