Amazon Search がAWS Inferentiaで検索システムのコストを85%削減

この記事は How Amazon Search reduced ML inference costs by 85% with AWS Inferentia を翻訳したものです。
Amazon の製品検索エンジンは、何十億もの製品をインデックスし、世界中の何億人もの顧客にサービスを提供している、世界で最も利用されているサービスの1つです。Amazon Search チームは、 Amazon.com の検索エンジンを強化し、顧客が簡単に検索できるようにする機械学習 (ML) テクノロジーを開発しています。優れたカスタマーエクスペリエンスを提供し、 Amazon.com 検索エンジンが必要とする大規模な運用を実現するために、このチームは常に、リアルタイムのレイテンシーと高いスループットを実現する、より費用対効果の高いシステムを構築する方法を模索しています。チームは、オペレーションコストを全面的に削減しつつ、モデルのトレーニングと推論を高速化するためにディープラーニング用に最適化されたハードウェアとコンパイラを常に模索しています。
この投稿では、Amazon Search で AWS Inferentia をどのように使用しているかを説明します。AWS Inferentia は、ディープラーニングの推論ワークロードを高速化するために AWS が構築した高性能アクセラレータです。
チームは、AWS Inferentia ベースの Amazon Elastic Compute Cloud (Amazon EC2) Inf1 インスタンスで、Transformer ベースの自然言語処理（NLP）モデルを使用して低レイテンシーの ML 推論を実行し、強力なスループットとレイテンシーのパフォーマンスを維持しながら、インフラストラクチャコストを最大 85% 節約します。

類似商品の識別とクエリの意図を予測するためのディープラーニング

Amazonマーケットプレイスの検索は、ASIN（Amazon標準識別番号、商品を一意に識別する10桁の英数字）、商品画像、テキストの説明、クエリなど、複数の入力を処理するマルチタスク、マルチモーダルな問題です。個別化されたユーザー体験を実現するために、検索のさまざまな側面でさまざまなモデルからの予測が使われています。検索システムには何千ものモデルがあり、ピーク負荷時に毎秒数万トランザクションが発生するため、これを実現するのは一筋縄ではいきません。私たちはその中の2つの要素に焦点を当てています。

類似商品の識別 — ユーザーのクエリに一致する最も関連性の高い製品リストを表示するには、区別するのが難しい製品を特定することが重要です。
クエリの意図の予測 — 検索ページと製品のレイアウトを顧客が探しているものに合わせて調整するには、ユーザーのクエリの意図とタイプ（たとえば、メディア関連のクエリ、ヘルプクエリ、その他のクエリタイプ）を予測することが重要です。

これらの予測は両方とも、Transformer モデルアーキテクチャ、つまり BERT ベースのモデルを使用して行われます。実際、どちらも同じ BERT ベースのモデルを基本としており、それぞれがこのバックボーンの上で分類/回帰処理を行っています。

類似商品の識別は、評価対象の2つの商品のさまざまな特徴（製品タイプ、タイトル、説明など）を入力とし、大規模なデータセットで定期的に計算されます。このモデルは、マルチタスク方式でエンドツーエンドでトレーニングされます。これらのバッチワークロードの定期的な起動を自動化するために Amazon SageMaker Processing job が使用され、使用された処理時間に対してのみ支払いが行われます。このバッチワークロードのユースケースでは、推論スループットの要件は合計 8,800 TPS でした。

意図の予測は、Amazonマーケットプレイスでのユーザーエクスペリエンスを向上させるために、ユーザーのテキストクエリを入力として受け取り、毎日のトラフィックをリアルタイムに処理することが必要です。このモデルはマルチクラス分類の問題としてトレーニングされており、 Amazon Elastic Container Service (Amazon ECS) にデプロイされています。これにより、迅速な自動スケーリングが可能になり、デプロイの定義と管理も簡単になります。これはリアルタイムのユースケースであるため、快適なユーザーエクスペリエンスを実現するため、P99のレイテンシを10ミリ秒未満にする必要がありました。

AWS Inferentia と AWS Neuron SDK

EC2 Inf1 インスタンスは、ディープラーニングの推論ワークロードを高速化するために AWS が開発した最初の ML アクセラレータである AWS Inferentia を搭載しています。Inf1 インスタンスは、同等の GPU ベースの EC2 インスタンスよりもスループットが最大 2.3 倍高く、推論あたりのコストが最大 70% 削減されます。好きなフレームワーク (PyTorch、TensorFlow、MXNet) を使用してモデルをトレーニングし、そのモデルを AWS Inferentia に簡単にデプロイして、組み込みのパフォーマンス最適化の恩恵を受けることができます。Inf1 インスタンスを使用して、画像認識、オブジェクト検出、NLP、最新のレコメンダーモデルなど、さまざまなモデルタイプをデプロイできます。
AWS Neuron は、EC2 Inf1 インスタンスの ML 推論パフォーマンスを最適化するコンパイラ、ランタイム、およびプロファイリングツールで構成されるソフトウェア開発キット (SDK) です。Neuronは、TensorFlow や PyTorch などの一般的な ML フレームワークとネイティブに統合されています。そのため、選択したフレームワークで提供されるのと同じ使い慣れた API を使用してディープラーニングモデルを AWS Inferentia にデプロイでき、パフォーマンスが向上し、推論あたりのコストがクラウドで最小というメリットが得られます。
Neuron SDKは発表以来、パフォーマンスの向上と推論コストの削減を続けながら、サポートするモデルを増やし続けてきました。これには、自然言語処理モデル (BERT)、画像分類モデル (ResNet、VGG)、およびオブジェクト検出モデル (OpenPose および SSD) が含まれます。

Inf1 インスタンスにデプロイすることで、低レイテンシー、高スループット、コスト削減を実現

Amazon Search チームは、類似商品の識別の高スループット要件と意図予測の低レイテンシー要件を満たしながら、コストを節約したいと考えていました。彼らは AWS Inferentia ベースの Inf1 インスタンスにデプロイすることを選択し、高いパフォーマンス要件を満たしただけでなく、推論コストを最大 85% 節約できました。

顧客認識の類似商品の識別

Inf1 を使用する前は、専用の Amazon EMR クラスターで CPU ベースのインスタンスを使用して実行していました。ハードウェアアクセラレーションに頼らずに1 秒あたり合計 8,800 トランザクションという高いスループット要件を満たすには、多数のインスタンスが必要でした。チームは、それぞれ 4 つの AWS Inferentia アクセラレーターと 16 の Neuron コア (AWS Inferentia チップあたり 4 コア) を搭載した inf1.6xlarge インスタンスに切り替えました。Transformer ベースのモデルを 1 つの Neuron コアでトレースし、スループットを最大化するために Neuron コアごとに 1 つずつモデルをロードしました。16 個の Neuron コアを利用することで、推論コストを 85% 削減しました（当時のAmazon EC2 オンデマンド価格に基づく）。

クエリの意図の予測

P99 のレイテンシー要件が 10 ミリ秒以下であることを考慮して、チームは inf1.6xlarge インスタンスで利用可能なすべての Neuron コアにモデルをロードしました。PyTorch Neuronでは、torch.neuron.DataParallel API を使用してこれを簡単に行うことができます。Inf1導入により、モデルのレイテンシは3ミリ秒、エンドツーエンドのレイテンシは約10ミリ秒、ピーク負荷時の最大スループットは16,000 TPSに達しました。

コンパイルとデプロイコードのサンプル

以下は、Inf1 インスタンスの使用を開始し、Amazon Search チームのようにパフォーマンスとコストのメリットを実感するのに役立つサンプルコードです。 PyTorch Neuron を使用して、PyTorchモデルをコンパイルして推論を実行する方法を示します。まず、モデルは torch.neuron.trace() でコンパイルされます：


m = torch.jit.load(f="./cpu_model.pt", map_location=torch.device('cpu'))
m.eval()
model_neuron = torch.neuron.trace(
    m,
    inputs,
    compiler_workdir="work_" + str(cores) + "_" + str(batch_size),
    compiler_args=[
'--fp32-cast=all', '--Neuron コア-pipeline-cores=' + str(cores)
])
model_neuron.save("m5_batch" + str(batch_size) + "_cores" + str(cores) +
"_with_extra_op_and_fp32cast.pt")

traceメソッドに指定できる引数の全リストについては、 PyTorch-Neuron trace Python API を参照してください。ご覧のとおり、コンパイラーの引数は torch.neuron API に直接渡すことができます。FP32 オペレータはすべて --fp32-cast=all で BF16 にキャストされ、ダイナミックレンジを維持しながら最高のパフォーマンスを実現します。性能とモデル精度のトレードオフを制御できるように、他にもcastingオプションが用意されています。今回使用されたモデルは両方とも、単一のNeuron コア用にコンパイルされました（パイプラインなし)。次に、 torch.jit.load を使用してモデルを Inferentia にロードし、推論に使用します。Newronランタイムは、モデルを Neuron コアに自動的にロードします。


cm_cpd_preprocessing_jit = torch.jit.load(f=CM_CPD_PROC,
                                          map_location=torch.device('cpu'))
cm_cpd_preprocessing_jit.eval()
m5_model = torch.jit.load(f=CM_CPD_M5)
m5_model.eval()
input = get_input()
with torch.no_grad():
    batch_cm_cpd = cm_cpd_preprocessing_jit(input)
    input_ids, attention_mask, position_ids, valid_length, token_type_ids = (
        batch_cm_cpd['input_ids'].type(torch.IntTensor),
        batch_cm_cpd['attention_mask'].type(torch.HalfTensor),
        batch_cm_cpd['position_ids'].type(torch.IntTensor),
        batch_cm_cpd['valid_length'].type(torch.IntTensor),
        batch_cm_cpd['token_type_ids'].type(torch.IntTensor))
    model_res = m5_model(input_ids, attention_mask, position_ids, valid_length,
                         token_type_ids)

まとめ

Amazon Search チームは、トラフィックが多くパフォーマンス要件が厳しい状況において、AWS Inferentia ベースの Inf1 インスタンスを使用して推論コストを 85% 削減することができました。AWS Inferentia と Neuron SDK のおかげで、チームはトレーニングとは独立してデプロイプロセスを最適化できる柔軟性が得られました。また、充実したツールと使い慣れたフレームワーク API により、短期間で習得できました。
この記事で提供されているサンプルコードから始めることで、パフォーマンスとコスト上のメリットを引き出すことができます。また、 PyTorch と TensorFlow を使って Inferentia で ML モデルを実行するためのエンドツーエンドのチュートリアルもチェックしてください。

翻訳は Prototyping Engineer の秋田が担当しました。

Amazon Web Services ブログ