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

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

このガイドは、バッチ ワークロードの実行に Kubernetes コンテナ オーケストレーション機能を使用することを検討している ML エンジニア、プラットフォーム管理者、オペレーター、データおよび AI のスペシャリストを対象としています。 Cloud de Confiance by S3NS のコンテンツで使用されている一般的なロールとタスクの例の詳細については、一般的な 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 では、このドキュメントのコマンドを実行できない場合があります。

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

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

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

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

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

1. Flex Start を使用してノードプールを作成します。

   単一ホスト

   gcloud container node-pools create NODE_POOL_NAME \
       --cluster=CLUSTER_NAME \
       --location=LOCATION_NAME \
       --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: クラスタの名前。
   • LOCATION_NAME: クラスタ コントロール プレーンのコンピューティング リージョン。
   • NODE_ZONES: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
   • MACHINE_TYPE: ノードに使用するマシンのタイプ。TPU 互換マシンタイプの詳細については、TPU バージョンを選択するの表をご覧ください。

   マルチホスト

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

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

   • NODE_POOL_NAME: ノードプールに付ける名前。
   • CLUSTER_NAME: クラスタの名前。
   • LOCATION_NAME: クラスタ コントロール プレーンのコンピューティング リージョン。
   • NODE_ZONES: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
   • MACHINE_TYPE: ノードに使用するマシンのタイプ。たとえば、TPU Trillium には ct6e-standard-4t を使用できます。使用可能なマシンタイプの詳細については、TPU バージョンを選択するをご覧ください。
   • TPU_TOPOLOGY: TPU スライスの物理トポロジ。トポロジの形式は TPU のバージョンによって異なります。TPU トポロジの詳細については、トポロジを選択するの表をご覧ください。

   上記のコマンドでは、Flex Start VM を使用してノードプールを作成するときに次の必須フラグを使用します。

   • --enable-autoscaling: Flex Start は、ワークロードの実行時に必要なコンピューティング リソースのみをプロビジョニングします。次のパラメータを設定する必要があります。

     • --num-nodes=0
     • --min-nodes=0

     • --max-nodes は、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 では予約を使用することも、必要とすることもありません。

2. ノードプールで Flex-Start のステータスを確認します。

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

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

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

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

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

2. 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 のバージョンをご覧ください。

3. dws-flex-start.yaml マニフェストを適用します。

   kubectl apply -f dws-flex-start.yaml
   

4. 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>
   

クリーンアップ

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

プロジェクトを削除する

1. In the Cloud de Confiance console, go to the Manage resources page.

   Go to Manage resources

2. In the project list, select the project that you want to delete, and then click Delete.
3. In the dialog, type the project ID, and then click Shut down to delete the project.

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

1. ジョブを削除します。

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

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

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

3. クラスタを削除します。

   gcloud container clusters delete CLUSTER_NAME
   

次のステップ

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