Cloud TPU v5e 推論の概要

注: Cloud TPU を初めて使用する場合は、Cloud TPU の概要をご覧ください。

概要と利点

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 でモデルを提供するには、次の操作を行う必要があります。

1. TensorFlow SavedModel 形式でモデルをシリアル化する
2. Inference Converter を使用して、保存したモデルをサービス提供用に準備する
3. 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 の要件

1. モデルは、TensorFlow または JAX から SavedModel 形式でエクスポートする必要があります。

2. TPU 関数の関数エイリアスを定義する必要があります。詳細については、Inference Converter ユーザーガイドをご覧ください。このガイドの例では、TPU 関数のエイリアスとして tpu_func を使用します。

3. Tensorflow ライブラリ（Cloud TPU Inference Converter の依存関係）は AVX 命令を使用するようにコンパイルされるため、マシンの CPU が Advanced Vector eXtensions（AVX）命令をサポートしていることを確認してください。ほとんどの CPU が AVX をサポートします。

   注: lscpu | grep avx を実行して、AVX 命令セットがサポートされているかどうかを確認できます。

JAX モデルの推論とサービス提供

このセクションでは、jax2tf と TensorFlow Serving を使用して JAX モデルを提供する方法について説明します。

1. jax2tf を使用してモデルを SavedModel 形式にシリアル化する
2. Inference Converter を使用して、サービス提供用の保存済みモデルを準備する
3. 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 モデルのサービス提供例

前提条件

1. 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
   

2. SSH を使用して TPU VM に接続し、推論のデモコードをインストールします。

   gsutil -m cp -r \
   "gs://cloud-tpu-inference-public/demo" \
   .
   

3. JAX デモの依存関係をインストールします。

   pip install -r ./demo/jax/requirements.txt
   

推論用に JAX BERT モデルを提供する

事前トレーニング済みの BERT モデルを Hugging Face からダウンロードできます。

注: 現在、PyTorch の保存済みモデルのエクスポートはサポートされていません。

1. Flax BERT モデルから TPU 互換 TensorFlow 保存済みモデルをエクスポートします。

   cd demo/jax/bert
   python3 export_bert_model.py
   

2. 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 ...
   

3. 推論リクエストをモデルサーバーに送信します。

   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?."
   

4. クリーンアップする。

   他のデモを実行する前に、必ず 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 モデルをダウンロードできます。

1. Stable Diffusion モデルを、TPU 互換の TF2 保存済みモデル形式でダウンロードします。

   注: ダウンロードには 3 分ほどかかります。

   cd demo/jax/stable_diffusion
   python3 export_stable_diffusion_model.py
   

2. モデル用の 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 ...
   

3. モデルサーバーにリクエストを送信します。

   python3 stable_diffusion_request.py
   

   このスクリプトは、プロンプトとして「Painting of a squirrel squerrel in New York」でプロンプトを送信します。出力画像は、現在のディレクトリに stable_diffusion_images.jpg として保存されます。

4. クリーンアップする。

   他のデモを実行する前に、必ず 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 サービス提供ワークフロー

1. 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}
   

2. 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
   

3. デモコードをダウンロードします。

   gsutil -m cp -r \
   "gs://cloud-tpu-inference-public/demo" \
   .
   

4. TensorFlow デモの依存関係をインストールします。

   pip install -r ./demo/tf/requirements.txt
   

5. 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}
   

6. Serving Client API を使用してモデルをクエリします。

   • REST クライアント API
   • gRPC クライアント API

TensorFlow ResNet-50 サービス提供デモを実行する

1. Keras ResNet-50 モデルから TPU 互換の TF2 保存済みモデルをエクスポートします。

   cd demo/tf/resnet-50
   python3 export_resnet_model.py
   

2. モデルの 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 ...
   

3. リクエストをモデルサーバーに送信します。

   リクエスト画像は、https://i.imgur.com/j9xCCzn.jpeg からのバナナです。

   python3 resnet_request.py
   

   出力は次のようになります。

   Predict result: [[('n07753592', 'banana', 0.94921875), ('n03532672', 'hook', 0.022338867), ('n07749582', 'lemon', 0.005126953)]]
   

4. クリーンアップする。

   他のデモを実行する前に、必ず Docker コンテナをクリーンアップしてください。

   CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
   docker stop ${CONTAINER_ID}
   

   モデルのアーティファクトをクリーンアップします。

   sudo rm -rf /tmp/tf/
   

PyTorch モデルの推論とサービス提供

PyTorch で作成されたモデルの場合、ワークフローは次のとおりです。

1. TorchDynamo と PyTorch/XLA を使用して読み込みと推論を行う Python モデルハンドラを作成する
2. TorchModelArchiver を使用してモデル アーカイブを作成する
3. 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 の詳細については、以下をご覧ください。

• Pytorch デベロッパー ディスカッションの投稿
• TorchDynamo のドキュメント
• 詳しくは、PyTorch 2.0 と XLAをご覧ください

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 に従って独自のカスタムモデル ハンドラを開発します。 トレーニング済みモデルとモデル定義ファイルの提供も必要になる場合があります。

注: すべての TorchServe デフォルト モデル ハンドラは PyTorch/XLA をサポートします

次の Densenet 161 の例では、モデル アーティファクトと、TorchServe が提供するデフォルトの画像分類ハンドラを使用します。

• model.py
• index_to_name.json
• image_classifier.py

1. いくつかの環境変数を構成します。

   CWD="$(pwd)"
   
   WORKDIR="${CWD}/densenet_161"
   
   mkdir -p ${WORKDIR}/model-store
   mkdir -p ${WORKDIR}/logs
   

2. 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}
   

3. モデルの重みをダウンロードします。

   wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth
   
   mv densenet161-8d451a50.pth ${WORKDIR}
   

4. 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 モデルサーバーを起動して推論リクエストを処理できます。

1. 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
   

2. モデルサーバーの健全性をクエリします。

   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
   }
   

3. モデルサーバーのログ

   ログにアクセスするには、次のコマンドを使用します。

   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 の使用率を分析できます。プロファイリングの詳細については、次をご覧ください。

• Cloud TPU でのプロファイリング
• TensorFlow プロファイリング
• PyTorch プロファイリング
• JAX プロファイリング