Cloud TPU VM ユーザーガイド

このガイドでは、Cloud TPU VM を使用するための Google Cloud Platform プロジェクトの設定方法について説明します。Cloud TPU VM を使用する際のコマンドと、Cloud TPU VM の使用を開始する際に発生する可能性がある一般的な問題の解決策について説明します。

Cloud TPU VM は、TPU ホストマシン（Cloud TPU デバイスに接続しているマシン）で実行され、TPU を扱う際のパフォーマンスとユーザビリティを大幅に改善します。

Cloud TPU を初めて使用する場合は、TPU 入門ガイドをご覧ください。

Cloud TPU Pod で実行する予定の場合は、フレームワークに応じた Pod のトレーニング ドキュメントをご覧ください。

• TensorFlow

• JAX ポッド

• PyTorch XLA ユーザーガイド

Cloud TPU VM では新しい Cloud TPU アーキテクチャが導入されています。Cloud TPU アーキテクチャの詳細については、システム アーキテクチャをご覧ください。

GCP プロジェクトの設定

重要: gcloud alpha コマンドを実行する前に、Google Cloud CLI をインストールする必要があります。Google Cloud CLI のインストールの手順に従って、Google Cloud CLI をインストールします。必要に応じて、ディストリビューション固有の手順ではなく、Linux のインストール手順を使用します。

Google Cloud CLI をインストールしたら、次のコマンドを使用して gcloud の alpha コンポーネントをインストールします。

gcloud components install alpha

gcloud コンポーネントの詳細については、Google Cloud CLI コンポーネントの管理をご覧ください。

GCP プロジェクトを準備する

Google アカウントにログインします。Google アカウントをまだお持ちでない場合は、新しいアカウントを登録します。Google Cloud コンソールで、プロジェクト セレクタページから Cloud プロジェクトを選択するか作成します。プロジェクトで課金が有効になっていることを確認します。Cloud Shell で gcloud を使用してプロジェクト ID を設定します。プロジェクト ID は、Cloud Console に表示されるプロジェクトの名前です。

$ gcloud config set project project-id

Cloud TPU API を有効にする

Cloud Shell で、次の gcloud コマンドを使用して Cloud TPU API を有効にします（Google Cloud Console から有効にすることもできます）。

$ gcloud services enable tpu.googleapis.com

gcloud コマンドを設定する

次のコマンドを実行して、GCP プロジェクトを使用するように gcloud を構成し、TPU VM プレビュー版の必要なコンポーネントをインストールします。

$ gcloud config set account your-email-account
$ gcloud config set project your-project

TPU の管理

Cloud TPU VM は、gcloud か curl を使用して操作します。詳細については、Cloud TPU の操作をご覧ください。

gcloud を使用した Cloud TPU VM の作成

$ gcloud alpha compute tpus tpu-vm create tpu-name \
  --zone=zone \
  --accelerator-type=v3-8 \
  --version=tpu-vm-tf-2.8.0

必須フィールド

zone

Cloud TPU を作成するゾーン。

accelerator-type

作成する Cloud TPU のタイプ。

version

Cloud TPU ランタイムのバージョン。

TPU VM を作成するときに、--metadata startup-script フラグを使用して起動スクリプトを指定できます。次に例を示します。

$ gcloud alpha compute tpus tpu-vm create tpu-name \
--zone=zone \
--accelerator-type=v3-8 \
--version=tpu-vm-tf-2.8.0 \
--metadata startup-script=your-script

起動スクリプトは、TPU VM がプロビジョニングされるたびにだけでなく、メンテナンス イベントによって TPU VM が再起動されても実行されます。

curl を使用した Cloud TPU VM の作成

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" -d "{accelerator_type: 'v2-8', \
  runtime_version:'[runtime_version]', \
  network_config: {enable_external_ips: true}}" \
  https://tpu.googleapis.com/v2alpha1/projects/[project]/locations/[zone]/nodes?node_id=[node_name]

必須フィールド

runtime_version

使用するランタイムのバージョン。

project

登録済みの GCP プロジェクトの名前。

zone

Cloud TPU を作成するゾーン。

node_name

作成する TPU VM の名前。

Cloud TPU VM への接続

1. （オプション）SSH 用にファイアウォールを設定します。

デフォルトのネットワークは、すべての VM への SSH アクセスを許可するようにあらかじめ構成されています。デフォルトのネットワークを使用しない場合や、デフォルトのネットワーク設定が編集されている場合は、ファイアウォール ルールを追加して、SSH アクセスを明示的に有効にする必要があります。

$ gcloud compute firewall-rules create --network=network allow-ssh --allow=tcp:22

1. TPU VM に SSH 接続します。

$ gcloud alpha compute tpus tpu-vm ssh tpu-name --zone zone --project project-id

必須フィールド

tpu_name

接続する TPU VM の名前。

zone

Cloud TPU を作成するゾーン。 現在は、europe-west4-a と us-central1-a がサポートされています。

project-id

GCP プロジェクト ID。

省略可能フィールド

user

SSH 経由で Cloud TPU VM に接続する際は、TPU 名に $USER@ 接頭辞を使用して、認証するユーザー名を選択できます（例: my-email-account@tpu-node-1）。

worker

Cloud TPU Pod の場合、SSH 接続するワーカー VM を選択できます。デフォルトはワーカー 0 で、TPU Pod に関連付けられた最初の VM です。

ssh-key-file

SSH 認証鍵ファイルのパスデフォルトでは、~/.ssh/google_compute_engine です。

internal-ip

内部 IP アドレスを使用して TPU VM に接続します。この接続を機能させるには、接続する TPU VM の内部 IP アドレスへの SSH 接続を許可するようにネットワークとファイアウォールを構成する必要があります。

command

TPU VM で実行するコマンド。このコマンドはターゲット TPU VM 上で実行され、終了します。

tunnel-through-iap

TCP 転送用に Cloud Identity-Aware Proxy を介して SSH 接続をトンネルします。詳しくは、TCP 転送の概要をご覧ください。

TPU Pod に関連付けられた他の TPU VM に SSH 接続するには、コマンドに --worker ${WORKER_NUMBER} を追加します。ここで、WORKER_NUMBER は、0 から始まるインデックスです。

Cloud TPU リソースの一覧表示

Cloud TPU は、指定したゾーン内にあるすべてを一覧表示できます。

$ gcloud alpha compute tpus tpu-vm list --zone=zone

必須フィールド

zone

Cloud TPU を作成するゾーン。

このコマンドは、指定されたゾーン内の Cloud TPU リソースを一覧表示します。設定されたリソースがない場合、出力では VM と TPU にダッシュだけが表示されます。

Cloud TPU に関する情報の取得

特定の Cloud TPU に関する情報は、次のコマンドを使用して取得できます。

$ gcloud alpha compute tpus tpu-vm describe tpu-name \
  --zone=zone

必須フィールド

tpu-name

作成する Cloud TPU の名前。

zone

Cloud TPU が作成されたゾーン。

Cloud TPU リソースの停止

単一の Cloud TPU は、次のコマンドを使用して停止できます。TPU Pod は停止できません。

$ gcloud alpha compute tpus tpu-vm stop tpu-name \
  --zone=zone

必須フィールド

tpu-name

停止する Cloud TPU の名前。

zone

Cloud TPU を作成したゾーン。

Cloud TPU リソースの起動

Cloud TPU が停止した場合は、次のコマンドを使用して再起動できます。

$ gcloud alpha compute tpus tpu-vm start tpu-name --zone  zone

コマンドフラグの説明

tpu-name

起動する Cloud TPU の名前。

zone

Cloud TPU が作成されたゾーン。

VM と Cloud TPU リソースの削除

Cloud TPU は、使用が終了すると削除できます。

$ gcloud alpha compute tpus tpu-vm delete tpu-name \
  --zone=zone

必須フィールド

zone

Cloud TPU が作成されたゾーン。

パフォーマンス指標の取得

パフォーマンス プロファイルは、コマンドライン スクリプトまたは TensorBoard を使用して取得できます。TensorBoard のインストール手順については、TensorBoard の設定をご覧ください。

TensorFlow モデルの場合、標準の TensorFlow プロファイリング コールバック メソッドを使用してプロファイル データを自動的に取得できます。

TensorFlow モデルのプロファイル データを手動で取得するには、TPU VM で次のコマンドを使用します。

重要: 次のコマンドでは、port を TensorBoard サーバーがリッスンするポートに置き換えます。

$ python3 -c "import tensorflow as tf; tf.profiler.experimental.client.trace('grpc://localhost:port', 'gs://model-dir', 1000)"

コマンドラインを使用して PyTorch モデルのプロファイル データを取得するには、TPU VM で次のコマンドを使用します。

$ python3 -c "import torch_xla.debug.profiler as xp; xp.trace('localhost:port', '/tmp/tb', 1000)"

JAX モデルのプロファイル データを取得する方法については、JAX プログラムのプロファイリングをご覧ください。

プロファイル データの表示

1. Cloud Shell を開きます。
2. TensorBoard がインストールされていることを確認します。
3. TensorBoard を実行する
4. Cloud Shell で、[ウェブでプレビュー] ボタンをクリックし、[ポートを変更] を選択して 6006 と入力します。
5. [プロフィール] をクリックして、概要ページを表示します。
6. ツールの [Trace Viewer] に移動します。

TPU の割り当ての追加をリクエストする

Cloud TPU は、デフォルトの TPU 割り当てをプロジェクトに割り当てます。さらに割り当てが必要な場合は、追加の割り当てをリクエストするをご覧ください。

プロファイル情報の生成と表示

トレーニング指標は、プロファイル情報を生成して TensorBoard を使用することで可視化できます。

gcloud の設定のトラブルシューティング

問題

gcloud components update を実行すると、次のエラー メッセージが表示される。

ERROR: (gcloud.components.update)
You cannot perform this action because the gcloud CLI component manager
is disabled for this installation.

解決策

TPU VM で gcloud を使用するには、パッケージ マネージャーで管理されない gcloud インストールを使用する必要があります。ソースコードから gcloud をインストールするには、次の手順を行います。

sudo apt-get remove google-cloud-sdk
curl -O https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
tar -xzf google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
./google-cloud-sdk/install.sh
source ~/.bashrc

問題

gcloud alpha compute tpus tpu-vm で始まるコマンドを実行すると、次のエラー メッセージが表示される。

ERROR: (gcloud.alpha.compute.tpus) Invalid choice: 'tpu-vm'.

解決策

この問題は、コンポーネント リポジトリが適切に更新されていない場合に発生します。これを確認するには、gcloud --version を実行します。出力の最初の行は「Google Cloud CLI HEAD」です。出力が異なる場合は、更新は行われていません。その場合は、次のコマンドを使用して gcloud コンポーネントの更新を試してみてください。

gcloud components update

引き続き同じエラーが発生する場合は、次のコマンドを使用して gcloud を再インストールしてみてください。

gcloud components reinstall

問題

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} コマンドを実行すると、次のエラー メッセージが表示される。

Waiting for SSH key to propagate.
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ERROR: (gcloud.alpha.compute.tpus.tpu-vm.ssh) Could not SSH into the instance.  It is possible that your SSH key has not propagated to the instance yet. Try running this command again.  If you still cannot connect, verify that the firewall and instance are set to accept ssh traffic.

解決策

SSH 認証鍵の伝播に問題がある可能性があります。自動生成された鍵をバックアップの場所に移動して、gcloud で強制的に再作成してみてください。

mv ~/.ssh/google_compute_engine ~/.ssh/old-google_compute_engine
mv ~/.ssh/google_compute_engine.pub ~/.ssh/old-google_compute_engine.pub

よくある質問

V1Alpha1 API と V1 API を使用して Cloud TPU VM を操作できますか？

Get / List は使用できますが、ミューテーションは V2Alpha1 API バージョンでのみ使用できます。

TPU で Cloud TPU VM が使用されているかどうかを確認するにはどうすればよいですか？

TPU で、APIVersion フィールドを V2_ALPHA1 に設定して、GetNode を呼び出します。