TPU と Flex Start VM で小規模なバッチワークロードを実行する

Autopilot Standard

このガイドでは、Flex Start VM を使用して、中規模および小規模のトレーニングワークロードの TPU プロビジョニングを最適化する方法について説明します。Flex Start VM は、Flex Start 使用オプションを使用して作成されます。このガイドでは、Flex Start VM を使用して、TPU スライスノードプールで構成されるワークロードをデプロイします。

このガイドは、バッチワークロードの実行に Kubernetes コンテナオーケストレーション機能を使用することを検討している ML エンジニア、プラットフォーム管理者、オペレーター、データおよび AI のスペシャリストを対象としています。 Google Cloud のコンテンツで使用されている一般的なロールとタスクの例の詳細については、一般的な GKE ユーザーのロールとタスクをご覧ください。

Flex Start の料金

Flex Start は、リソースを必要に応じて動的にプロビジョニングする必要があり、最長 7 日間の短期間での利用が想定され、複雑な割り当て管理が必要なく、コスト効率よく使いたいワークロードにおすすめです。Flex Start は Dynamic Workload Scheduler によって動いており、料金も Dynamic Workload Scheduler に基づいたものになります。

vCPU、GPU、TPU が最大 53% 割引になります。
従量課金制です。

始める前に

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

Google Kubernetes Engine API を有効にする。

Google Kubernetes Engine API の有効化

このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。gcloud CLI をインストール済みの場合は、gcloud components update コマンドを実行して最新のバージョンを取得します。以前のバージョンの gcloud CLI では、このドキュメントのコマンドを実行できない場合があります。
注: gcloud CLI がすでにインストールされている場合には、必ず compute/region プロパティを設定してください。主にゾーンクラスタを使用する場合は、代わりに compute/zone を設定します。デフォルトのロケーションを設定することで、gcloud CLI のエラー（One of [--zone, --region] must be supplied: Please specify location など）を防止できます。クラスタのロケーションが設定したデフォルトと異なる場合は、特定のコマンドでロケーションの指定が必要になることがあります。

バージョン 1.33.0-gke.1712000 以降を実行している Autopilot クラスタまたは Standard クラスタがあることを確認します。
Flex Start の制限事項を理解していることを確認します。
Standard クラスタを使用する場合は、クラスタが正常に機能するように、Flex Start を有効にしていないノードプールを少なくとも 1 つ維持していることを確認します。
ノードのロケーションにプリエンプティブル TPU の割り当てがあることを確認します。

Flex Start を使用してノードプールを作成する

Autopilot モードでクラスタを使用している場合は、このセクションをスキップして、バッチワークロードを実行するに進みます。

既存の Standard クラスタで Flex Start を有効にしてノードプールを作成するには、gcloud CLI を使用します。

単一ホスト TPU スライスノードプールを作成する

Flex Start を使用して、単一ホストの TPU スライスノードプールを作成できます。

Flex Start を使用してノードプールを作成します。
```
gcloud container node-pools create NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --location=CONTROL_PLANE_LOCATION \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    --reservation-affinity=none \
    --enable-autoscaling \
    --flex-start \
    --num-nodes 0 \
    --min-nodes=0 \
    --max-nodes=1
```
次のように置き換えます。
- NODE_POOL_NAME: ノードプールに付ける名前。
- CLUSTER_NAME: クラスタの名前。
- CONTROL_PLANE_LOCATION: クラスタコントロールプレーンのコンピューティングリージョン。
- NODE_ZONES: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
- MACHINE_TYPE: ノードに使用するマシンのタイプ。TPU 互換マシンタイプの詳細については、TPU のバージョンを選択するの表をご覧ください。
  
  たとえば、ノードプールの作成コマンドに次のパラメータを含めることができます。
```
...
--machine-type=ct6e-standard-4t \
--tpu-topology=4x4 \
--enable-autoscaling \
--num-nodes=0 \
--max-nodes=4 \
```
  このコマンドは、4x4 トポロジが 16 個のチップで構成され、各 ct6e-standard-4t VM に 4 個のチップがあるため、--max-nodes フィールドを 4 に設定します。
  
  クラスタオートスケーラーは、ワークロードに必要なノード数までスケールアップします。ワークロードが完了すると、クラスタオートスケーラーはゼロノードにスケールダウンします。
- --reservation-affinity=none: Flex Start では予約を使用することも、必要とすることもありません。

マルチホスト TPU スライスノードプールを作成する

マルチホスト TPU スライスノードプールを作成する手順は、Ironwood（TPU7x）を使用するか、以前の TPU バージョンを使用するかによって異なります。

Ironwood（TPU7x）

Google Cloud CLI または Terraform を使用して、Ironwood（TPU7x）バージョンのマルチホスト TPU スライスノードプールを作成できます。

gcloud

Ironwood（TPU7x）でマルチホスト TPU スライスノードプールを作成するには、まずワークロードポリシーを作成する必要があります。

ワークロードポリシーを作成します。
```
gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
    --type=HIGH_THROUGHPUT \
    --accelerator-topology=TPU_TOPOLOGY \
    --project=PROJECT_ID \
    --region=REGION
```
次のように置き換えます。
- WORKLOAD_POLICY_NAME: ワークロードポリシーの名前。
- TPU_TOPOLOGY: TPU Ironwood（TPU7x）トポロジ。例: 2x2x2。サポートされているすべての Ironwood（TPU7x）トポロジについては、トポロジのセクションをご覧ください。
- PROJECT_ID: 実際の Google Cloud プロジェクト ID。
- REGION: ワークロードポリシーのリージョン。ワークロードポリシーはリージョンリソースであり、同じトポロジを共有するノードプール間で再利用できます。
ワークロードポリシーを使用してノードプールを作成します。
```
gcloud container node-pools create NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --location=us-central1 \
    --node-locations=us-central1-c \
    --machine-type=tpu7x-standard-4t \
    --reservation-affinity=none \
    --enable-autoscaling \
    --num-nodes=0 --min-nodes=0 --max-nodes=MAX_NODES \
    --flex-start \
    --placement-policy=WORKLOAD_POLICY
```
次のように置き換えます。
- NODE_POOL_NAME: 新しいノードプールの名前。
- WORKLOAD_POLICY: 作成したワークロードポリシーの名前。
- MAX_NODES: ノードプールの最大サイズ。--enable-autoscaling が指定されている場合、--max-nodes フラグは必須です。TPU_TOPOLOGY（{A}x{B}x{C}）で定義された値の積を、各 VM のチップ数で割った数と同じにする必要があります。たとえば、TPU_TOPOLOGY が 2x2x2 の場合、積は 8 です。tpu7x-standard-4t の各 VM には 4 つのチップがあるため、ノード数は 2 です。
このコマンドは、次の特性を持つ NODE_POOL_NAME という名前のノードプールを作成します。
- --machine-type=tpu7x-standard-4t は、Ironwood（TPU7x）マシンタイプを指定します。
- --flex-start は、Flex Start を有効にします。

Terraform

google プロバイダのバージョン 4.84.0 以降を使用していることを確認します。
ワークロードポリシーを作成します。
```
resource "google_compute_resource_policy" {
  name   = "WORKLOAD_POLICY_NAME"
  region = CLUSTER_LOCATION
  workload_policy {
    type = "HIGH_THROUGHPUT"
    accelerator_topology = "TPU_TOPOLOGY"
  }
}
```
次のように置き換えます。
- WORKLOAD_POLICY_NAME: ワークロードポリシーの名前。
- CLUSTER_LOCATION: クラスタのコンピューティングのロケーション。Kubernetes コントロールプレーンの信頼性を高めるため、リージョンクラスタを使用することをおすすめします。ゾーンクラスタを使用することもできます。詳細については、TPU のバージョンとトポロジの選択をご覧ください。
- TPU_TOPOLOGY: TPU Ironwood（TPU7x）トポロジ。例: 2x2x2。サポートされているすべての Ironwood（TPU7x）トポロジについては、TPU を計画するをご覧ください。
google_compute_resource_policy リファレンスの詳細については、Terraform プロバイダをご覧ください。
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]
    }
    flex_start = true
  }

  placement_policy {
    policy_name = WORKLOAD_POLICY_NAME
  }
}
```
次のように置き換えます。
- NODE_POOL_RESOURCE_NAME: Terraform テンプレートのノードプールリソースの名前。
- PROJECT_ID: プロジェクト ID。
- CLUSTER_NAME: ノードプールを追加する既存のクラスタの名前。
- POOL_NAME: 作成するノードプールの名前。
- 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 トポロジの詳細については、トポロジを選択するの表をご覧ください。
必要に応じて、次の変数を使用することもできます。
- RESERVATION_NAME: TPU 予約を使用する場合は、ノードプールの作成時に使用する予約リソースのラベルのリストを指定します。reservation_affinity フィールドに RESERVATION_LABEL_VALUES を挿入する方法については、Terraform プロバイダをご覧ください。
- autoscaling: 自動スケーリングが有効なノードプールを作成します。GKE がマルチホスト TPU スライスノードプールをスケーリングすると、ノードプールがゼロから最大サイズまでアトミックにスケールアップされます。
  - MAX_NODES: ノードプールの最大サイズ。TPU_TOPOLOGY（{A}x{B}x{C}）で定義された値の積を各 VM のチップ数で割った数と同じにする必要があります。たとえば、TPU_TOPOLOGY が 2x2x2 の場合、積は 8 です。tpu7x-standard-4t の各 VM には 4 つのチップがあるため、ノード数は 2 です。
- spot: TPU スライスノードに Spot VM を使用するノードプール。この設定は、ノードプールの作成後に変更することはできません。詳細については、Spot VM をご覧ください。
- flex_start: Flex Start 消費オプションを使用するノードプール。spot が有効になっている場合、この設定を true に設定することはできません。

その他の TPU バージョン

Google Cloud CLI、Terraform、または Google Cloud コンソールを使用して、バージョン v3、v4、v5p、v5e、Trillium（v6e）でマルチホスト TPU スライスノードプールを作成できます。

gcloud

  gcloud container node-pools create NODE_POOL_NAME \
      --cluster=CLUSTER_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --node-locations=NODE_ZONES \
      --machine-type=MACHINE_TYPE \
      --tpu-topology=TPU_TOPOLOGY \
      --reservation-affinity=none \
      --enable-autoscaling \
      --num-nodes=0 --min-nodes=0 --max-nodes=MAX_NODES \
      --flex-start

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

NODE_POOL_NAME: 新しいノードプールの名前。
CLUSTER_NAME: クラスタの名前。
CONTROL_PLANE_LOCATION: 使用する TPU バージョンに基づくゾーンの名前。使用可能なロケーションについては、GKE での TPU の可用性をご覧ください。
NODE_ZONES: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
MACHINE_TYPE: ノードに使用するマシンのタイプ。TPU 互換マシンタイプの詳細については、TPU のバージョンを選択するの表をご覧ください。
TPU_TOPOLOGY: TPU トポロジ。例: 2x2x2。サポートされているすべての TPU トポロジについては、トポロジのセクションをご覧ください。
MAX_NODES: ノードプールの最大サイズ。--enable-autoscaling が指定されている場合、--max-nodes フラグは必須です。TPU_TOPOLOGY（{A}x{B}x{C}）で定義された値の積を、各 VM のチップ数で割った数と同じにする必要があります。

このコマンドは、Flex Start が有効な NODE_POOL_NAME という名前のノードプールを作成します。

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]
    }
    flex_start = 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 トポロジの詳細については、トポロジを選択するの表をご覧ください。
必要に応じて、次の変数を使用することもできます。
- 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 をご覧ください。
- flex_start: Flex Start 使用オプションを使用するようにノードプールを設定します。spot が有効になっている場合、true に設定することはできません。

コンソール

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

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

Google Kubernetes Engine に移動
クラスタのリストで、変更するクラスタの名前をクリックします。
[ノードプールを追加] をクリックします。
[ノードプールの詳細] セクションで、[ノードのロケーションを指定する] チェックボックスをオンにします。
使用する TPU バージョンに基づくゾーンの名前を選択します。使用可能なロケーションについては、GKE での TPU の可用性をご覧ください。
ナビゲーションパネルで [ノード] をクリックします。
[マシンの構成] セクションで、[TPU] を選択します。
[シリーズ] プルダウンメニューで、次のいずれかを選択します。
- CT3: TPU v3、単一ホストデバイス
- CT3P: TPU v3、マルチホスト Pod スライス
- CT4P: TPU v4
- CT5LP: TPU v5e
- CT5P: TPU v5p
- CT6E: TPU Trillium（v6e）
[マシンタイプ] プルダウンメニューで、ノードに使用するマシンの名前を選択します。TPU バージョンを選択するの表で、マルチホスト TPU スライスノードプールを作成するマシンタイプと TPU トポロジを定義する方法を確認します。
[TPU トポロジ] プルダウンメニューで、TPU スライスの物理トポロジを選択します。
[変更が必要です] ダイアログで [変更を適用] をクリックします。
[ブートディスクの種類] が [標準永続ディスク] と [SSD 永続ディスク] のいずれかであることを確認します。
必要に応じて、[Spot VM 上にノードを作成する] チェックボックスをオンにして、ノードプール内のノードで Spot VM を使用します。
[作成] をクリックします。

ノードプールで Flex Start のステータスを確認する

次のコマンドを実行します。

gcloud container node-pools describe NODE_POOL_NAME \
    --cluster CLUSTER_NAME \
    --location CONTROL_PLANE_LOCATION \
    --format="get(config.flexStart)"

ノードプールで Flex Start が有効になっている場合、フィールド flexStart は True に設定されます。

バッチワークロードを実行する

このセクションでは、Flex Start VM を使用して TPU ノードのスケジュールを設定する Job を作成します。Kubernetes の Job コントローラは、1 つ以上の Pod を作成し、特定のタスクが正常に実行されるようにします。

Google Cloud コンソールで、 [Cloud Shell をアクティブにする] をクリックして、Cloud Shell セッションを起動します。 Google Cloud コンソールの下部ペインでセッションが開きます。

dws-flex-start.yaml という名前のファイルを作成します。

単一ホスト

dws-flex-start.yaml ファイルには次のマニフェストを使用します。

apiVersion: batch/v1
kind: Job
metadata:
  name: job-1
spec:
  template:
    spec:
      nodeSelector:
        cloud.google.com/gke-flex-start: "true"
        cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
        cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
      containers:
      - name: container-1
        image: gcr.io/k8s-staging-perf-tests/sleep:latest
        args: ["3600s"] # Sleep for 1 hour
        resources:
          requests:
              google.com/tpu: NUM_CHIPS
          limits:
              google.com/tpu: NUM_CHIPS
      restartPolicy: OnFailure

マルチホスト

dws-flex-start.yaml ファイルには次のマニフェストを使用します。

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: job-1
---
apiVersion: batch/v1
kind: Job
metadata:
  name: job-1
spec:
  backoffLimit: 0
  completions: 2
  parallelism: 2
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
          cloud.google.com/gke-flex-start: "true"
          cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
          cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        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
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            google.com/tpu: NUM_CHIPS
          limits:
            google.com/tpu: NUM_CHIPS

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

ACCELERATOR_TYPE: ノードプールの作成時に使用した TPU アクセラレータのタイプ。たとえば、tpu-v4-podslice や tpu-v5-lite-podslice です。
TPU_TOPOLOGY: TPU スライスの物理トポロジ。たとえば、TPU のバージョンに応じて 4x4x4 または 2x2 になります。
NUM_CHIPS: 各 VM の TPU チップ数は 1、4、8 個です。詳細については、TPU のバージョンをご覧ください。

dws-flex-start.yaml マニフェストを適用します。
```
kubectl apply -f dws-flex-start.yaml
```

Job が同じノードで実行されていることを確認します。

kubectl get pods

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

NAME    READY   STATUS      RESTARTS   AGE   IP       NODE               NOMINATED NODE   READINESS GATES
job-1   0/1     Completed   0          19m   10.(...) gke-flex-zonal-a2  <none>           <none>

クリーンアップ

このページで使用したリソースについて、 Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

プロジェクトを削除する

注意: プロジェクトを削除すると、次のような影響があります。

プロジェクト内のすべてのものが削除されます。このドキュメントのタスクで既存のプロジェクトを使用した場合、それを削除すると、そのプロジェクトで行った他の作業もすべて削除されます。
カスタムプロジェクト ID が失われます。このプロジェクトを作成したときに、将来使用するカスタムプロジェクト ID を作成した可能性があります。そのプロジェクト ID を使用した URL（たとえば、appspot.com）を保持するには、プロジェクト全体ではなくプロジェクト内の選択したリソースだけを削除します。

In the Google Cloud console, go to the Manage resources page.
Go to Manage resources
In the project list, select the project that you want to delete, and then click Delete.
In the dialog, type the project ID, and then click Shut down to delete the project.

個々のリソースを削除する

ジョブを削除します。

kubectl delete job -l "job-name in (job-1,job-2)"

ノードプールを削除します。

gcloud container node-pools delete NODE_POOL_NAME \
      --location CONTROL_PLANE_LOCATION

クラスタを削除します。

gcloud container clusters delete CLUSTER_NAME

次のステップ

GKE の TPU の詳細を確認する。
ノードの自動プロビジョニングの詳細を確認する。
GKE でバッチワークロードを実行するためのベストプラクティスの詳細を確認する。

TPU と Flex Start VM で小規模なバッチ ワークロードを実行する コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

Flex Start の料金

始める前に

Flex Start を使用してノードプールを作成する

単一ホスト TPU スライス ノードプールを作成する

マルチホスト TPU スライス ノードプールを作成する

Ironwood（TPU7x）

gcloud

Terraform

その他の TPU バージョン

gcloud

Terraform

コンソール

ノードプールで Flex Start のステータスを確認する

バッチ ワークロードを実行する

単一ホスト

マルチホスト

クリーンアップ

プロジェクトを削除する

個々のリソースを削除する

次のステップ

TPU と Flex Start VM で小規模なバッチワークロードを実行する

単一ホスト TPU スライスノードプールを作成する

マルチホスト TPU スライスノードプールを作成する

バッチワークロードを実行する