GKE に TPU マルチスライスをデプロイする

このページでは、Google Kubernetes Engine（GKE）の TPU マルチスライス構成について説明します。GKE でマルチスライスを構成する前に、次のコンセプトを理解しておく必要があります。

  1. Cloud TPU の概要
  2. Cloud TPU システム アーキテクチャ
  3. GKE の TPU について

TPU マルチスライスとは

TPU マルチスライスは、2 つ以上の Cloud TPU スライスがデータセンター ネットワーク（DCN）上で通信する TPU VM のアーキテクチャ組織です。マルチスライスは、フルスタックで費用対効果に優れた大規模なトレーニングを可能にします。最大数万の TPU チップまでほぼ線形にスケーリングできます。マルチスライス構成では、GKE は複数の TPU スライスにマルチスライス ワークロードをデプロイします。スライス内のチップ間の通信は、チップ間相互接続（ICI）上で行われます。スライス間の通信は DCN 上で行われます。

Job が大きすぎて 1 つの TPU スライスに収まらない場合は、マルチスライスを使用することをおすすめします。

GKE でのマルチスライスの可用性

  • GKE は、バージョン 1.27.4-gke.900 以降でマルチスライスをサポートしています。
  • マルチスライスは、JAX と PyTorch フレームワークをサポートしています。サポートされている最小の JAX のバージョンは 2.1 です。
  • マルチスライスは、マルチホスト TPU スライス ノードプールのみをサポートします。たとえば、2x2x1 トポロジの ct4p-hightpu-4t や、2x2 トポロジの ct5lp-hightpu-4t ではマルチスライスを使用できません。これらは、単一ホスト TPU スライス ノードプールであるためです。
  • マルチスライスは、同期マルチコントローラ トレーニングのみをサポートします。
  • マルチスライス ワークロードは、同じ TPU タイプ、サイズ、トポロジを共有する TPU スライス間でのみ実行できます。

始める前に

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

  • Google Kubernetes Engine API を有効にします。
    • Google Kubernetes Engine API の有効化
  • このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化します。すでに gcloud CLI をインストールしている場合は、gcloud components update を実行して最新のバージョンを取得します。

マルチスライスでワークロードを実行する

このセクションでは、次の手順について説明します。

  1. 3 つのマルチホスト TPU スライス ノードプールを作成します。
  2. JobSet を使用してマルチスライスでワークロードを実行します。

TPU ノードプールを作成する

複数のマルチホストの TPU ノードプールを作成できます。このガイドでは、マルチスライス ワークロードを実行する 3 つのマルチホスト 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 を使用します。詳細については、TPU のバージョンとトポロジを選択するをご覧ください。
  • CLUSTER_NAME: クラスタの名前。
  • NODE_ZONE: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
  • MACHINE_TYPE: ノードに使用するマシンのタイプ。使用可能なマシンタイプの詳細については、TPU 構成のマッピングをご覧ください。
  • TPU_TOPOLOGY: TPU スライスの物理トポロジ。トポロジの形式は、次のように TPU のバージョンによって異なります。
    • TPU v4: トポロジを 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: ノードグループの最大サイズ。--max-nodes フラグは、--enable-autoscaling が指定されている場合は必須で、TPU_TOPOLOGY（{A}x{B}x{C}）で定義された値の積を、各 VM のチップ数で割った数と同じにする必要があります。

Terraform

  1. google プロバイダのバージョン 4.84.0 以降を使用していることを確認します。

  2. 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 を使用してノードプールを作成するには:

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

    Google Kubernetes Engine に移動

  2. クラスタのリストで、変更するクラスタの名前をクリックします。

  3. [ノードプールを追加] をクリックします。

  4. [ノードプールの詳細] セクションで、[ノードのロケーションを指定する] チェックボックスをオンにします。

  5. 使用する 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 を使用します。

  6. ナビゲーション パネルで [ノード] をクリックします。

  7. [マシンの構成] セクションで、[TPU] を選択します。

  8. [シリーズ] プルダウン メニューで、次のいずれかを選択します。

    • CT4P: TPU v4 の場合。
    • CT5LP: TPU v5e の場合。

  9. [マシンタイプ] プルダウン メニューで、ノードに使用するマシンの名前を選択します。TPU 構成のマッピングの表で、マルチホスト TPU ノードプールを作成するマシンタイプと TPU トポロジを定義する方法を確認します。

  10. [TPU トポロジ] プルダウン メニューで、TPU スライスの物理トポロジを選択します。

  11. [変更が必要です] ダイアログで [変更を適用] をクリックします。

  12. [ブートディスクの種類] が [標準永続ディスク] と [SSD 永続ディスク] のいずれかであることを確認します。

  13. 必要に応じて、[Spot VM 上にノードを作成する] チェックボックスをオンにして、ノードプール内のノードで Spot VM を使用します。

  14. [作成] をクリックします。

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

  1. kubectl を使用してクラスタにアクセスできるように、認証情報を取得します。

    gcloud container clusters get-credentials CLUSTER_NAME \
    --project=PROJECT_ID

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

    • CLUSTER_NAME: クラスタの名前。
    • PROJECT_ID: プロジェクト ID。

  2. Cloud Shell で kubectl を使用して、TPU ノードを表示します。

    kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=TPU_ACCELERATOR \
   -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY

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

    • TPU_ACCELERATOR: ノードプールの作成時に使用した TPU アクセラレータのタイプ。たとえば、tpu-v4-podslice、tpu-v5-lite-device、tpu-v5-lite-podslice です。
    • TPU_TOPOLOGY: TPU スライスの物理トポロジ。

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

     NAME                                    STATUS   ROLES    AGE    VERSION
 gke-tpu-20ee2cce-5tv6                   Ready    <none>   34h     v1.28.1-gke.1066000

マルチスライス ワークロードを実行する

このセクションでは、TPU スライス内の TPU チップのグローバル数を示す JAX ワークロードを実行してから終了します。

JAX ワークロードを実行する手順は次のとおりです。

  1. 次の tpu-multislice.yaml マニフェストを作成します。

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: NUM_SLICES
      template:
        spec:
          parallelism: NUM_NODES
          completions: NUM_NODES
          backoffLimit: 0
          template:
            spec:
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                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("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                limits:
                  google.com/tpu: NUM_CHIPS

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

    • NUM_SLICES: TPU ノードプールの数。この場合、NUM_SLICES は 3 です。
    • TPU_ACCELERATOR: ノードプールの作成時に使用した TPU アクセラレータのタイプ。たとえば、tpu-v4-podslice、tpu-v5-lite-device、tpu-v5-lite-podslice です。
    • TPU_TOPOLOGY: TPU スライスの物理トポロジ。たとえば、TPU のバージョンに応じて 4x4x4 または 2x2 を使用します。
    • NUM_NODES: ノードプール内のノード数。ゼロにするか、TPU_TOPOLOGY（{A}x{B}x{C}）で定義された値の積を各 VM のチップ数で割った数と同じにする必要があります。マルチホスト TPU v4 の場合、各 VM のチップ数は 4 です。マルチホスト TPU v5e の場合、各 VM のチップ数は 1、4、または 8 です。したがって、TPU_TOPOLOGY が 2x4x4（各 VM に 4 つのチップがある TPU v4）の場合、NUM_NODES は 32÷4 で 8 になります。
    • NUM_CHIPS: マルチホスト TPU v4 の場合、各 VM のチップ数は 4 です。マルチホスト TPU v5e の場合、各 VM のチップ数は 1、4、または 8 です。詳細については、TPU VM 上の TPU チップをご覧ください。

    このマニフェストの内容:

    • JobSet は、JobSet 名と同じ名前（この場合は multislice-job）のヘッドレス サービスです。
    • maxRestarts: 4 は、子 Job が失敗したときに GKE が JobSet を再起動する最大回数を示します。JobSet の再起動回数が定義された最大値に達すると、JobSet は失敗としてマークされます。
    • parallelism フィールドと completions フィールドは、各ノードプール内のノード数と同じです。
    • マルチスライスは同期マルチコントローラ トレーニングのみをサポートするため、backoff は 0 です。0 に設定する必要があります。Pod が失敗した場合にジョブを失敗させます。
    • アフィニティ セクションの値により、マルチスライスのグループ内で TPU マルチスライス ワークロードが 1 つだけ実行されることが保証されます。
    • containerPort: 8080 は MXLA コーディネーターのポートです。
    • securityContext: privileged: true は、ノードに TPU にアクセスする特権があることを示します。GKE バージョン 1.28 以降のノードでは、TPU にアクセスする特権がある必要はありません。詳細については、特権モードなしでコンテナを実行するをご覧ください。

  2. 次のようにマニフェストを適用します。

    kubectl apply -f tpu-multislice.yaml

  3. ワークロードが許可されたことを確認します。

    kubectl get jobsets

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

    NAME            RESTARTS   COMPLETED   AGE
multislice-job                         3s

  4. プロビジョニングされた Pod のステータスをモニタリングします。

    kubectl get pods

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

     NAME                                READY   STATUS      RESTARTS   AGE
 multislice-job-slice-0-0-wzq9t      0/1     Completed   0          2m31s
 multislice-job-slice-0-1-zf4dp      0/1     Completed   0          2m30s
 multislice-job-slice-1-0-hbfn5      0/1     Completed   0          2m31s
 multislice-job-slice-1-1-45fgl      0/1     Completed   0          2m30s
 multislice-job-slice-2-0-wjbp4      0/1     Completed   0          2m30s
 multislice-job-slice-2-1-lwnvs      0/1     Completed   0          2m30s

multislice-job JobSet は、Pod をスケジュール設定、作成、完了するまで実行します。Pod 名の形式は <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix> です。jobsetName 接頭辞により、Pod が属する JobSet が決まります。

その他の構成

以降のセクションでは、マルチスライスに適用できる追加の構成について説明します。

GKE Pod で hostNetwork を有効にする

TPU スライス間のネットワーク パフォーマンスを改善するには、hostNetworking を有効にすることをおすすめします。Pod 仕様で hostNetwork: true を使用して、すべての Kubernetes ネットワーク スタックをスキップし、Kubernetes Pod が VM 間通信にホスト ネットワークを直接使用できるようにします。

hostNetworking を有効にするには、Pod 仕様から次の 2 行を削除します。

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

hostNetwork でワーカーノードの検出に podHostnames を引き続き使用するには、dnsPolicy: ClusterFirstWithHostNet を設定します。これは、トレーニングを自動再開するジョブを実行しているときに、同じチェックポイントを再読み込みするのに同じ名前である必要がある場合に重要です。

ロギング

TPU VM を含む GKE ノードで実行されるコンテナによって出力されたログは、クラスタで GKE システム ロギングが有効になっている場合にログ エクスプローラに表示されます。

次のフィルタでログ エクスプローラを使用して、GKE からのログを表示して、ワークロードのコンテナログを表示できます。

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME

TPU スライスとワーカーには、次のフィルタを使用します。

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>

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

一般的な TPU 指標に加えて、マルチスライス固有の TPU ランタイム指標が 4 つあります。これらの指標は、GKE バージョン 1.29.1-gke.1016000 以降で使用できます。TPU ワークロードは JAX バージョン 0.4.24 を使用する必要があります。

利用可能なマルチスライス指標は次のとおりです。

  • DCN（データセンター ネットワーク）転送レイテンシ: マルチスライス トラフィックのネットワーク転送レイテンシの分布。
  • 全体的なレイテンシ: マルチスライス トラフィックのエンドツーエンドの全体的なレイテンシの分布。
  • ホストからデバイスへの転送レイテンシ: マルチスライス トラフィックのデータチャンクごとの、ホストからデバイスへの転送レイテンシの分布。
  • デバイス間転送のレイテンシ: マルチスライス トラフィックのデータチャンクごとの、デバイスからホストへの転送レイテンシの分布。

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

  • kubernetes.io/container/multislice/network/dcn_transfer_latencies
  • kubernetes.io/container/multislice/network/collective_end_to_end_latencies
  • kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies
  • kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies

TPU スライスとマルチスライス

次の表は、TPU スライスとマルチスライスのアーキテクチャ組織の差を示しています。

TPU スライス マルチスライス
相互接続 ワークロードは単一の TPU スライスで実行されます。スライス内のすべての TPU チップは ICI で接続されます。 ワークロードは複数の TPU スライスで実行されます。スライス内の通信は ICI 上で行われます。スライス間の通信は DCN 上で行われます。
サポートされているノードプール 単一ホスト TPU スライスとマルチホスト TPU スライス マルチホスト TPU スライスのグループ
推奨されるワークロード タイプ IndexedJob または JobSet JobSet

次のステップ