Cloud TPU VM にモデルをデプロイする

Google Cloud では、Tensor Processing Unit（TPU）と呼ばれるカスタム設計の ML アクセラレータにアクセスできます。TPU は、ML モデルのトレーニングと推論を高速化するように最適化されており、自然言語処理、コンピュータビジョン、音声認識などのさまざまなアプリケーションに最適です。

このページでは、Vertex AI でオンライン予測を行うため、モデルを単一のホストの Cloud TPU v5e にデプロイする方法について説明します。

Cloud TPU バージョン v5e のみがサポートされています。他の世代の Cloud TPU はサポートされていません。

モデルをインポートする

Cloud TPU にデプロイするには、モデルを Vertex AI にインポートし、次のいずれかのコンテナを使用するように構成する必要があります。

構築済みの最適化された TensorFlow ランタイムコンテナ（nightly バージョンまたはバージョン 2.15 以降）
ビルド済みの PyTorch TPU コンテナバージョン 2.1 以降
TPU をサポートする独自のカスタムコンテナ

構築済みの最適化された TensorFlow ランタイムコンテナ

Cloud TPU で TensorFlow SavedModel をインポートして実行するには、モデルを TPU に対して最適化する必要があります。TensorFlow SavedModel が TPU に対して最適化されていない場合は、次の 3 つの方法でモデルを最適化できます。

手動によるモデル最適化 - 推論コンバータを使用してモデルを最適化し、保存します。次に、モデルを upload するときに --saved_model_tags='serve,tpu' フラグと --disable_optimizer=true フラグを渡す必要があります。次に例を示します。

model = aiplatform.Model.upload(
    display_name='Manually optimized model',
    artifact_uri="gs://model-artifact-uri",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
    serving_container_args=[
        "--saved_model_tags=serve,tpu",
        "--disable_optimizer=true",
    ]
)

自動パーティショニングによる自動モデル最適化 - モデルをインポートすると、Vertex AI は自動パーティショニングアルゴリズムを使用して最適化されていないモデルの最適化を試みます。この最適化は、すべてのモデルで機能するわけではありません。最適化が失敗した場合は、モデルを手動で最適化するか、手動パーティショニングによる自動モデル最適化を選択する必要があります。次に例を示します。
```
model = aiplatform.Model.upload(
    display_name='TPU optimized model with automatic partitioning',
    artifact_uri="gs://model-artifact-uri",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
    serving_container_args=[
    ]
)
```

手動パーティショニングによるモデルの自動最適化。--converter_options_string フラグを指定し、ニーズに合わせて ConverterOptions.TpuFunction を調整します。例については、コンバータイメージをご覧ください。手動パーティショニングに必要な ConverterOptions.TpuFunction のみサポートされています。次に例を示します。

model = aiplatform.Model.upload(
display_name='TPU optimized model with manual partitioning',
  artifact_uri="gs://model-artifact-uri",
  serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
  serving_container_args=[
      "--converter_options_string='tpu_functions { function_alias: \"partitioning function name\" }'"
  ]
)

モデルのインポートの詳細については、Vertex AI へのモデルのインポートをご覧ください。

ビルド済みの PyTorch コンテナ

Cloud TPU に PyTorch モデルをインポートして実行する手順は、PyTorch モデルをインポートして実行する手順と同じです。

たとえば、Cloud TPU v5e 推論の TorchServe は、Torch Model Archiver を使用して Densenet 161 モデルをモデルアーティファクトにパッケージ化する方法を示しています。

次に、モデルアーティファクトを Cloud Storage フォルダにアップロードし、次に示すようにモデルをアップロードします。

model = aiplatform.Model.upload(
    display_name='DenseNet TPU model from SDK PyTorch 2.1',
    artifact_uri="gs://model-artifact-uri",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/pytorch-tpu.2-1:latest",
    serving_container_args=[],
    serving_container_predict_route="/predictions/model",
    serving_container_health_route="/ping",
    serving_container_ports=[8080]
)

詳細については、PyTorch のモデルアーティファクトをエクスポートすると Jupyter ノートブックでビルド済みコンテナを使用して PyTorch モデルを提供するをご覧ください。

カスタムコンテナ

カスタムコンテナの場合、モデルは TensorFlow モデルである必要はありませんが、TPU 向けに最適化されている必要があります。TPU 向けに最適化されたモデルの生成については、一般的な ML フレームワークに関する次のガイドをご覧ください。

Cloud TPU v5e での JAX、TensorFlow、PyTorch を使用してトレーニングされたモデルの提供については、Cloud TPU v5e 推論をご覧ください。

カスタムコンテナがカスタムコンテナの要件を満たしていることを確認します。

ドライバが直接メモリアクセス（DMA）を介して TPU チップと通信できるように、ロックされたメモリの上限を引き上げる必要があります。次に例を示します。

コマンドライン

ulimit -l 68719476736

Python

import resource

resource.setrlimit(
    resource.RLIMIT_MEMLOCK,
    (
        68_719_476_736_000,  # soft limit
        68_719_476_736_000,  # hard limit
    ),
  )

カスタムコンテナを使用したモデルのインポートについては、予測にカスタムコンテナを使用するをご覧ください。前処理または後処理のロジックを実装する場合は、カスタム予測ルーチンの使用を検討してください。

エンドポイントを作成する

Cloud TPU のエンドポイントを作成する手順は、エンドポイントを作成する手順と同じです。

たとえば、次のコマンドは endpoint リソースを作成します。

endpoint = aiplatform.Endpoint.create(display_name='My endpoint')

レスポンスには、新しいエンドポイントの ID が含まれます。この ID は、次の手順で使用します。

エンドポイントの作成の詳細については、エンドポイントにモデルをデプロイするをご覧ください。

モデルのデプロイ

Cloud TPU にモデルをデプロイする手順は、サポートされている次の Cloud TPU マシンタイプのいずれかを指定することを除き、任意のモデルをデプロイする手順と同じです。

マシンタイプ	TPU チップの数
`ct5lp-hightpu-1t`	1
`ct5lp-hightpu-4t`	4
`ct5lp-hightpu-8t`	8

TPU アクセラレータはマシンタイプに組み込まれています。アクセラレータタイプやアクセラレータ数を指定する必要はありません。

たとえば、次のコマンドは deployModel を呼び出してモデルをデプロイします。

machine_type = 'ct5lp-hightpu-1t'

deployed_model = model.deploy(
    endpoint=endpoint,
    deployed_model_display_name='My deployed model',
    machine_type=machine_type,
    traffic_percentage=100,
    min_replica_count=1
    sync=True,
)

詳細については、エンドポイントにモデルをデプロイするをご覧ください。

オンライン予測を取得する

Cloud TPU からオンライン予測を取得する手順は、オンライン予測を取得する手順と同じです。

たとえば、次のコマンドは predict を呼び出してオンライン予測リクエストを送信します。

deployed_model.predict(...)

カスタムコンテナについては、カスタムコンテナの予測リクエストとレスポンスの要件をご覧ください。

容量の確保

デフォルトでは、Custom model serving TPU v5e cores per region の割り当ては 0 です。

割り当ての増加をリクエストするには、割り当て上限の引き上げをリクエストするをご覧ください。

料金

Vertex Prediction の他のすべてのマシンタイプと同様に、TPU マシンタイプに対しては時間単位で課金されます。詳細については、予測料金をご覧ください。

次のステップ

オンライン予測を取得する方法を確認する。