Cloud TPU 上のソフトウェア バージョンの切り替え

概要

TPU で実行されているフレームワークのソフトウェア バージョンは、ローカル VM で実行されているバージョンと一致する必要があります。このソフトウェア バージョンは、TPU を削除して再作成せずに、実行中の Cloud TPU で切り替えられるようになりました。そのため、特定の夜間バージョンのソフトウェア フレームワークで Cloud TPU を構成することもできます。これらのフレームワークのサポートされているバージョンを選択することをおすすめします。

用途

バージョンを切り替えるには、cloud-tpu-client Python ライブラリを使用することをおすすめします。

TensorFlow の使用例

from cloud_tpu_client import Client
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--tpu-name',
                    type=str,
                    required=True,
                    help='Name of the TPU Instance')
parser.add_argument('--target-version',
                    type=str,
                    required=True,
                    help='target TPU Runtime version')

args = parser.parse_args()

c = Client(args.tpu_name)
c.configure_tpu_version(args.target_version, restart_type='ifNeeded')
c.wait_for_healthy()

これにより、ローカル VM で実行される TensorFlow バージョンと一致するように Cloud TPU が構成されます。これには、公式リリースと、日付が付加されたナイトリー ビルドが含まれます。

configure_tpu_version API の restart_type パラメータは、バージョンを切り替えるときの TPU 再起動動作を定義します。オプションは「always」（デフォルト）と「ifNeeded」です。

  • 「always」は、前回の実行でリソースが漏洩したことが原因で、ステータスが「UNHEALTHY_TENSORFLOW」になっている、またはメモリ不足（OOM）エラーを返す状態などの TPU を修正する場合に使用できます。このオプションを設定すると、新しいフレームワーク バージョンがインストールされていない場合でも TPU は再起動されます。

  • 「ifNeeded」は、すでに適切なバージョンに設定されている場合はランタイムが再起動されないため、トレーニング スクリプトに著しいスタートアップ時間が追加されないことから有用です。このオプションを設定すると、TPU は正しいフレームワーク バージョンがインストールされていない場合にのみ再起動されます。

このライブラリは Cloud TPU と直接通信します。そのため、このコードは同じネットワーク内の VM で実行する必要があります。残りのモデルのコード内で実行することをおすすめします。

その他のソフトウェア オプション

TensorFlow には、正しいバージョンを構成する最も簡単な方法である tf.__version__ 文字列が含まれています。その他のソフトウェア オプションは次のとおりです。

  • PyTorch - pytorch-1.13、pytorch-nightly-dev20220930、pytorch-nightly
  • Jax - tpu_driver、tpu_driver0.1-dev20200320、tpu_driver_nightly

たとえば、PyTorch の最新の夜間ビルドで実行するように TPU を構成します。

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()