Amazon EC2 Inf1、Inf2 インスタンスにおける FastAPI と PyTorch モデルの AWS Inferentia 利用時の最適化

深層学習モデルを大規模にデプロイする際には、パフォーマンスとコストのメリットを最大化するために、基盤となるハードウェアを効果的に活用することが重要です。高スループットかつ低レイテンシであることが必要となる本番ワークロードの場合、Amazon Elastic Compute Cloud (EC2) インスタンス、モデルサービングのスタック、デプロイアーキテクチャの選択が非常に重要です。非効率的なアーキテクチャはアクセラレーターの最適化がなされていないことから、不必要に高い本番コストにつながる可能性があります。

この記事では、AWS Inferentia アクセラレータ（Amazon EC2 Inf1 およびAmazon EC2 Inf2 インスタンスに搭載されています）に FastAPI モデルサーバーをデプロイする流れを説明します。また、ハードウェア効率の最大化を実現するために、すべてのNeuronCore に並列にデプロイされたサンプルモデルのホスティングも示します。

ソリューションの概要

FastAPI は、Flask や Django などの従来のフレームワークよりもはるかに高速な、Python アプリケーションをホストするためのオープンソースの Web フレームワークです。これらは広く使用されている Web Server Gateway Interface (WSGI) ではなく、Asynchronous Server Gateway Interface (ASGI) を活用することで実現しています。ASGI は、WSGI がリクエストを順次処理するのとは異なり、非同期に処理します。このことから、FastAPI は低レイテンシが要求されるリクエストを処理するのに最適な選択肢となっています。FastAPI を使用することで、指定されたポートを介してクライアントからのリクエストを受け取るサーバーを Inferentia (Inf1/Inf2) インスタンスにデプロイできます。

私たちの目標は、ハードウェアを最大限に活用することで、最低コストで最高のパフォーマンスを達成することです。これにより、より少ないアクセラレーターでより多くの推論要求を処理できます。各 AWS Inferentia1 デバイスには 4 つの NeuronCore-v1 が含まれ、各 AWS Inferentia2 デバイスには 2 つの NeuronCore-v2 が含まれます。AWS Neuron SDK を使用すると、各 NeuronCore を並列に使用できるため、スループットを犠牲にすることなく、4 つかそれ以上のモデルを並列にロードおよび推論する際の制御が可能となります。

FastAPI を使用する際には、Python の Web サーバー (Gunicorn、Uvicorn、Hypercorn、Daphne) の選択肢があります。これらの Web サーバーは、基盤となる機械学習 (ML)  モデルの上に抽象化レイヤーを提供します。要求元クライアントは、ホストされたモデルを認識する必要がありません。クライアントは、サーバー下でデプロイされたモデルの名前やバージョンを知る必要はありません。すなわち、エンドポイント名はモデルをロードして実行する関数へのプロキシになります。対照的に、フレームワーク固有のサービングツール、たとえば TensorFlow Serving では、モデルの名前とバージョンがエンドポイント名の一部となります。サーバー側でモデルが変更された場合、クライアントは API 呼び出しを新しいエンドポイントに変更する必要があります。したがって、A/B テストなどの状況でモデルのバージョンを継続的に改善させている場合、FastAPI を備えた汎用 Python Web サーバーを使用すると、エンドポイント名が静的であるため、モデル提供時には便利な方法であるといえます。

ASGI サーバーの役割は、指定された数のワーカーを生成してクライアントからのリクエストを受信し、推論コードを実行することです。サーバーの重要な機能は、要求された数のワーカーが使用可能かつアクティブであることを確認することです。ワーカーが停止された場合、サーバーは新しいワーカーを起動する必要があります。この状況においては、サーバーおよびワーカーは Unix のプロセス ID (PID) によって識別される場合があります。このブログでは、Python Web サーバーとして人気のある Hypercorn サーバを使用します。

このブログでは、AWS Inferentia NeuronCore 上で FastAPI を使用し、深層学習モデルをデプロイするためのベストプラクティスを共有します。ここでは、複数のモデルを個別の NeuronCore にデプロイして、同時に呼び出すことができることを示します。このセットアップにおいては、複数のモデルが同時に推論できるため、スループットが向上し、NeuronCore の使用率が完全に最適化されます。コードは GitHub リポジトリに公開されています。次の図は、EC2 Inf2 インスタンスに今回のソリューションを構築したアーキテクチャを示しています。

EC2 Inf1 インスタンスタイプでも同じアーキテクチャが適用されますが、4 つのコアがあります。そのためアーキテクチャ図が少し変更されます。

AWS Inferentia NeuronCore

NeuronCore を使用するための AWS Neuron が提供するツールについて詳しく見てみましょう。次の表は、各 Inf1 および Inf2 インスタンスタイプに搭載されている NeuronCore の数を示しています。ホスト vCPU とシステムメモリは、すべての利用可能な NeuronCore で共有されます。

Inf2 インスタンスは、Inf1 インスタンスの NeuronCore-v1 に比べて新しい NeuronCores-v2 を含んでいます。コアの数は少ないですが、Inf1 インスタンスよりも4倍大きいスループットと10倍小さいレイテンシを提供することができます。Inf2インスタンスは、Generative AI や OPT/GPT 系統の大規模言語モデル、Stable Diffusion や vision transformers などの深層学のワークロードに最適な選択肢となります。

Neuron Runtimeは、Neuron デバイス上でモデルを実行する役割を担っています。Neuron Runtime は、どの NeuronCore がどのモデルをどのように実行するかを決定します。Neuron Runtime の設定は、プロセスレベルで環境変数を使用して制御されます。デフォルトでは、Neuron フレームワーク拡張機能がユーザーの代わりに Neuron Runtime の設定を行います。ただし、より最適化された動作を実現するために、明示的な設定も可能です。

よく使われる環境変数は NEURON_RT_NUM_CORES と NEURON_RT_VISIBLE_CORES です。これらの環境変数を使用すると、Python プロセスを NeuronCore に関連付けることができます。NEURON_RT_NUM_CORES では、指定された数のコアをプロセスに予約することができ、NEURON_RT_VISIBLE_CORES では、NeuronCore の範囲を予約することができます。例えば、NEURON_RT_NUM_CORES=2 myapp.py とすると、myapp.py に2つのコアが予約され、NEURON_RT_VISIBLE_CORES='0-2' myapp.py とすると、myapp.py に0、1、2の3つのコアが予約されます。NeuronCore はデバイス（AWS Inferentia チップ）間でも予約することができます。したがって、NEURON_RT_VISIBLE_CORES='0-5' myapp.py とすると、EC2 Inf1 インスタンスタイプの device1 の最初の4つのコアと device2 の1つのコアが予約されます。同様に、EC2 Inf2インスタンスタイプでは、この設定により device1 と device2 の2つのコアが予約され、device3 に1つのコアが予約されます。以下の表はこれらの変数の設定をまとめたものです。

すべての環境変数の一覧については、Neuron Runtime Configuration を参照してください。

デフォルトでは、モデルをロードする際に、モデルは明示的に上記の環境変数で指定されていない限り、NeuronCore 0にロードされ、次に NeuronCore 1 にロードされます。先述したように、NeuronCore は利用可能なホスト vCPU とシステムメモリを共有します。したがって、各 NeuronCore に展開されたモデルは利用可能なリソースを競合します。モデルが NeuronCore を大幅に利用している場合は、これは問題になりません。しかし、モデルが NeuronCore の一部をのみ使用し、残りをホスト vCPU で実行している場合は、NeuronCore ごとのCPUの利用可能性を考慮することが重要になります。これはインスタンスの選択にも影響を与えます。

次の表は、各 NeuronCore に 1 つのモデルを展開した場合のホスト vCPU とシステムメモリの数を示しています。アプリケーションの NeuronCore の使用状況、vCPU の使用状況、およびメモリの使用状況に応じて、どの設定がアプリケーションにとって最もパフォーマンスが高いかを確認するためにテストを実行することが推奨されています。Neuron Top ツールを使用すると、コアの利用状況やデバイスおよびホストメモリの利用状況を視覚化するのに役立ちます。これらのメトリックに基づいて適切な判断を行うことができます。このブログの最後で Neuron Top の使用方法を示します。

自分で Neuron SDK の機能をテストしてみるには、最新の PyTorch 向け Neuron 機能をチェックしてください。

システムセットアップ

このソリューションで使用されたシステムのセットアップは以下になります。

ソリューションのセットアップ

ソリューションをセットアップするために必要な手順がいくつかあります。まず、Amazon Elastic Container Registry からの push と pull を許可する IAM ロールを作成して、EC2 インスタンスがそれを担えるようにします。

ステップ1：IAMロールのセットアップ

注意：このポリシーのアタッチ EC2 インスタンスは Amazon ECR へのフルアクセスが可能になります。本番環境のワークロードに対しては、最小特権の原則に従うことを強く推奨します。

ステップ2：AWS CLI のセットアップ

上記で指定された Deep Learning AMI を使用している場合、AWS CLI がインストールされています。別のAMI（Amazon Linux 2023, Base Ubuntuなど）を使用している場合は、このガイドに従って CLI ツールをインストールしてください。

CLI ツールがインストールされたら、 aws configure コマンドを使用して CLI を設定します。アクセスキーがある場合はここに追加できますが、AWS サービスとのやり取りには必ずしも必要ありません。今回の工程には IAM ロールを使用しています。

注意：default プロファイルを作成するために少なくとも1つの値（default region または default format）を入力する必要があります。この例では、region に us-east-2 を、default format にjsonを設定します。

GitHubリポジトリをクローンする

GitHub リポジトリには、FastAPI を使用して AWS Inferentia インスタンスの NeuronCore 上にモデルを展開するために必要なスクリプトが提供されています。この例では、再利用可能なソリューションを作成するために Docker コンテナを使用しており、ユーザーが入力を提供するための config.properties ファイルが含まれています。

# Docker Image and Container Name
docker_image_name_prefix=<Docker image name>
docker_container_name_prefix=<Docker container name>

# Deployment Setup
path_to_traced_models=<Path to traced model>
compiled_model=<Compiled model file name>
num_cores=<Number of NeuronCores to Deploy a Model Server>
num_models_per_server=<Number of Models to Be Loaded Per Server>

設定ファイルには、Docker イメージ と Docker コンテナのためのユーザーが定義するプレフィックスが必要です。fastapi フォルダと trace-model フォルダの build.sh スクリプトでは、これを使用して Docker イメージを作成します。

AWS Inferentia に向けたモデルコンパイル

モデルをトレースして、PyTorch Torchscript の .pt 形式ファイルを生成することから始めます。まず、trace-model ディレクトリにアクセスして、.env ファイルを変更します。選択したインスタンスのタイプに応じて、.env ファイル内のCHIP_TYPE を変更します。例として、Inf2 を選択しますが、同じ手順は Inf1 のデプロイプロセスでも適用されます。

次に、同じファイル内で default region を設定します。この region は ECR リポジトリの作成に使用され、Docker イメージがこのリポジトリにプッシュされます。また、このフォルダには、AWS Inferentia 上で bert-base-uncased モデルをトレースするために必要なすべてのスクリプトが含まれています。このスクリプトは、Hugging Face で利用可能なほとんどのモデルに使用できます。Dockerfile には、Neuron でモデルを実行するためのすべての依存関係が含まれており、trace-model.py コードがエントリーポイントとして実行されます。

Neuron コンパイルについて

Neuron SDK の API は、PyTorch の Python API に非常に近いです。PyTorch の torch.jit.trace() は、モデルとサンプルとなる入力テンソルを引数として取ります。サンプル入力はモデルに引き渡され、モデルのレイヤーを通過する際に呼び出される操作が TorchScript として記録されます。PyTorch における JIT トレースについての詳細は、以下のドキュメントを参照してください。

torch.jit.trace() と同様に、Inf1 の場合は以下のコードを使用して AWS Inferentia 上でモデルをコンパイルできるかどうかを確認できます。

import torch_neuron
model_traced = torch.neuron.trace(model, 
                                  example_inputs,
                                  compiler_args = 
                                  [‘--fast-math’, ‘fp32-cast-matmul’,
                                   ‘--neuron-core-pipeline-cores’,’1’],
                         optimizations=[torch_neuron.Optimization.FLOAT32_TO_FLOAT16])

Inf2 の場合、ライブラリは torch_neuronx と呼ばれます。以下は、モデルのコンパイルを Inf2 インスタンスにおいて確認する方法です。

import torch
import torch_neuronx
model_traced = torch.neuronx.trace(model, 
                                   example_inputs,
                                   compiler_args = 
                                   [‘--fast-math’, ‘fp32-cast-matmul’,
                                    ‘--neuron-core-pipeline-cores’,’1’],
          optimizations=[torch_neuronx.Optimization.FLOAT32_TO_FLOAT16])

モデルのトレースを行うと、次のようにしてサンプル入力となるテンソルを入力することができます。

answer_logits = model_traced(*example_inputs)

最後に生成された TorchScript をローカルディスクに保存します。

model_traced.save('./compiled-model-bs-{batch_size}.pt')

前述のコードに示されているように、compiler_args と optimizations を使用しデプロイを最適化することができます。torch.neuron.trace API の引数の詳細なリストについては、PyTorch-Neuron trace python API を参照してください。

また、次の重要なポイントに注意して下さい。

ビルドプロセス

build.sh を実行することでコンテナをビルドします。ビルドスクリプトは、ベースとなる Deep Learning Container Image を取得し、HuggingFace の transformers パッケージをインストールし Docker イメージを作成する処理を行います。 .env ファイルで指定された CHIP_TYPE に基づいて、docker.properties ファイルが適切な BASE_IMAGE を決定します。このBASE_IMAGE は、AWSが提供する Neuron Runtime 用の Deep Learning Container Image を指します。

これらの処理はプライベートな ECR リポジトリを利用します。そのため、イメージを取得する前にログインを行い、一時的なAWSクレデンシャルを取得する必要があります。

aws ecr get-login-password --region <region> | docker login --username AWS --password-stdin 763104351884.dkr.ecr.<region>.amazonaws.com

注意：<region> フラグで指定されたコマンドおよびリポジトリ URI 内のリージョンは、.env ファイルに記載されたリージョンで置き換える必要があります。

このプロセスを簡易に実行するために、fetch-credentials.sh ファイルを使用することができます。リージョンは自動的に .env ファイルから取得されます。

次に、 push.sh スクリプトを使用してイメージをプッシュします。このスクリプトは Amazon ECR にリポジトリを作成し、コンテナイメージをプッシュします。

最後に、イメージがビルドされてプッシュされたら、run.sh を実行してコンテナとして実行し、logs.sh を使用して実行中のログを表示できます。コンパイラのログ（以下のスクリーンショットを参照）では、Neuron でコンパイルされた算術演算子の割合と、モデルサブグラフの割合が表示されます。スクリーンショットは、bert-base-uncased-squad2 モデルのログを示しています。ログによれば、算術演算子の 95.64% がコンパイルされ、Neuron でコンパイルされた演算子とサポートされていない演算子のリストも示されています。

最新の PyTorch Neuron パッケージでサポートされているすべての演算子のリストはこちらにあります。同様に、最新のPyTorch Neuronx パッケージでサポートされているすべての演算子のリストもこちらです。

FastAPI を使用してモデルをデプロイする

モデルがコンパイルされた後、トレースされたモデルは trace-model フォルダに保存されます。この例では、バッチサイズ 1 としてトレースモデルを配置しています。ここでは、より大きいバッチサイズが実現不可能または必要ないユースケースを考えており、バッチサイズ 1 を使用しています。より大きいバッチサイズが必要なユースケースでは、torch.neuron.DataParallel (Inf1 の場合) または torch.neuronx.DataParallel (Inf2 の場合) APIが役立つ場合があります。

fast-api フォルダには、FastAPI を使用してモデルをデプロイするためのすべての必要なスクリプトが含まれています。変更なしでモデルをデプロイするには、単に deploy.sh スクリプトを実行するだけで、FastAPI コンテナイメージがビルド、指定されたコア数でコンテナが実行され、各 FastAPI モデルサーバーで指定された数のモデルがデプロイされます。このフォルダには .env ファイルも含まれており、これを変更して正しい CHIP_TYPE と AWS_DEFAULT_REGION を反映させることができます。

注意：FastAPI スクリプトは、イメージのビルド、プッシュ、コンテナとしての実行に使用されるものと同じ環境変数に依存しています。FastAPI デプロイメントスクリプトは、これらの変数からの最後に知られている値を使用します。したがって、もし最後にトレースしたモデルが Inf1 インスタンスタイプの場合、そのモデルがこれらのスクリプトを介してデプロイされます。

fastapi-server.py ファイルは、サーバーのホスティングとモデルへのリクエストの送信を担当しており、以下のことを行います。

このセットアップでは、各 NeuronCore に格納されているモデルとその数を一覧表示するAPIを簡単に設定することができます。同様に、特定の NeuronCore からモデルを削除するための API も作成できます。

FastAPI コンテナをビルドするため Dockerfile は、モデルのトレース用にビルドした Docker イメージをベースにして構築されています。これが、docker.properties ファイルがモデルのトレース用の Docker イメージの ECR パスを指定している理由です。今回のセットアップでは、すべての NeuronCore 上の Docker コンテナは同種のものであり、1つのイメージをビルドし、それから複数のコンテナを実行できます。エントリーポイントのエラーを回避するために、 Dockerfile で ENTRYPOINT ["/usr/bin/env"] を指定してから、hypercorn fastapi-server:app -b 0.0.0.0:8080 というスタートアップシェルスクリプトを実行します。このスタートアップスクリプトはすべてのコンテナで同じです。モデルのトレース用と同じベースイメージを使用している場合、build.sh スクリプトを実行するだけでこのコンテナをビルドできます。push.sh スクリプトは、モデルのトレース用と同じままです。変更された Docker イメージとコンテナ名は、docker.properties ファイルで提供されます。

run.sh ファイルは以下の操作を行います。

NeuronCore 0 で Inf1 用に展開するための Docker 実行コマンドは次のようになります。

docker run -t -d \
	    --name $ bert-inf-fastapi-nc-0 \
	    --env NEURON_RT_VISIBLE_CORES="0-0" \
	    --env CHIP_TYPE="inf1" \
	    -p ${port_num}:8080 --device=/dev/neuron0 ${registry}/ bert-inf-fastapi

NeuronCore 5 で展開するための実行コマンドは次のようになります。

docker run -t -d \
	    --name $ bert-inf-fastapi-nc-5 \
	    --env NEURON_RT_VISIBLE_CORES="5-5" \
	    --env CHIP_TYPE="inf1" \
	    -p ${port_num}:8080 --device=/dev/neuron0 ${registry}/ bert-inf-fastapi

コンテナが展開された後、 run_apis.py スクリプトを使用して、API を並列スレッドで呼び出します。コードは、各 NeuronCore に1つずつ展開された6つのモデルを呼び出すように設定されていますが、簡単に異なる設定に変更できます。クライアント側から API を呼び出す方法は以下の通りです。

import requests

url_template = http://localhost:%i/predictions_neuron_core_%i/model_%i

# NeuronCore 0
response = requests.get(url_template % (8081,0,0))

# NeuronCore 5
response = requests.get(url_template % (8086,5,0))

NeuronCore の監視

モデルサーバーが展開された後、NeuronCore の利用状況を監視するために、neuron-top を使用して各 NeuronCore の利用率をリアルタイムで観測することができます。neuron-top は、Neuron SDK 内の CLI ツールで、NeuronCore、vCPU、メモリの利用状況などの情報を提供します。別のターミナルで、次のコマンドを入力します。

neuron-top

このコマンドによって、下の画像のような出力が得られます。このシナリオでは、Inf2.xlarge インスタンスでサーバーごとに2つの NeuronCore と 2 つのモデルを使用するように指定しています。以下のスクリーンショットでは、2 つのサイズが 287.8 MB のモデルが 2 つの NeuronCore にロードされていることが表示されています。合計 4 つのモデルがロードされているため、使用されているデバイスメモリは 1.3 GB です。矢印キーを使用して異なるデバイスの NeuronCore 間を移動できます。

同様に、Inf1.6xlarge インスタンスタイプでは、合計 12 のモデル（ 1 つのコアあたり 2 つのモデル、6 つのコア合計）がロードされていることが表示されます。合計 2.1 GB のメモリが消費され、各モデルのサイズは 177.2 MB です。

run_apis.py スクリプトを実行した後、各 NeuronCore の利用率の割合を確認できます（以下のスクリーンショットを参照）。また、システムの vCPU 使用率とランタイムの vCPU 使用率も確認できます。

以下のスクリーンショットは、Inf2 インスタンスのコア使用率の割合を示しています。

同様に、このスクリーンショットは、Inf1.6xlarge インスタンスタイプにおけるコアの利用率を示しています。

クリーンアップ

作成したすべての Docker コンテナをクリーンアップするために、すべての実行中および停止したコンテナを削除する cleanup.sh スクリプトを提供しています。このスクリプトはすべてのコンテナを削除しますので、実行中の一部のコンテナを保持したい場合には使用しないでください。

結論

プロダクションのワークロードはしばしば高いスループット、低レイテンシ、コストの要件を持っています。アクセラレータを最適でない方法で使用する非効率なアーキテクチャでは、不必要に高いコストにつながる可能性があります。この記事では、 FastAPI と NeuronCore を最適に活用し、スループットを最大化しながら最小のレイテンシで動作すること方法を示しました。詳しくは GitHubのリポジトリに手順を公開しています。このソリューションを使用すると、各 NeuronCore に複数のモデルを展開し、異なる NeuronCore 上で複数のモデルを並列に運用することができ、パフォーマンスを損なうことなく行えます。Amazon Elastic Kubernetes Service （Amazon EKS）などのサービスを使用してモデルをスケールで展開する方法については、「AWS Inferentiaを使用して Amazon EKS で 3,000種類のディープラーニングモデルを 1 時間あたり 50 USD 以下で提供」の記事を参照してください。

翻訳はソリューションアーキテクトの松崎が担当しました。原文はこちらです。