TPU リソースを管理する

このページでは、以下を使用して Cloud TPU リソースを管理する方法について説明します。

Cloud TPU には、TPU ノードと TPU VM の 2 つの VM アーキテクチャがあります。この 2 つの VM アーキテクチャについては、システム アーキテクチャをご覧ください。 このドキュメントで説明する gcloud コマンドは、両方の TPU 構成で使用できます。使用する gcloud コマンドは、使用する TPU 構成によって異なります。各 gcloud コマンドは、タブ付きのセクションに表示されます。使用する TPU 構成のタブを選択すると、対応する gcloud コマンドがウェブページに表示されます。TPU ノードを使用する必要がある場合を除き、TPU VM を使用することをおすすめします。Cloud TPU v4 以降では、TPU VM アーキテクチャのみがサポートされています。

前提条件

以下の手順を実行する前に、Google Cloud CLI をインストールし、Google Cloud プロジェクトを作成して、Cloud TPU API を有効にする必要があります。手順については、プロジェクトを設定して Cloud TPU API を有効にするをご覧ください。

Google Cloud CLI を使用している場合は、Google Cloud Shell または Compute Engine VM を使用するか、Google Cloud CLI をローカルにインストールできます。Google Cloud Shell を使用すると、ソフトウェアをインストールせずに Cloud TPU を操作できます。非アクティブな状態が一定期間続くと Google Cloud Shell の接続が解除される可能性があります。長時間実行コマンドを実行する場合は、ローカルマシンに Google Cloud CLI をインストールすることをおすすめします。Google Cloud CLI の詳細については、gcloud リファレンスをご覧ください。

Cloud TPU をプロビジョニングする

Cloud TPU をプロビジョニングするには、gcloud、Google Cloud コンソール、または Cloud TPU API を使用します。

gcloud を使用して TPU をプロビジョニングする方法は 2 つあります。

  1. キューに格納されたリソースの使用: gcloud alpha compute tpus queued-resources create
  2. Create Node API の使用: gcloud alpha compute tpus tpu-vm create

キューに格納されたリソースを使用してプロビジョニングすることをおすすめします。キューに格納されたリソースをリクエストすると、リクエストは Cloud TPU サービスによって管理されるキューに追加されます。リクエストされたリソースが使用可能になると、すぐに Google Cloud プロジェクトに割り当てられ、すぐに専用で使用できるようになります。

キューに格納されたリソースを使用して TPU を作成するには、キューに格納されたリソースをご覧ください。

マルチスライスを使用している場合は、マルチスライスの概要をご覧ください。

マルチスライスを使用する場合は、キューに格納されたリソースをリクエストするときに、次の追加パラメータを指定します。

export NODE_COUNT=node_count
export NODE_PREFIX=your_tpu_prefix # Optional

Create Node API を使用して Cloud TPU を作成する

Create Node API を使用して TPU を作成するには、gcloud compute tpus tpu-vm create コマンドを実行します。

TPU VM

TensorCore または TPU チップの観点から TPU の構成を指定できます。 詳細については、システム アーキテクチャで使用している TPU バージョンのセクションをご覧ください。

次のコマンドは、TensorCore ベースの構成を使用します。

$ gcloud alpha compute tpus tpu-vm create tpu-name \
  --zone=us-central2-b \
  --accelerator-type=v4-8 \
  --version=tpu-software-version

コマンドフラグの説明

zone
Cloud TPU を作成するゾーン
accelerator-type
作成する Cloud TPU のタイプ
version
TPU ソフトウェアのバージョン
shielded-secure-boot(オプション)
セキュアブートを有効にして TPU インスタンスを作成することを指定します。 これにより、暗黙的に Shielded VM インスタンスにされます。 Shielded VM とは をご確認ください。

次のコマンドは、Create Node API を使用して TPU を作成します。

  $ gcloud alpha compute tpus tpu-vm create tpu-name \
    --zone=us-central2-b \
    --type=v4 \
    --topology=2x2x1 \
    --version=tpu-software-version

必須フラグ

tpu-name
作成する TPU VM の名前。
zone
Cloud TPU を作成するゾーン
type
サポートされている TPU タイプの詳細については、TPU タイプをご覧ください。
topology
サポートされているトポロジについては、トポロジのセクションをご覧ください。
version
使用する TPU ソフトウェア バージョン。詳細については、TPU ソフトウェアのバージョンをご覧ください。

サポートされている TPU のタイプとトポロジの詳細については、サポートされている TPU の構成をご覧ください。

TPU ノード

$ gcloud compute tpus execution-groups create --name=tpu-name \
  --zone=us-central1-a \
  --tf-version=2.12.0 \
  --machine-type=n1-standard-1 \
  --accelerator-type=v3-8

コマンドフラグの説明

zone
Cloud TPU を作成するゾーン

tf-version
gcloud コマンドが VM にインストールする TensorFlow のバージョン。

machine-type
作成する Compute Engine VM のマシンタイプ

accelerator-type
作成する Cloud TPU のタイプ

Google Cloud コンソールでの Cloud TPU の作成

  1. Google Cloud Console に移動します。
  2. 左側のナビゲーション メニューから、[Compute Engine] > [TPU] を選択します。
  3. [Create TPU Node] をクリックします。
  4. [名前] ボックスに TPU インスタンス名を入力します。
  5. [ゾーン] ボックスで、TPU を作成するゾーンを選択します。
  6. [TPU 設定] で、[TPU VM アーキテクチャ] または [TPU ノード アーキテクチャ] を選択します。TPU の構成に応じて、TPU を TPU VM として作成するか、TPU ノードとして作成するかが決まります。詳細については、システム アーキテクチャをご覧ください。
  7. [TPU タイプ] で、作成する TPU タイプを選択します。
  8. [TPU ソフトウェア バージョン] で、ソフトウェア バージョンを選択します。Cloud TPU VM の作成時には、インストールされる TPU ランタイム バージョンが TPU ソフトウェア バージョンによって指定されます。Cloud TPU ノードを作成するときに、TPU ソフトウェア バージョンを使用すると、ノードの VM にインストールされている ML フレームワークを選択できます。その他の設定は必要ありません。詳細については、サポートされているモデルをご覧ください。
  9. [作成] をクリックしてリソースを作成します。

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

次のコマンドは、curl を使用して TPU を作成します。

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
  runtime_version:'tpu-vm-tf-2.16.1-pjrt', \
  network_config: {enable_external_ips: true}, \
  shielded_instance_config: { enable_secure_boot: true }}" \
  https://tpu.googleapis.com/v2/projects/project-id/locations/us-central2-b/nodes?node_id=node_name

必須項目

runtime_version
使用する Cloud TPU ランタイムのバージョン。
project
: Google Cloud プロジェクトの名前。
zone
Cloud TPU を作成するゾーン
node_name
作成する TPU VM の名前。

起動スクリプトを実行する

TPU VM の作成時に --metadata startup-script パラメータを指定すると、各 TPU VM で起動スクリプトを実行できます。次のコマンドは、起動スクリプトを使用して TPU VM を作成します。

$ gcloud compute tpus tpu-vm create tpu-name \
    --zone=us-central2-b \
    --accelerator-type=tpu-type \
    --version=tpu-vm-tf-2.16.1-pjrt \
    --metadata startup-script='#! /bin/bash
      pip3 install numpy
      EOF'

Cloud TPU への接続

SSH を使用して TPU に接続できます。

TPU VM

TPU VM を使用する場合は、SSH を使用して明示的に TPU VM に接続する必要があります。

  • gcloud compute tpus tpu-vm ssh コマンドを使用して、SSH を介して TPU VM に接続します。

    4 チップを超えるスライスをリクエストすると、Cloud TPU は 4 つのチップのグループごとに TPU VM を作成します。

    バイナリのインストールやコードを実行するには、tpu-vm ssh command を使用して各 TPU VM に接続します。

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}
    

    特定の TPU VM にアクセスするか、SSH を使用して各 TPU VM にバイナリをインストールするには、0 ベースのインデックスに続けて --worker フラグを使用します。

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} --worker=1
    

    複数の TPU VM がある場合は、--worker=all フラグと --command フラグを使用して、すべての TPU VM で同時にコマンドを実行します。 次に例を示します。

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
    --zone  ${ZONE} \
    --worker=all \
    --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
    

    マルチスライスの場合、(列挙型 tpu-name を使用して)単一の VM でコマンドを実行するか、--node=all--worker=all、または --command フラグと、オプションの --batch-size フィールドを使用して、マルチスライスのすべてのスライスのすべての TPU VM でコマンドを実行できます。

    gcloud alpha compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
    --project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
    --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html' --batch-size=4
    
    $ gcloud compute tpus tpu-vm ssh tpu-name --zone=zone
  • 次の手順で SSH-in-browser を使用します。

    1. Google Cloud コンソール で、[TPU] ページに移動します。

      TPU に移動

    2. TPU VM のリストで、接続する TPU VM の行にある [SSH] をクリックします。

TPU ノード

デフォルトでは、TPU ノードの作成に使用する gcloud コマンドは、TPU ノードに自動的に SSH 接続を試みます。TPU ノードを使用していて、gcloud コマンドで Compute Engine インスタンスに接続していない場合は、次のコマンドを実行して接続できます。

$ gcloud compute ssh tpu-name \
  --zone=zone

TPU VM を作成したら、SSH を使用して TPU VM に接続し、次のコマンドを実行します。これにより、起動スクリプトからログを表示できます。

$ cat /var/log/syslog | grep startup-script

Cloud TPU リソースの一覧表示

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

gcloud を使用した Cloud TPU リソースの一覧表示

使用するコマンドは、TPU VM を使用しているか、TPU ノードを使用しているかによって異なります。詳細については、システム アーキテクチャをご覧ください。

TPU VM

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

TPU ノード

$ gcloud compute tpus execution-groups list --zone=zone

このコマンドは、指定されたゾーン内の Cloud TPU リソースを一覧表示します。設定されたリソースがない場合、出力では VM と TPU にダッシュだけが表示されます。 一方のリソースがアクティブで、もう一方がアクティブでなければ、正常なステータスではないことを伝えるメッセージが表示されます。この場合、実行されていないリソースの起動または再起動が必要になります。

Google Cloud コンソールでの Cloud TPU リソースの一覧表示

  1. Google Cloud コンソールに移動します。

  2. 左側のナビゲーション メニューから、[Compute Engine] > [TPU] を選択します。 コンソールに TPU ページが表示されます。

Cloud TPU に関する情報の取得

指定した Cloud TPU に関する情報を取得できます。

gcloud を使用して Cloud TPU に関する情報を取得する

使用するコマンドは、TPU VM を使用しているか、TPU ノードを使用しているかによって異なります。詳細については、システム アーキテクチャをご覧ください。

TPU VM

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

TPU ノード

$ gcloud compute tpus execution-groups describe tpu-name \
  --zone=zone

Google Cloud コンソールで Cloud TPU に関する情報を取得する

  1. Google Cloud コンソールに移動します。
  2. 左側のナビゲーション メニューから、[Compute Engine] > [TPU] を選択します。 コンソールに TPU ページが表示されます。
  3. Cloud TPU の名前をクリックします。Cloud TPU の詳細ページが表示されます。

Cloud TPU リソースの停止

単一の Cloud TPU を停止することで、VM の構成とソフトウェアを失わずに課金を停止できます。TPU Pod またはキューに格納されたリソース API を介して割り当てられた TPU の停止はサポートされていません。キューに格納されたリソース API を介して割り当てられた TPU に対する課金を停止するには、TPU を削除する必要があります。

gcloud を使用した Cloud TPU の停止

Cloud TPU の停止に使用するコマンドは、TPU VM と TPU ノードのどちらを使用しているかによって異なります。詳細については、システム アーキテクチャをご覧ください。

TPU VM

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

TPU ノード

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

Google Cloud コンソールで Cloud TPU を停止する

  1. Google Cloud コンソールに移動します。

  2. 左側のナビゲーション メニューから、[Compute Engine] > [TPU] を選択します。 コンソールに TPU ページが表示されます。

  3. Cloud TPU の横にあるチェックボックスをオンにして、[停止] をクリックします。

Cloud TPU リソースの起動

停止した Cloud TPU は再開できます。

gcloud を使用した Cloud TPU の起動

停止された Cloud TPU を起動して、使用を再開できます。

停止した Cloud TPU の開始に使用するコマンドは、TPU VM と TPU ノードのどちらを使用しているかによって異なります。詳細については、システム アーキテクチャをご覧ください。

TPU VM

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

TPU ノード

$ gcloud compute tpus start tpu-name --zone=zone

Google Cloud コンソールで Cloud TPU を起動する

  1. Google Cloud コンソールに移動します。

  2. 左側のナビゲーション メニューから、[Compute Engine] > [TPU] を選択します。 コンソールに TPU ページが表示されます。

  3. Cloud TPU の横にあるチェックボックスをオンにして、[開始] をクリックします。

TPU VM を削除する

使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。 詳細については、システム アーキテクチャをご覧ください。

TPU VM

セッションの終了時に TPU VM スライスを削除します。

 gcloud compute tpus tpu-vm delete ${TPU_NAME} --project=${PROJECT_ID}
 --zone=${ZONE} --quiet

コマンドフラグの説明

zone
Cloud TPU を削除するゾーン

TPU ノード

$ gcloud compute tpus execution-groups delete tpu-name \
  --zone=zone

コマンドフラグの説明

zone
Cloud TPU を削除するゾーン

Google Cloud コンソールでの Cloud TPU の削除

  1. Google Cloud コンソールに移動します。

  2. 左側のナビゲーション メニューから、[Compute Engine] > [TPU] を選択します。 コンソールに TPU ページが表示されます。

  3. Cloud TPU の横にあるチェックボックスをオンにして、[削除] をクリックします。

高度な設定

カスタム ネットワーク リソース

TPU を作成するときに、ネットワークかサブネットワークのいずれかまたは両方を指定できます。これを行うには、gcloud コマンドまたは curl 呼び出しを送信します。

gcloud CLI でネットワークまたはサブネットワークを指定するには、次のコマンドを使用します。

--network [NETWORK] --subnetwork [SUBNETWORK]

curl 呼び出しでネットワークまたはサブネットワークを指定するには、次のコマンドを使用します。

network_config: {network: '[NETWORK]', subnet: '[SUBNETWORK]', enable_external_ips: true}

ネットワーク

TPU に使用するネットワークを指定することもできます。指定しない場合は、default ネットワークが使用されます。

有効なネットワークのフォーマット:

https://www.googleapis.com/compute/{version}/projects/{proj-id}/global/networks/{network}
compute/{version}/projects/{proj-id}/global/networks/{network}
compute/{version}/projects/{proj-##}/global/networks/{network}
projects/{proj-id}/global/networks/{network}
projects/{proj-##}/global/networks/{network}
global/networks/{network}
{network}

サブネットワーク

特定のサブネットワークを使用するように、サブネットワークを指定できます。指定するサブネットワークは、TPU を実行するゾーンと同じリージョンに存在する必要があります。

有効な形式:

https://www.googleapis.com/compute/{version}/projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
compute/{version}/projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
compute/{version}/projects/{proj-##}/regions/{region}/subnetworks/{subnetwork}
projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
projects/{proj-##}/regions/{region}/subnetworks/{subnetwork}
regions/{region}/subnetworks/{subnetwork}
{subnetwork}

限定公開の Google アクセス

TPU VM に SSH 接続するには、TPU VM のアクセス構成を追加するか、TPU VM が接続されたサブネットワークの限定公開の Google アクセスを有効にする必要があります。

アクセス構成を追加するには、enable_external_ips を設定する必要があります。TPU を作成すると、enable_external_ips がデフォルトで設定されます。無効にするには、次のコマンドを入力します。

--internal-ips

または、curl 呼び出しを使用します。

network_config: {enable_external_ips: true}

限定公開の Google アクセスを構成したら、SSH 経由で VM に接続します。

カスタム サービス アカウント

各 TPU VM には、API リクエストを行うために使用するサービス アカウントが関連付けられています。TPU VM は、このサービス アカウントを使用して Cloud TPU API を呼び出し、Cloud Storage やその他のサービスにアクセスします。デフォルトでは、TPU VM はデフォルトの Compute Engine サービス アカウントを使用します。

--service-account フラグを使用して TPU VM を作成するときに、カスタム サービス アカウントを指定できます。サービス アカウントは、TPU VM を作成する Google Cloud プロジェクトと同じプロジェクトで定義する必要があります。TPU VM に使用されるカスタム サービス アカウントには、Cloud TPU API を呼び出すための TPU 閲覧者のロールが必要です。TPU VM で実行されるコードが他の Google Cloud サービスを呼び出す場合は、これらのサービスにアクセスするために必要なロールが必要です。

TPU を作成するときに、--service-account フラグを使用してカスタム サービス アカウントを指定できます。サービス アカウントの詳細については、サービス アカウントをご覧ください。

次のコマンドを使用して、カスタム サービス アカウントを指定します。

gcloud CLI を使用して TPU VM を作成する

$ gcloud alpha compute tpus tpu-vm create tpu-name \
    --zone=us-central2-b \
    --accelerator-type=tpu-type \
    --version=tpu-vm-tf-2.16.1-pjrt \
    --service-account=your-service-account

curl を使用して TPU VM を作成する

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
  runtime_version:'tpu-vm-tf-2.16.1-pjrt', \
  network_config: {enable_external_ips: true}, \
  shielded_instance_config: { enable_secure_boot: true }}" \
  service_account: {email: 'your-service-account'} \
  https://tpu.googleapis.com/v2/projects/project-id/locations/us-central2-b/nodes?node_id=node_name

カスタム サービス アカウントを使用するには、Google Cloud Storage バケット用のサービス アカウントを承認する必要があります。詳細については、Cloud Storage バケットへの接続をご覧ください。

カスタム VM SSH メソッド

  1. SSH 用にファイアウォールを設定します。

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

    $ gcloud CLI compute firewall-rules create \
    --network=network allow-ssh \
    --allow=tcp:22
    
  2. TPU VM に SSH 接続します。

    $ gcloud compute tpus tpu-vm ssh tpu-name \
    --zone=us-central2-b \
    --project=project-id
    

    必須フィールド

    • tpu-name: TPU ノードの名前
    • zone: TPU ノードのロケーション。現在、us-central2-b だけがサポートされています。
    • project-id: 上記で作成したプロジェクト。

    省略可能なフィールドの一覧については、gcloud API のドキュメントをご覧ください。