このページの情報の一部またはすべては、S3NS の Trusted Cloud には適用されない場合があります。

TPU と Flex Start プロビジョニングモードで小規模なバッチワークロードを実行する

Autopilot Standard

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

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

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 がすでにインストールされている場合には、必ず 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 を使用してノードプールを作成する

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

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

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

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 でノードプールを作成するときに次の必須フラグを使用します。
- --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 では予約を使用することも、必要とすることもありません。
ノードプールで 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 で TPU ノードのスケジュールを設定する Job を作成します。Kubernetes の Job コントローラは、1 つ以上の Pod を作成し、特定のタスクが正常に実行されるようにします。

Trusted Cloud コンソールで、 [Cloud Shell をアクティブにする] をクリックして、Cloud Shell セッションを起動します。 Trusted 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: 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:
            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>

クリーンアップ

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

プロジェクトを削除する

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

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

In the Trusted 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 LOCATION_NAME

クラスタを削除します。

gcloud container clusters delete CLUSTER_NAME

次のステップ

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

TPU と Flex Start プロビジョニング モードで小規模なバッチ ワークロードを実行する

Flex Start の料金

始める前に

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

単一ホスト

マルチホスト

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

単一ホスト

マルチホスト

クリーンアップ

プロジェクトを削除する

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

次のステップ

TPU と Flex Start プロビジョニングモードで小規模なバッチワークロードを実行する

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