GKE Standard に TPU ワークロードをデプロイする

Standard

このページでは、Google Kubernetes Engine（GKE）で Cloud TPU アクセラレータ（TPU）を使用するワークロードをリクエストしてデプロイする方法について説明します。

GKE で TPU ワークロードを構成してデプロイする前に、次のコンセプトを理解しておく必要があります。

始める前に

始める前に、次の作業が完了していることを確認してください。

Google Kubernetes Engine API を有効にする。

Google Kubernetes Engine API を有効化

このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。すでに gcloud CLI をインストールしている場合は、gcloud components update を実行して最新のバージョンを取得する。
注: gcloud CLI がすでにインストールされている場合には、必ず compute/region と compute/zone のプロパティを設定してください。デフォルトのロケーションを設定することで、gcloud CLI のエラー（One of [--zone, --region] must be supplied: Please specify location など）を防ぐことができます。

GKE での TPU の可用性

GKE を使用して、TPU によるノードプールを作成、管理します。こうした専用アクセラレータを使用することで、大規模な AI モデルのトレーニングとチューニングを行い、推論を実行できます。

GKE でサポートされている TPU バージョンのリストをご覧ください。

TPU 構成を計画する

ML モデルと必要なメモリ量に基づいて TPU 構成を計画します。TPU 構成を計画する場合に関連する手順は次のとおりです。

オンデマンド VM または Spot VM に十分な割り当てを確保する

オンデマンド VM または Spot VM で TPU ノードプールを作成する場合は、使用するリージョンに十分な TPU の割り当てが必要です。

TPU 予約を使用する TPU ノードプールの作成には、TPU の割り当ては必要ありません¹。予約済みの TPU については、この手順はスキップしてもかまいません。

GKE でオンデマンドまたは Spot TPU ノードプールを作成するには、Compute Engine API の割り当てが必要です。Compute Engine API の割り当て（compute.googleapis.com）は、Cloud TPU API の割り当て（tpu.googleapis.com）とは異なります。これは、Cloud TPU API を使用して TPU を作成するときに必要です。

TPU 用の Compute Engine API の割り当ての上限と現在の使用量を確認するには、次の操作を行います。

Google Cloud コンソールで [割り当て] ページに移動します。

[割り当て] に移動

[ フィルタ] ボックスで次の操作を行います。

[サービス] プロパティを選択し、「Compute Engine API」と入力して Enter キーを押します。
[タイプ] プロパティを選択し、[Quota] を選択します。

[名前] プロパティを選択し、TPU のバージョンとマシンタイプに基づいて割り当ての名前を入力します。たとえば、マシンタイプが「ct5lp-」で始まるオンデマンド TPU v5e ノードを作成する場合は、「TPU v5 Lite PodSlice chips」と入力します。

TPU バージョン	次で始まるマシンタイプ	オンデマンドインスタンスの割り当て名	Spot² インスタンスの割り当ての名前
TPU v4	`ct4p-`	`TPU v4 PodSlice chips`	`Preemptible TPU v4 PodSlice chips`
TPU v5e	`ct5l-`	`TPU v5 Lite Device chips`	`Preemptible TPU v5 Lite Device chips`
TPU v5e	`ct5lp-`	`TPU v5 Lite PodSlice chips`	`Preemptible TPU v5 Lite PodSlice chips`
TPU v5p	`ct5p-`	`TPU v5p chips`	`Preemptible TPU v5p chips`

[項目（ロケーションなど）] プロパティを選択し、「region:」に続けて、GKE で TPU を作成するリージョンの名前を入力します。たとえば、ゾーン us-west4-a で TPU ノードを作成する場合は、「region:us-west4」と入力します。TPU の割り当てはリージョン単位であるため、同じリージョン内のすべてのゾーンで同じ TPU の割り当てが使用されます。

入力したフィルタに一致する割り当てがない場合、プロジェクトには目的のリージョンで指定した割り当てのいずれも付与されていないため、TPU 割り当ての増加をリクエストする必要があります。

TPU ノードプールを作成する場合は、--reservation フラグと --reservation-affinity=specific フラグを使用して、予約済みインスタンスを作成します。TPU の予約は、コミットメントを購入した場合に利用できます。↩
Spot インスタンスは、TPU ノードプールの作成時に --spot フラグを使用して作成します。↩

予約の可用性を確保する

予約済みの TPU ノードプールの作成、つまり予約を使用する TPU ノードプールの作成には TPU の割り当ては必要ありません。ただし、ノードプールの作成時には十分な未使用または使用可能なチップが必要です。

プロジェクト内に存在する予約を確認するには、予約のリストを表示します。

TPU 予約で使用可能なチップ数を確認するには、予約の詳細を表示します。

クラスタを作成する

使用可能な TPU があるリージョンに Standard モードで GKE クラスタを作成します。Kubernetes コントロールプレーンの高可用性を実現するリージョンクラスタを使用することをおすすめします。Google Cloud CLI または Google Cloud コンソールを使用できます。

gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --cluster-version VERSION

次のように置き換えます。

CLUSTER_NAME: 新しいクラスタの名前。
LOCATION: TPU の容量が使用可能なリージョン。
VERSION: GKE バージョン。使用するマシンタイプをサポートしている必要があります。デフォルトの GKE バージョンは、ターゲット TPU で利用できない場合があります。TPU マシンタイプで使用できる最小 GKE バージョンについては、GKE での TPU の可用性をご覧ください。

ノードプールを作成する

単一ホスト TPU スライス

Google Cloud CLI、Terraform、または Google Cloud コンソールを使用して、単一ホスト TPU スライスのノードプールを作成できます。

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    [--num-nodes=NUM_NODES \]
    [--spot \]
    [--enable-autoscaling \]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME \]
    [--total-min-nodes TOTAL_MIN_NODES \]
    [--total-max-nodes TOTAL_MAX_NODES \]
    [--location-policy=ANY]

次のように置き換えます。

POOL_NAME: 新しいノードプールの名前。
LOCATION: 使用する TPU バージョンに基づくゾーンの名前。
- TPU v4 の場合は、us-central2-b を使用します。
- ct5l- で始まる TPU v5e マシンタイプの場合は、us-central1-a または europe-west4-b を使用します。
- ct5lp- で始まる TPU v5e マシンタイプの場合は、us-west1-c、us-west4-a、us-west4-b、us-central1-a、us-east1-c、us-east5-b、または europe-west4-a を使用します。
- TPU v5p の場合は、us-east1-d、us-east5-a、または us-east5-c を使用します。
詳細については、GKE での TPU の提供状況をご覧ください。
CLUSTER_NAME: クラスタの名前。
NODE_ZONE: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
MACHINE_TYPE: ノードに使用するマシンのタイプ。TPU 互換マシンタイプの詳細については、TPU 構成のマッピングの表をご覧ください。

必要に応じて、次のフラグも使用できます。

NUM_NODES: 各ゾーンのノードプール内のノードの初期数。このフラグを省略した場合、デフォルトは 3 です。--enable-autoscaling フラグを使用してノードプールに対して自動スケーリングが有効になっている場合は、ワークロードで必要になるとオートスケーラーが追加のノードを直ちにプロビジョニングするため、NUM_NODES を 0 に設定することをおすすめします。
RESERVATION_NAME: ノードプールの作成時に GKE が使用する予約の名前。このフラグを省略すると、GKE は使用可能な TPU を使用します。TPU 予約の詳細については、TPU 予約をご覧ください。
--enable-autoscaling: 自動スケーリングが有効なノードプールを作成します。
- TOTAL_MIN_NODES: ノードプール内のノードの最小数。自動スケーリングを指定しない場合は、このフィールドを省略します。
- TOTAL_MAX_NODES: ノードプール内のノードの最大数。自動スケーリングを指定しない場合は、このフィールドを省略します。
--spot: ノードプール内のノードに Spot VM を使用するようにノードプールを設定します。これは、ノードプールの作成後に変更することはできません。

Terraform

google プロバイダのバージョン 4.84.0 以降を使用していることを確認します。
Terraform 構成に次のブロックを追加します。

resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES
  autoscaling {
    total_min_node_count = TOTAL_MIN_NODES
    total_max_node_count = TOTAL_MAX_NODES
    location_policy      = "ANY"
  }

  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
  }
}

次のように置き換えます。

NODE_POOL_RESOURCE_NAME: Terraform テンプレートのノードプールリソースの名前。
PROJECT_ID: プロジェクト ID。
CLUSTER_NAME: 既存のクラスタの名前。
POOL_NAME: 作成するノードプールの名前。
CLUSTER_LOCATION: クラスタのコンピューティングゾーン。TPU バージョンを利用できるリージョンを指定します。詳細については、TPU のバージョンとトポロジの選択をご覧ください。
NODE_ZONES: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
NUM_NODES: 各ノードプールのゾーン内にあるノードプールの初期ノード数。省略した場合のデフォルトは 3 です。自動スケーリングテンプレートを使用してノードプールの自動スケーリングが有効になっている場合は、ワークロードの要求があれば GKE が直ちに追加の TPU ノードをプロビジョニングするため、NUM_NODES を 0 に設定することをおすすめします。
MACHINE_TYPE: 使用する TPU マシンのタイプ。TPU 互換マシンタイプを確認するには、TPU 構成のマッピングの表をご覧ください。

必要に応じて、次の変数を使用することもできます。

autoscaling: 自動スケーリングが有効なノードプールを作成します。単一ホストの TPU スライスの場合、GKE は TOTAL_MIN_NODES と TOTAL_MAX_NODES の値の間でスケーリングを行います。
- TOTAL_MIN_NODES: ノードプール内のノードの最小数。自動スケーリングを指定しない場合、このフィールドは省略可能です。
- TOTAL_MAX_NODES: ノードプール内のノードの最大数。自動スケーリングを指定しない場合、このフィールドは省略可能です。
RESERVATION_NAME: TPU 予約を使用する場合、これはノードプールの作成時に使用する予約リソースのラベルのリストです。reservation_affinity フィールドに RESERVATION_LABEL_VALUES を挿入する方法については、Terraform プロバイダをご覧ください。
spot: TPU ノードに Spot VM を使用するようにノードプールを設定します。これは、ノードプールの作成後に変更することはできません。詳細については、Spot VM をご覧ください。

コンソール

TPU を使用してノードプールを作成するには:

Google Cloud コンソールで Google Kubernetes Engine のページに移動します。

Google Kubernetes Engine に移動
クラスタのリストで、変更するクラスタの名前をクリックします。
[ノードプールを追加] をクリックします。
[ノードプールの詳細] セクションで、[ノードのロケーションを指定する] チェックボックスをオンにします。
使用する TPU のバージョンに基づいてゾーンを選択します。
- TPU v4 の場合は、us-central2-b を使用します。
- ct5l- で始まる TPU v5e マシンタイプの場合は、us-central1-a または europe-west4-b を使用します。
- ct5lp- で始まる TPU v5e マシンタイプの場合は、us-west1-c、us-west4-a、us-west4-b、us-central1-a、us-east1-c、us-east5-b、または europe-west4-a を使用します。
- TPU v5p の場合は、us-east1-d、us-east5-a、または us-east5-c を使用します。
ナビゲーションパネルで [ノード] をクリックします。
[マシンの構成] セクションで、[TPU] を選択します。
[シリーズ] プルダウンメニューで、次のいずれかを選択します。
- CT4P: TPU v4
- CT5LP: TPU v5e
- CT5P: TPU v5p
[マシンタイプ] プルダウンメニューで、ノードに使用するマシンの名前を選択します。TPU 構成のマッピングの表で、単一ホスト TPU ノードプールを作成するマシンタイプと TPU トポロジを定義する方法を確認します。
[TPU トポロジ] プルダウンメニューで、TPU スライスの物理トポロジを選択します。
[変更が必要です] ダイアログで [変更を適用] をクリックします。
[ブートディスクの種類] が [標準永続ディスク] と [SSD 永続ディスク] のいずれかであることを確認します。
必要に応じて、[Spot VM 上にノードを作成する] チェックボックスをオンにして、ノードプール内のノードで Spot VM を使用します。
[作成] をクリックします。

マルチホスト TPU スライス

Google Cloud CLI、Terraform、または Google Cloud コンソールを使用して、マルチホスト TPU スライスのノードプールを作成できます。

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONE \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    --num-nodes=NUM_NODES \
    [--spot \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME]

次のように置き換えます。

POOL_NAME: 新しいノードプールの名前。
LOCATION: 使用する TPU バージョンに基づくゾーンの名前。
- TPU v4 の場合は、us-central2-b を使用します。
- ct5l- で始まる TPU v5e マシンタイプはマルチホストにはなりません。
- ct5lp- で始まる TPU v5e マシンタイプの場合は、us-west1-c、us-west4-a、us-west4-b、us-central1-a、us-east1-c、us-east5-b、または europe-west4-a を使用します。
- ct5p- で始まる TPU v5p マシンタイプの場合は、us-east1-d、us-east5-a、または us-east5-c を使用します。
詳細については、GKE での TPU の提供状況をご覧ください。
CLUSTER_NAME: クラスタの名前。
NODE_ZONE: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
MACHINE_TYPE: ノードに使用するマシンのタイプ。使用可能なマシンタイプの詳細については、TPU 構成のマッピングをご覧ください。
TPU_TOPOLOGY: TPU スライスの物理トポロジ。トポロジの形式は、次のように TPU のバージョンによって異なります。
- TPU v4 または v5p: トポロジを 3 タプル（{A}x{B}x{C}）で定義します。例: 4x4x4。
- TPU v5e: トポロジを 2 タプル（{A}x{B}）で定義します。例: 2x2。
詳細については、トポロジをご覧ください。
NUM_NODES: ノードプール内のノード数。ゼロにするか、TPU_TOPOLOGY（{A}x{B}x{C}）で定義された値の積を各 VM のチップ数で割った数と同じにする必要があります。マルチホスト TPU v4 と TPU v5e の場合、各 VM のチップ数は 4 です。したがって、TPU_TOPOLOGY が 2x4x4（各 VM に 4 つのチップがある TPU v4）の場合、NUM_NODES は 32÷4 で 8 になります。

必要に応じて、次のフラグも使用できます。

RESERVATION_NAME: ノードプールの作成時に GKE が使用する予約の名前。このフラグを省略すると、GKE は使用可能な TPU ノードプールを使用します。TPU 予約の詳細については、TPU 予約をご覧ください。
--spot: TPU ノードに Spot VM を使用するようにノードプールを設定します。これは、ノードプールの作成後に変更することはできません。詳細については、Spot VM をご覧ください。
--enable-autoscaling: 自動スケーリングが有効なノードプールを作成します。GKE がマルチホスト TPU スライスノードプールをスケーリングすると、ノードプールがゼロから最大サイズまでアトミックにスケールアップされます。
- MAX_NODES: ノードグループの最大サイズ。--enable-autoscaling が指定されている場合、--max-nodes フラグは必須です。TPU_TOPOLOGY（{A}x{B}x{C}）で定義された値の積を、各 VM のチップ数で割った数と同じにする必要があります。

Terraform

google プロバイダのバージョン 4.84.0 以降を使用していることを確認します。
Terraform 構成に次のブロックを追加します。
```
resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES

  autoscaling {
    max_node_count = MAX_NODES
    location_policy      = "ANY"
  }
  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = TPU_TOPOLOGY
  }
}
```
次のように置き換えます。
- NODE_POOL_RESOURCE_NAME: Terraform テンプレートのノードプールリソースの名前。
- PROJECT_ID: プロジェクト ID。
- CLUSTER_NAME: ノードプールを追加する既存のクラスタの名前。
- POOL_NAME: 作成するノードプールの名前。
- CLUSTER_LOCATION: クラスタのコンピューティングのロケーション。Kubernetes コントロールプレーンの信頼性を高めるため、リージョンクラスタを使用することをおすすめします。ゾーンクラスタを使用することもできます。詳細については、TPU のバージョンとトポロジの選択をご覧ください。
- NODE_ZONES: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
- NUM_NODES: ノードプール内のノード数。この値はゼロにするか、TPU チップの数を 4 で割った数にする必要があります。これは、マルチホスト TPU スライスでは各 TPU ノードに 4 つのチップがあるためです。たとえば、TPU_TOPOLOGY が 4x8 の場合、32 チップがあるため、NUM_NODES は 8 にする必要があります。TPU トポロジの詳細については、TPU 構成のマッピングの表をご覧ください。
- TPU_TOPOLOGY: TPU スライスに必要な物理トポロジを示します。トポロジの形式は、使用している TPU のバージョンによって異なります。
  - TPU v4 の場合: トポロジを 3 タプル（{A}x{B}x{C}）で定義します。例: 4x4x4。
  - TPU v5e の場合: トポロジを 2 タプル（{A}x{B}）で定義します。例: 2x2。
必要に応じて、次の変数を使用することもできます。
- RESERVATION_NAME: TPU 予約を使用する場合、これはノードプールの作成時に使用する予約リソースのラベルのリストです。reservation_affinity フィールドに RESERVATION_LABEL_VALUES を挿入する方法については、Terraform プロバイダをご覧ください。
- autoscaling: 自動スケーリングが有効なノードプールを作成します。GKE がマルチホスト TPU スライスノードプールをスケーリングすると、ノードプールがゼロから最大サイズまでアトミックにスケールアップされます。
  - MAX_NODES: ノードプールの最大サイズです。TPU_TOPOLOGY（{A}x{B}x{C}）で定義された値の積を各 VM のチップ数で割った数と同じにする必要があります。
- spot: TPU ノードに Spot VM を使用するようにノードプールを設定します。これは、ノードプールの作成後に変更することはできません。詳細については、Spot VM をご覧ください。

コンソール

TPU を使用してノードプールを作成するには:

Google Cloud コンソールで Google Kubernetes Engine のページに移動します。

Google Kubernetes Engine に移動
クラスタのリストで、変更するクラスタの名前をクリックします。
[ノードプールを追加] をクリックします。
[ノードプールの詳細] セクションで、[ノードのロケーションを指定する] チェックボックスをオンにします。
使用する TPU のバージョンに基づいてゾーンを選択します。
- TPU v4 の場合は、us-central2-b を使用します。
- ct5l- で始まる TPU v5e マシンタイプはマルチホストにはなりません。
- ct5lp- で始まる TPU v5e マシンタイプの場合は、us-west1-c、us-west4-a、us-west4-b、us-central1-a、us-east1-c、us-east5-b、または europe-west4-a を使用します。
- ct5p- で始まる TPU v5p マシンタイプの場合は、us-east1-d、us-east5-a、または us-east5-c を使用します。
ナビゲーションパネルで [ノード] をクリックします。
[マシンの構成] セクションで、[TPU] を選択します。
[シリーズ] プルダウンメニューで、次のいずれかを選択します。
- CT4P: TPU v4 の場合。
- CT5LP: TPU v5e の場合。
[マシンタイプ] プルダウンメニューで、ノードに使用するマシンの名前を選択します。TPU 構成のマッピングの表で、マルチホスト TPU ノードプールを作成するマシンタイプと TPU トポロジを定義する方法を確認します。
[TPU トポロジ] プルダウンメニューで、TPU スライスの物理トポロジを選択します。
[変更が必要です] ダイアログで [変更を適用] をクリックします。
[ブートディスクの種類] が [標準永続ディスク] と [SSD 永続ディスク] のいずれかであることを確認します。
必要に応じて、[Spot VM 上にノードを作成する] チェックボックスをオンにして、ノードプール内のノードで Spot VM を使用します。
[作成] をクリックします。

プロビジョニング状態

使用可能な TPU 容量が不足しているために GKE が TPU スライスノードプールを作成できない場合、GKE は、容量不足が原因で TPU ノードを作成できないことを示すエラーメッセージを返します。

単一ホストの TPU スライスノードプールを作成している場合は、次のようなエラーメッセージが表示されます。

2 nodes cannot be created due to lack of capacity. The missing nodes will be
created asynchronously once capacity is available. You can either wait for the
nodes to be up, or delete the node pool and try re-creating it again later.

マルチホスト TPU スライスノードプールを作成している場合は、次のようなエラーメッセージが表示されます。

The nodes (managed by ...) cannot be created now due to lack of capacity. They
will be created asynchronously once capacity is available. You can either wait
for the nodes to be up, or delete the node pool and try re-creating it again
later.

TPU プロビジョニングリクエストは長時間キューに留まり、キューにある間は「プロビジョニング」状態のままになります。

容量が利用可能になると、GKE は作成されていない残りのノードを作成します。

すぐに容量が必要な場合は、Spot VM の使用を検討してください。ただし、Spot VM は、オンデマンドインスタンスとは異なる割り当てを消費します。

キューに格納された TPU リクエストを削除するには、TPU スライスノードプールを削除します。

TPU ノードでワークロードを実行する

ワークロードの準備

TPU ワークロードには、次の準備要件があります。

JAX、PyTorch、TensorFlow などのフレームワークは、libtpu 共有ライブラリを使用して TPU VM にアクセスします。libtpu には、XLA コンパイラ、TPU ランタイムソフトウェア、TPU ドライバが含まれています。PyTorch と JAX の各リリースには、特定の libtpu.so バージョンが必要です。GKE で TPU を使用するには、次のバージョンを使用してください。

TPU タイプ	`libtpu.so` のバージョン
TPU v5e `tpu-v5-lite-podslice` `tpu-v5-lite-device`	推奨される jax[tpu] バージョン: v0.4.9 以降推奨される torchxla[tpuvm] のバージョン: v2.1.0 以降
TPU v5p	推奨される jax[tpu] バージョン: 0.4.19 以降。推奨される torchxla[tpuvm] バージョン: 毎晩更新されるバージョンの 2023 年 10 月 23 日付のビルドを使用することをおすすめします。
TPU v4 `tpu-v4-podslice`	jax[tpu]: v0.4.4 以降 torchxla[tpuvm]: v2.0.0 以降

TPU リソースをリクエストするコンテナに、次の環境変数を設定します。
- TPU_WORKER_ID: 各 Pod の一意の整数。この ID は、TPU スライス内の一意のワーカー ID を示します。このフィールドでサポートされる値の範囲は、0 から Pod 数から 1 を引いた値までです。
- TPU_WORKER_HOSTNAMES: スライス内で相互に通信する必要がある TPU VM ホスト名または IP アドレスのカンマ区切りのリスト。スライス内の TPU VM ごとにホスト名または IP アドレスが必要です。IP アドレスまたはホスト名のリストは、TPU_WORKER_ID によって順序付けされ、ゼロのインデックスが付けられます。
- ワークロードマニフェストに Kubernetes ノードセレクタを追加して、定義した TPU マシンタイプと TPU トポロジで、GKE が TPU ワークロードをスケジュールできるようにします。
```
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
  
```
  次のように置き換えます。
  - TPU_ACCELERATOR: TPU アクセラレータの名前:
    - TPU v4 の場合は、tpu-v4-podslice を使用します。
    - ct5l- で始まる TPU v5e マシンタイプの場合は、tpu-v5-lite-device を使用します。
    - ct5lp- で始まる TPU v5e マシンタイプの場合は、tpu-v5-lite-podslice を使用します。
    - TPU v5p の場合は、tpu-v5p-slice を使用します。
  - TPU_TOPOLOGY: TPU スライスの物理トポロジ。トポロジの形式は、次のように TPU のバージョンによって異なります。
    - TPU v4: トポロジを 3 タプル（{A}x{B}x{C}）で定義します。例: 4x4x4。
    - TPU v5e: トポロジを 2 タプル（{A}x{B}）で定義します。例: 2x2。
    - TPU v5p: トポロジを 3 タプル（{A}x{B}x{C}）で定義します。例: 4x4x4。

ワークロードの準備が完了したら、TPU を使用する Job を実行できます。

以降のセクションでは、TPU で単純な計算を実行する Job の実行例を示します。

例 1: TPU ノードプールで使用可能な TPU チップの数を表示するワークロードを実行する

次のワークロードは、マルチホスト TPU スライス内のノード全体の TPU チップ数を返します。マルチホストスライスを作成する場合、ワークロードに次のパラメータを設定します。

TPU バージョン: TPU v4
トポロジ: 2x2x4

このバージョンとトポロジの選択により、マルチホストスライスが作成されます。

次のマニフェストを available-chips-multihost.yaml として保存します。

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        securityContext:
          privileged: true
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4

マニフェストをデプロイします。
```
kubectl create -f available-chips-multihost.yaml
```
GKE は、4 つの TPU VM（マルチホスト TPU スライス）を使用して TPU v4 スライスを実行します。スライスには、相互接続された 16 個のチップがあります。

Job によって 4 つの Pod が作成されたことを確認します。

kubectl get pods

出力は次のようになります。

NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s

いずれかの Pod のログを取得します。
```
kubectl logs POD_NAME
```
POD_NAME は、作成した Pod の名前に置き換えます。例: tpu-job-podslice-0-5cd8r

出力は次のようになります。
```
TPU cores: 16
```

例 2: TPU VM で使用可能な TPU チップの数を表示するワークロードを実行する

次のワークロードは、特定のノードに関連付けられている TPU チップの数を表示する静的 Pod です。単一ホストノードを作成するには、ワークロードに次のパラメータを設定します。

TPU バージョン: TPU v5e
トポロジ: 2x4

このバージョンとトポロジを選択すると、単一ホストのスライスが作成されます。

次のマニフェストを available-chips-singlehost.yaml として保存します。

apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    securityContext:
      privileged: true
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8

マニフェストをデプロイします。
```
kubectl create -f available-chips-singlehost.yaml
```
GKE は、TPU v5e を使用する 8 つの単一ホスト TPU スライスを含むノードをプロビジョニングします。各 TPU VM には 8 つのチップ（単一ホストの TPU スライス）があります。
Pod のログを取得します。
```
kubectl logs tpu-job-jax-v5
```
出力は次のようになります。
```
Total TPU chips: 8
```

アクセラレータ（GPU と TPU）を使用してノードプールをアップグレードする

GKE は、ノードプールを含む Standard クラスタを自動的にアップグレードします。ノードを新しいバージョンに早急に移行する必要がある場合は、ノードプールを手動でアップグレードすることもできます。クラスタのアップグレードの動作を制御するには、リリースチャンネル、メンテナンスの時間枠と除外、ロールアウトのシーケンスを使用します。

サージアップグレードや Blue/Green アップグレードなど、ノードプールにノードアップグレード戦略を構成することもできます。これらの戦略を構成することで、環境の速度と中断の最適なバランスが実現されるようにノードプールをアップグレードできます。マルチホスト TPU スライスノードプールの場合、GKE は、構成済みのノードアップグレード戦略を使用せずに、1 つのステップでノードプール全体をアトミックに再作成します。詳細については、GKE の TPU に関連する用語の「アトミック性」の定義をご覧ください。

ノードのアップグレード戦略を使用する場合は、構成に応じて、GKE が一時的に追加のリソースをプロビジョニングする必要があります。Google Cloud でノードプールのリソースの容量が限られている場合（例: GPU または TPU を使用してノードを追加しようとするとリソースの可用性エラーが表示される）は、リソースが制限された環境でのアップグレードをご覧ください。

クリーンアップ

このガイドで使用したリソースに関して、Google Cloud アカウントに課金されないようにするには、ワークロードのスケジュールがなくなった TPU ノードプールの削除を検討してください。実行中のワークロードを正常に終了する必要がある場合は、ノードを削除する前に kubectl drain を使用してワークロードをクリーンアップします。

TPU ノードプールを削除します。
```
gcloud container node-pools delete POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME
```
次のように置き換えます。
- POOL_NAME: ノードプールの名前。
- CLUSTER_NAME: クラスタの名前。
- LOCATION: クラスタの Compute Engine のロケーション。

その他の構成

以降のセクションでは、TPU ワークロードに適用できる追加の構成について説明します。

マルチスライス

小さなスライスをマルチスライスに集約して、より大きなトレーニングワークロードを処理できます。詳細については、GKE のマルチスライス TPU をご覧ください。

TPU 予約を移行する

既存の TPU 予約がある場合は、まず TPU 予約を新しい Compute Engine ベースの予約システムに移行する必要があります。移行が不要な Compute Engine ベースの予約システムを作成することもできます。TPU 予約を移行する方法については、TPU 予約をご覧ください。

ロギング

TPU VM を含む GKE ノード上で実行されているコンテナによって出力されたログは、GKE ロギングエージェントが収集し、Logging に送信して、Logging に表示します。

GKE ノードの自動プロビジョニングを使用する

TPU ワークロードのリソース需要を満たすように、ノードプールを自動的に作成および削除するように GKE を構成できます。詳細については、Cloud TPU の構成をご覧ください。

TPU ノードの自動修復

マルチホスト TPU スライスノードプール内の TPU ノードが異常な場合は、ノードプール全体が再作成されます。TPU ノードが異常な状態になる条件は次のとおりです。

共通のノード条件を持つ TPU ノード。
割り当てできない TPU 数が 0 より大きい TPU ノード。
プリエンプションにより停止または終了した TPU VM インスタンス。
ノードのメンテナンス: マルチホスト TPU スライスノードプール内の TPU ノード（VM）がホストメンテナンスのために停止した場合、GKE は TPU スライス全体を再作成します。

修復ステータス（失敗の理由を含む）はオペレーション履歴で確認できます。割り当て不足が原因でエラーが発生する場合は、Google Cloud アカウント担当者に連絡して、対応する割り当ての増加をリクエストしてください。

TPU ノードの正常終了を構成する

コントロールプレーンで 1.29.1-gke.1425000 以降が実行されている GKE クラスタでは、TPU ノードはシャットダウンの差し迫ったノードに警告する SIGTERM シグナルをサポートしています。TPU ノードでは、最大 5 分前のシャットダウン通知を構成できます。

この通知期間内に ML ワークロードを正常に終了するように GKE を構成できます。正常終了中に、ワークロードはデータ損失を抑えるためにワークロードデータを保存するなどのクリーンアッププロセスを実行できます。最大通知時間を取得するには、Pod マニフェストで spec.terminationGracePeriodSeconds フィールドを次のように 300 秒（5 分）に設定します。

    spec:
      terminationGracePeriodSeconds: 300

GKE は、これらの Pod を正常に終了し、トレーニング状態の保存など、定義した終了アクションを実行するために最善を尽くします。

特権モードを使用せずコンテナを実行する

TPU ノードがバージョン 1.28 以前で実行されている場合は、次のセクションをご覧ください。

TPU VM 上で実行されるコンテナは、ドライバが直接メモリアクセス（DMA）を介して TPU チップと通信できるように、ロックされたメモリの上限を引き上げる必要があります。これを有効にするには、上位の ulimit を構成する必要があります。コンテナの権限スコープを縮小するには、次の操作を行います。

securityContext を編集して次のフィールドを含めます。
```
securityContext:
  capabilities:
    add: ["SYS_RESOURCE"]
```
TPU リソースを使用するようにワークロードを設定する前に、コンテナ内で次のコマンドを実行して ulimit を増やします。
```
ulimit -l 68719476736
```

注: TPU v5e の場合、バージョン 1.27.4-gke.900 以降のクラスタで特権モードなしでコンテナを実行できます。

オブザーバビリティと指標

ダッシュボード

Google Cloud コンソールの [Kubernetes クラスタ] ページの [オブザーバビリティ] タブに、TPU オブザーバビリティ指標が表示されます。詳細については、GKE のオブザーバビリティ指標をご覧ください。

TPU ダッシュボードは、GKE クラスタでシステム指標が有効になっている場合にのみ表示されます。

ランタイム指標

GKE バージョン 1.27.4-gke.900 以降、JAX バージョン 0.4.14 以降を使用し containerPort: 8431 を指定する TPU ワークロードでは、TPU 使用率の指標を GKE システム指標としてエクスポートします。Cloud Monitoring では、TPU ワークロードのランタイムパフォーマンスをモニタリングするために、次の指標を使用できます。

デューティサイクル: 過去のサンプリング期間（60 秒）において、TensorCore が TPU チップでアクティブに処理していた時間の割合。割合が大きいほど、TPU 使用率が高くなります。
メモリ使用量: 割り当てられたアクセラレータメモリの量（バイト単位）。60 秒ごとにサンプリングされます。
メモリの総容量: アクセラレータの総メモリ（バイト単位）。60 秒ごとにサンプリングされます。

これらの指標は、Kubernetes ノード（k8s_node）と Kubernetes コンテナ（k8s_container）のスキーマにあります。

Kubernetes コンテナ:

kubernetes.io/container/accelerator/duty_cycle
kubernetes.io/container/accelerator/memory_used
kubernetes.io/container/accelerator/memory_total

Kubernetes ノード:

kubernetes.io/node/accelerator/duty_cycle
kubernetes.io/node/accelerator/memory_used
kubernetes.io/node/accelerator/memory_total

ホスト指標

GKE バージョン 1.28.1-gke.1066000 以降では、TPU VM は TPU 使用率の指標を GKE システム指標としてエクスポートします。Cloud Monitoring では、次の指標を使用して TPU ホストのパフォーマンスをモニタリングできます。

TensorCore の使用率: 使用されている TensorCore の現在の割合。TensorCore の値は、マトリックス乗算ユニット（MXU）およびベクトル単位の合計と等しくなります。TensorCore の使用率の値は、過去のサンプル期間（60 秒）に実行された TensorCore オペレーションを、同じ期間にサポートされている TensorCore オペレーションの数で割ったものです。値が大きいほど、使用率が高いことを意味します。
メモリ帯域幅の使用率: 現在使用されているアクセラレータメモリ帯域幅の割合。サンプル期間（60 秒）で使用されたメモリ帯域幅を、同じサンプル期間でサポートされる最大帯域幅で割って計算されます。

これらの指標は、Kubernetes ノード（k8s_node）と Kubernetes コンテナ（k8s_container）のスキーマにあります。

Kubernetes コンテナ:

kubernetes.io/container/accelerator/tensorcore_utilization
kubernetes.io/container/accelerator/memory_bandwidth_utilization

Kubernetes ノード:

kubernetes.io/container/node/tensorcore_utilization
kubernetes.io/container/node/memory_bandwidth_utilization

詳細については、Kubernetes の指標と GKE のシステム指標をご覧ください。

既知の問題

新しい TPU ノードが使用可能な TPU を報告する前に、クラスタオートスケーラーが新しい TPU ノードの容量を誤って計算することがあります。これにより、クラスタオートスケーラーが追加のスケールアップを実行して、結果として必要以上にノードを作成する場合があります。クラスタオートスケーラーは、通常のスケールダウンオペレーションの後、不要なノードをスケールダウンします。
クラスタオートスケーラーは、15 分以上待機状態になっている TPU ノードプールのスケールアップをキャンセルします。クラスタオートスケーラーは、このようなスケールアップオペレーションを後で再試行します。この動作により、予約を使用しないお客様の TPU の入手可能性が低下する可能性があります。
TPU taint の toleration がある TPU 以外のワークロードは、TPU ノードプールのドレイン中に再作成されると、ノードプールのスケールダウンを妨げることがあります。

GKE Standard に TPU ワークロードをデプロイする

始める前に

GKE での TPU の可用性

TPU 構成を計画する

オンデマンド VM または Spot VM に十分な割り当てを確保する

予約の可用性を確保する

クラスタを作成する

ノードプールを作成する

単一ホスト TPU スライス

gcloud

Terraform

コンソール

マルチホスト TPU スライス

gcloud

Terraform

コンソール

プロビジョニング状態

TPU ノードでワークロードを実行する

ワークロードの準備

例 1: TPU ノードプールで使用可能な TPU チップの数を表示するワークロードを実行する

例 2: TPU VM で使用可能な TPU チップの数を表示するワークロードを実行する

アクセラレータ（GPU と TPU）を使用してノードプールをアップグレードする

クリーンアップ

その他の構成

マルチスライス

TPU 予約を移行する

ロギング

GKE ノードの自動プロビジョニングを使用する

TPU ノードの自動修復

TPU ノードの正常終了を構成する

特権モードを使用せずコンテナを実行する

オブザーバビリティと指標

ダッシュボード

ランタイム指標

ホスト指標

既知の問題

次のステップ