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 環境を設定するをご覧ください。

注: TPU v5e を使用するには、gcloud alpha components をインストールする必要があります。

キューに格納されたリソースを使用して、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 で記述されている場合は、jax2tf を使用してモデルを SavedModel 形式でシリアル化する必要があります。

Inference Converter

Cloud TPU Inference Converter は、TPU 推論用に SavedModel 形式でエクスポートされたモデルを準備し、最適化します。推論コンバータは、ローカルシェルまたは 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 を使用して、保存したモデルをサービス提供用に準備する

推論コンバータの使用方法については、推論コンバータのガイドをご覧ください。

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 モデルを提供する

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

注: 現在、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. TPU 互換の TF2 保存済みモデル形式で Stable Diffusion モデルをダウンロードします。

   注: ダウンロードには 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 イメージを使用して、TPU VM 上でアーカイブ pytorch モデルを提供できます。

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 プロファイリング