Cloud TPU v5e 推論の概要
概要と利点
Cloud TPU v5e は、Google が開発した AI アクセラレータで、トランスフォーマー ベース、text-to-image、および CNN ベースのトレーニング、微調整、サービス提供(推論)のために最適化されています。TPU v5e スライスには、最大 256 個のチップを含めることができます。
サービス提供とは、トレーニング済み機械学習モデルを本番環境にデプロイするプロセスであり、本番環境を使用して推論に使用できます。レイテンシ SLO は提供の優先事項です。
このドキュメントでは、単一ホストの TPU でモデルを提供する方法について説明します。8 チップ以下の TPU スライスには 1 つの TPU VM またはホストがあり、単一ホストの TPU と呼ばれます。
使ってみる
v5e TPU には割り当てが必要です。オンデマンド TPU には tpu-v5s-litepod-serving
割り当てが必要です。予約済みの TPU には tpu-v5s-litepod-serving-reserved
割り当てが必要です。詳細については、Cloud セールスまでお問い合わせください。
Cloud TPU を使用するには、Google Cloud アカウントとプロジェクトが必要です。詳細については、Cloud TPU 環境を設定するをご覧ください。
キューに格納されたリソースを使用して、v5e TPU をプロビジョニングします。サービス提供に使用可能な v5e 構成の詳細については、サービス提供用の Cloud TPU v5e タイプをご覧ください。
Cloud TPU モデルの推論とサービス提供
推論のためにモデルを提供する方法は、モデルが書き込まれた ML フレームワークによって異なります。TPU v5e は、JAX、TensorFlow、PyTorch で作成されたモデルの提供をサポートしています。
JAX モデルの推論とサービス提供
TPU VM でモデルを提供するには、次の操作を行う必要があります。
- TensorFlow SavedModel 形式でモデルをシリアル化する
- Inference Converter を使用して、保存したモデルをサービス提供用に準備する
- TensorFlow Serving を使用してモデルを提供する
SavedModel の形式
SavedModel には、(トレーニング済みパラメータと計算を含む)完全な TensorFlow プログラムが含まれます。元のモデル構築コードを実行する必要はありません。
モデルが JAX で作成されている場合、SavedModel 形式でモデルをシリアル化するには、jax2tf
を使用する必要があります。
Inference Converter
Cloud TPU Inference Converter は、TPU 推論用に SavedModel 形式でエクスポートされたモデルを準備して最適化します。Inference Converter は、ローカルシェルまたは TPU VM で実行できます。 コンバータの実行に必要なすべてのコマンドライン ツールが含まれているため、TPU VM シェルの使用をおすすめします。Inference Converter の詳細については、Inference Converter ユーザーガイドをご覧ください。
Inference Converter の要件
モデルは、TensorFlow または JAX から SavedModel 形式でエクスポートする必要があります。
TPU 関数の関数エイリアスを定義する必要があります。詳細については、Inference Converter ユーザーガイドをご覧ください。このガイドの例では、TPU 関数のエイリアスとして
tpu_func
を使用します。Tensorflow ライブラリ(Cloud TPU Inference Converter の依存関係)は AVX 命令を使用するようにコンパイルされるため、マシンの CPU が Advanced Vector eXtensions(AVX)命令をサポートしていることを確認してください。ほとんどの CPU が AVX をサポートします。
JAX モデルの推論とサービス提供
このセクションでは、jax2tf
と TensorFlow Serving を使用して JAX モデルを提供する方法について説明します。
jax2tf
を使用してモデルを SavedModel 形式にシリアル化する- Inference Converter を使用して、サービス提供用の保存済みモデルを準備する
- TensorFlow Serving を使用してモデルを提供する
jax2tf
を使用して JAX モデルを SavedModel 形式にシリアル化する
次の Python 関数は、モデルコード内で jax2tf
を使用する方法を示しています。
# Inference function
def model_jax(params, inputs):
return params[0] + params[1] * inputs
# Wrap the parameter constants as tf.Variables; this will signal to the model
# saving code to save those constants as variables, separate from the
# computation graph.
params_vars = tf.nest.map_structure(tf.Variable, params)
# Build the prediction function by closing over the `params_vars`. If you
# instead were to close over `params` your SavedModel would have no variables
# and the parameters will be included in the function graph.
prediction_tf = lambda inputs: jax2tf.convert(model_jax)(params_vars, inputs)
my_model = tf.Module()
# Tell the model saver what the variables are.
my_model._variables = tf.nest.flatten(params_vars)
my_model.f = tf.function(prediction_tf, jit_compile=True, autograph=False)
tf.saved_model.save(my_model)
jax2tf
の詳細については、JAX と Cloud TPU の相互運用をご覧ください。
Inference Converter を使用して、保存したモデルをサービス提供用に準備する
Inference Converter の使用方法については、Inference Converter ガイドをご覧ください。
TensorFlow Serving を使用する
TensorFlow Serving の使用方法については、TensorFlow Serving をご覧ください。
JAX モデルのサービス提供例
前提条件
Docker の認証情報を設定し、Inference Converter と Cloud TPU Serving Docker イメージを pull します。
sudo usermod -a -G docker ${USER} newgrp docker gcloud auth configure-docker \ us-docker.pkg.dev docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0 docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
SSH を使用して TPU VM に接続し、推論のデモコードをインストールします。
gsutil -m cp -r \ "gs://cloud-tpu-inference-public/demo" \ .
JAX デモの依存関係をインストールします。
pip install -r ./demo/jax/requirements.txt
推論用に JAX BERT モデルを提供する
事前トレーニング済みの BERT モデルを Hugging Face からダウンロードできます。
Flax BERT モデルから TPU 互換 TensorFlow 保存済みモデルをエクスポートします。
cd demo/jax/bert python3 export_bert_model.py
Cloud TPU モデルサーバー コンテナを起動します。
docker run -t --rm --privileged -d \ -p 8500:8500 -p 8501:8501 \ --mount type=bind,source=/tmp/jax/bert_tpu,target=/models/bert \ -e MODEL_NAME=bert \ us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
コンテナの開始から約 30 秒後に、モデルサーバーのコンテナログを調べて、gRPC と HTTP サーバーが稼働していることを確認します。
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker logs ${CONTAINER_ID}
次の情報で終わるログエントリがある場合、サーバーはリクエストを処理する準備ができています。
2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ... [warn] getaddrinfo: address family for nodename not supported 2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ... [evhttp_server.cc : 245] NET_LOG: Entering the event loop ...
推論リクエストをモデルサーバーに送信します。
python3 bert_request.py
出力は次のようになります。
For input "The capital of France is [MASK].", the result is ". the capital of france is paris.." For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."
クリーンアップする。
他のデモを実行する前に、必ず Docker コンテナをクリーンアップしてください。
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker stop ${CONTAINER_ID}
モデルのアーティファクトをクリーンアップします。
sudo rm -rf /tmp/jax/
推論用に JAX Stable Diffusion を提供する
Hugging Face から事前トレーニング済みの Stable Diffusion モデルをダウンロードできます。
Stable Diffusion モデルを、TPU 互換の TF2 保存済みモデル形式でダウンロードします。
cd demo/jax/stable_diffusion python3 export_stable_diffusion_model.py
モデル用の Cloud TPU モデルサーバー コンテナを起動します。
docker run -t --rm --privileged -d \ -p 8500:8500 -p 8501:8501 \ --mount type=bind,source=/tmp/jax/stable_diffusion_tpu,target=/models/stable_diffusion \ -e MODEL_NAME=stable_diffusion \ us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
約 2 分後、モデルサーバーのコンテナログを調べて、gRPC サーバーと HTTP サーバーが実行されていることを確認します。
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker logs ${CONTAINER_ID}
次の情報で終わるログが表示されている場合、サーバーはリクエストを処理する準備ができています。
2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ... [warn] getaddrinfo: address family for nodename not supported 2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ... [evhttp_server.cc : 245] NET_LOG: Entering the event loop ...
モデルサーバーにリクエストを送信します。
python3 stable_diffusion_request.py
このスクリプトは、プロンプトとして「Painting of a squirrel squerrel in New York」でプロンプトを送信します。出力画像は、現在のディレクトリに
stable_diffusion_images.jpg
として保存されます。クリーンアップする。
他のデモを実行する前に、必ず Docker コンテナをクリーンアップしてください。
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker stop ${CONTAINER_ID}
モデル アーティファクトをクリーンアップする
sudo rm -rf /tmp/jax/
TensorFlow Serving
TPU VM で TensorFlow モデルを提供する手順は次のとおりです。
TensorFlow サービス提供ワークフロー
TPU VM の TensorFlow Serving Docker イメージをダウンロードします。
サンプル環境変数を設定する
export YOUR_LOCAL_MODEL_PATH=model-path export MODEL_NAME=model-name # Note: this image name may change later. export IMAGE_NAME=us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
Docker イメージをダウンロードする
docker pull ${IMAGE_NAME}
Docker の認証情報を設定し、Inference Converter と TensorFlow Serving Docker イメージを pull します。
sudo usermod -a -G docker ${USER} newgrp docker gcloud auth configure-docker \ us-docker.pkg.dev docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0 docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
デモコードをダウンロードします。
gsutil -m cp -r \ "gs://cloud-tpu-inference-public/demo" \ .
TensorFlow デモの依存関係をインストールします。
pip install -r ./demo/tf/requirements.txt
TPU VM で TensorFlow Serving Docker イメージを使用して TensorFlow モデルを提供します。
# PORT 8500 is for gRPC model server and 8501 is for HTTP/REST model server. docker run -t --rm --privileged -d \ -p 8500:8500 -p 8501:8501 \ --mount type=bind,source=${YOUR_LOCAL_MODEL_PATH},target=/models/${MODEL_NAME} \ -e MODEL_NAME=${MODEL_NAME} \ ${IMAGE_NAME}
Serving Client API を使用してモデルをクエリします。
TensorFlow ResNet-50 サービス提供デモを実行する
Keras ResNet-50 モデルから TPU 互換の TF2 保存済みモデルをエクスポートします。
cd demo/tf/resnet-50 python3 export_resnet_model.py
モデルの TensorFlow モデルサーバー コンテナを起動します。
docker run -t --rm --privileged -d \ -p 8500:8500 -p 8501:8501 \ --mount type=bind,source=/tmp/tf/resnet_tpu,target=/models/resnet \ -e MODEL_NAME=resnet \ us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
モデルサーバーのコンテナログを調べて、gRPC と HTTP Server が稼働していることを確認します。
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker logs ${CONTAINER_ID}
次の情報で終わるログが表示されている場合、サーバーはリクエストを処理する準備ができています。これには約 30 秒かかります。
2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ... [warn] getaddrinfo: address family for nodename not supported 2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ... [evhttp_server.cc : 245] NET_LOG: Entering the event loop ...
リクエストをモデルサーバーに送信します。
リクエスト画像は、https://i.imgur.com/j9xCCzn.jpeg からのバナナです。
python3 resnet_request.py
出力は次のようになります。
Predict result: [[('n07753592', 'banana', 0.94921875), ('n03532672', 'hook', 0.022338867), ('n07749582', 'lemon', 0.005126953)]]
クリーンアップする。
他のデモを実行する前に、必ず Docker コンテナをクリーンアップしてください。
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker stop ${CONTAINER_ID}
モデルのアーティファクトをクリーンアップします。
sudo rm -rf /tmp/tf/
PyTorch モデルの推論とサービス提供
PyTorch で作成されたモデルの場合、ワークフローは次のとおりです。
TorchDynamo
と PyTorch/XLA を使用して読み込みと推論を行う Python モデルハンドラを作成するTorchModelArchiver
を使用してモデル アーカイブを作成するTorchServe
を使用してモデルを提供する
TorchDynamo と PyTorch/XLA
TorchDynamo(Dynamo)は、PyTorch プログラムを高速化するために設計された Python レベルの JIT コンパイラです。コンパイラ バックエンドが接続するためのクリーンな API を提供します。実行直前に Python バイトコードを動的に変更します。PyTorch/XLA 2.0 リリースには、Dynamo を使用して推論とトレーニングを行うための試験運用版のバックエンドがあります。
Dynamo は、モデルパターンを認識するときに Torch FX(FX)グラフを提供し、PyTorch/XLA は lazy tensor アプローチを使用して FX グラフをコンパイルし、コンパイル済みの関数を返します。Dynamo の詳細については、以下をご覧ください。
torch.compile
を使用して densenet161 推論を実行する小規模なコード例を以下に示します。
import torch
import torchvision
import torch_xla.core.xla_model as xm
def eval_model(loader):
device = xm.xla_device()
xla_densenet161 = torchvision.models.densenet161().to(device)
xla_densenet161.eval()
dynamo_densenet161 = torch.compile(
xla_densenet161, backend='torchxla_trace_once')
for data, _ in loader:
output = dynamo_densenet161(data)
TorchServe
付属の torchserve-tpu
Docker イメージを使用して、アーカイブされた pytorch モデルを TPU VM で提供できます。
Docker の認証を設定します。
sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
us-docker.pkg.dev
Cloud TPU TorchServe Docker イメージを TPU VM に pull します。
CLOUD_TPU_TORCHSERVE_IMAGE_URL=us-docker.pkg.dev/cloud-tpu-images/inference/torchserve-tpu:v0.9.0-2.1
docker pull ${CLOUD_TPU_TORCHSERVE_IMAGE_URL}
モデル アーティファクトを収集する
まず、モデルハンドラを提供する必要があります。このハンドラは、TorchServe モデルサーバー ワーカーに、モデルを読み込み、入力データを処理して推論を実行するよう指示します。また、TorchServe デフォルトの推論ハンドラ(ソース)、または base_handler.py に従って独自のカスタムモデル ハンドラを開発します。 トレーニング済みモデルとモデル定義ファイルの提供も必要になる場合があります。
次の Densenet 161 の例では、モデル アーティファクトと、TorchServe が提供するデフォルトの画像分類ハンドラを使用します。
いくつかの環境変数を構成します。
CWD="$(pwd)" WORKDIR="${CWD}/densenet_161" mkdir -p ${WORKDIR}/model-store mkdir -p ${WORKDIR}/logs
TorchServe 画像分類サンプルからモデル アーティファクトをダウンロードしてコピーします。
git clone https://github.com/pytorch/serve.git cp ${CWD}/serve/examples/image_classifier/densenet_161/model.py ${WORKDIR} cp ${CWD}/serve/examples/image_classifier/index_to_name.json ${WORKDIR}
モデルの重みをダウンロードします。
wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth mv densenet161-8d451a50.pth ${WORKDIR}
Dynamo バックエンドを使用する TorchServe モデル構成ファイルを作成します。
echo 'pt2: "torchxla_trace_once"' >> ${WORKDIR}/model_config.yaml
次のファイルとディレクトリが表示されます。
>> ls ${WORKDIR} model_config.yaml index_to_name.json logs model.py densenet161-8d451a50.pth model-store
モデル アーカイブ ファイルを生成する
Cloud TPU TorchServe で PyTorch モデルを提供するには、Torch Model Archiver を使用してモデルハンドラとすべてのモデル アーティファクトをモデル アーカイブ ファイル (*.mar)
にパッケージ化する必要があります。
torch-model-archiver を使用してモデル アーカイブ ファイルを生成します。
MODEL_NAME=Densenet161
docker run \
--privileged \
--shm-size 16G \
--name torch-model-archiver \
-it \
-d \
--rm \
--mount type=bind,source=${WORKDIR},target=/home/model-server/ \
${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
torch-model-archiver \
--model-name ${MODEL_NAME} \
--version 1.0 \
--model-file model.py \
--serialized-file densenet161-8d451a50.pth \
--handler image_classifier \
--export-path model-store \
--extra-files index_to_name.json \
--config-file model_config.yaml
model-store ディレクトリにモデル アーカイブ ファイルが生成されます。
>> ls ${WORKDIR}/model-store
Densenet161.mar
推論リクエストを処理する
モデル アーカイブ ファイルが作成されたので、TorchServe モデルサーバーを起動して推論リクエストを処理できます。
TorchServe モデルサーバーを起動します。
docker run \ --privileged \ --shm-size 16G \ --name torchserve-tpu \ -it \ -d \ --rm \ -p 7070:7070 \ -p 7071:7071 \ -p 8080:8080 \ -p 8081:8081 \ -p 8082:8082 \ -p 9001:9001 \ -p 9012:9012 \ --mount type=bind,source=${WORKDIR}/model-store,target=/home/model-server/model-store \ --mount type=bind,source=${WORKDIR}/logs,target=/home/model-server/logs \ ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \ torchserve \ --start \ --ncs \ --models ${MODEL_NAME}.mar \ --ts-config /home/model-server/config.properties
モデルサーバーの健全性をクエリします。
curl http://localhost:8080/ping
モデルサーバーが稼働している場合は、次のように表示されます。
{ "status": "Healthy" }
現在登録されているモデルのデフォルト バージョンをクエリするには、次のコマンドを使用します。
curl http://localhost:8081/models
登録済みモデルが表示されます。
{ "models": [ { "modelName": "Densenet161", "modelUrl": "Densenet161.mar" } ] }
推論用に画像をダウンロードするには、次のコマンドを使用します。
curl -O https://raw.githubusercontent.com/pytorch/serve/master/docs/images/kitten_small.jpg mv kitten_small.jpg ${WORKDIR}
モデルサーバーに推論リクエストを送信するには、次のコマンドを使用します。
curl http://localhost:8080/predictions/${MODEL_NAME} -T ${WORKDIR}/kitten_small.jpg
次のようなレスポンスが表示されます。
{ "tabby": 0.47878125309944153, "lynx": 0.20393909513950348, "tiger_cat": 0.16572578251361847, "tiger": 0.061157409101724625, "Egyptian_cat": 0.04997897148132324 }
モデルサーバーのログ
ログにアクセスするには、次のコマンドを使用します。
ls ${WORKDIR}/logs/ cat ${WORKDIR}/logs/model_log.log
ログに次のメッセージが表示されます。
"Compiled model with backend torchxla\_trace\_once"
クリーンアップ
Docker コンテナを停止します。
rm -rf serve
rm -rf ${WORKDIR}
docker stop torch-model-archiver
docker stop torchserve-tpu
プロファイリング
推論を設定したら、プロファイラを使用してパフォーマンスと TPU の使用率を分析できます。プロファイリングの詳細については、次をご覧ください。