Google Cloud では、Tensor Processing Unit(TPU)と呼ばれるカスタム設計の ML アクセラレータにアクセスできます。TPU は、ML モデルのトレーニングと推論を高速化するように最適化されており、自然言語処理、コンピュータ ビジョン、音声認識などのさまざまなアプリケーションに最適です。
このページでは、Vertex AI でのオンライン推論のために、モデルを単一のホストの Cloud TPU v5e または v6e にデプロイする方法について説明します。
Cloud TPU バージョン v5e と v6e のみがサポートされています。他の世代の Cloud TPU はサポートされていません。
Cloud TPU バージョン v5e と v6e を使用できるロケーションについては、ロケーションをご覧ください。
モデルをインポートする
Cloud TPU にデプロイするには、モデルを Vertex AI にインポートし、次のいずれかのコンテナを使用するように構成する必要があります。
- 構築済みの最適化された TensorFlow ランタイム コンテナ(
nightly
バージョンまたはバージョン2.15
以降) - ビルド済みの PyTorch TPU コンテナバージョン
2.1
以降 - TPU をサポートする独自のカスタム コンテナ
構築済みの最適化された TensorFlow ランタイム コンテナ
Cloud TPU で TensorFlow SavedModel
をインポートして実行するには、モデルを TPU に対して最適化する必要があります。TensorFlow SavedModel
が TPU に対して最適化されていない場合は、モデルを自動的に最適化できます。これを行うには、モデルをインポートします。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=[
]
)
モデルのインポートの詳細については、Vertex AI へのモデルのインポートをご覧ください。
ビルド済みの PyTorch コンテナ
Cloud TPU に PyTorch モデルをインポートして実行する手順は、PyTorch モデルをインポートして実行する手順と同じです。
たとえば、Cloud TPU v5e 推論の TorchServe は、Torch Model Archiver を使用して Densenet 161 モデルをモデル アーティファクトにパッケージ化する方法を示しています。
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 のモデル アーティファクトをエクスポートするとチュートリアル ノートブックでビルド済みコンテナを使用して 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 チップの数 |
---|---|
ct6e-standard-1t |
1 |
ct6e-standard-4t |
4 |
ct6e-standard-8t |
8 |
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(...)
カスタム コンテナについては、カスタム コンテナの推論リクエストとレスポンスの要件をご覧ください。
容量の確保
ほとんどのリージョンでは、TPU v5e and v6e cores per region
カスタムモデル サービングの割り当ては 0 です。一部のリージョンでは制限されています。
割り当ての増加をリクエストするには、割り当ての調整をリクエストするをご覧ください。
料金
Vertex Prediction の他のすべてのマシンタイプと同様に、TPU マシンタイプに対しては時間単位で課金されます。詳細については、予測料金をご覧ください。
次のステップ
- オンライン推論の取得方法を確認する。