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 マシンタイプに対しては時間単位で課金されます。詳細については、予測料金をご覧ください。
次のステップ
- オンライン予測を取得する方法を確認する。