このページは Cloud Translation API によって翻訳されました。

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

Autopilot Standard

このページでは、費用対効果の高い大規模なトレーニングのために Cloud TPU マルチスライス構成を使用して、Google Kubernetes Engine（GKE）にワークロードをデプロイする方法について説明します。

GKE でマルチスライスを構成する前に、次のコンセプトを理解しておいてください。

重要な用語: このページでは、次の用語について説明します。

データセンターネットワーキング（DCN）: マルチスライス構成で Cloud TPU スライスを接続する、大容量、低レイテンシ、低スループットのネットワーク。
TPU スライスノード: 相互接続された TPU チップを備えた VM のセットを含む Kubernetes ノード。TPU スライスの各 VM には、1、4、8 個の TPU チップを含めることができます。詳細については、システムアーキテクチャをご覧ください。
マルチスライスワークロード: 複数の TPU スライスを必要とするワークロード。
TPU トポロジ: TPU スライス内の TPU チップの数と物理的な配置を定義します。詳しくは、GKE での TPU についてをご覧ください。

TPU マルチスライスとは

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

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

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

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

始める前に

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

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 など）を防止できます。

Standard クラスタまたは、マルチスライスをサポートするバージョンを実行する Autopilot クラスタを作成します。サポートされているバージョンについては、GKE でのマルチスライスの可用性をご覧ください。
GKE の Cloud TPU 用にプロジェクトに十分な割り当てがあることを確認します。
v0.2.3 以降の JobSet をインストールします。

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

このセクションでは、マルチスライスでワークロードを実行する方法について説明します。GKE Autopilot モードを使用している場合は、マルチスライスワークロードを実行するセクションに進みます。バージョン 1.29.2-gke.1521000 以降を実行する Autopilot クラスタでは、デフォルトで TPU が有効になります。

Standard モードのノードプールを準備する

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

3 つのマルチホスト TPU スライスノードプールを作成します。
ノードプールのステータスを確認します

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 バージョンに基づくゾーンの名前。使用可能なロケーションを確認するには、GKE での TPU の可用性をご覧ください。
CLUSTER_NAME: クラスタの名前。
NODE_ZONE: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
MACHINE_TYPE: ノードに使用するマシンのタイプ。使用可能なマシンタイプの詳細については、TPU バージョンを選択するをご覧ください。
TPU_TOPOLOGY: TPU スライスの物理トポロジ。トポロジの形式は TPU のバージョンによって異なります。TPU トポロジの詳細については、トポロジを選択するの表をご覧ください。

詳細については、トポロジをご覧ください。
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 トポロジの詳細については、トポロジを選択するの表をご覧ください。
必要に応じて、次の変数を使用することもできます。
- 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 バージョンに基づいてゾーンの名前を選択します。使用可能なロケーションを確認するには、GKE での TPU の可用性をご覧ください。
ナビゲーションパネルで [ノード] をクリックします。
[マシンの構成] セクションで、[TPU] を選択します。
[シリーズ] プルダウンメニューで、次のいずれかを選択します。
- CT3P: TPU v3 の場合。
- CT4P: TPU v4 の場合。
- CT5LP: TPU v5e の場合。
[マシンタイプ] プルダウンメニューで、ノードに使用するマシンの名前を選択します。TPU バージョンを選択するの表で、マルチホスト TPU スライスノードプールを作成するマシンタイプと TPU トポロジを定義する方法を確認します。
[TPU トポロジ] プルダウンメニューで、TPU スライスの物理トポロジを選択します。
[必要な変更] ダイアログで [変更する] をクリックします。
[ブートディスクの種類] が [標準永続ディスク] と [SSD 永続ディスク] のいずれかであることを確認します。
必要に応じて、Spot VM でノードを有効にするチェックボックスをオンにして、ノードプール内のノードで Spot VM を使用します。
[作成] をクリックします。

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

kubectl を使用してクラスタにアクセスできるように、認証情報を取得します。
```
gcloud container clusters get-credentials CLUSTER_NAME \
    --project=PROJECT_ID
```
次のように置き換えます。
- CLUSTER_NAME: クラスタの名前。
- PROJECT_ID: プロジェクト ID。
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 ワークロードを実行する手順は次のとおりです。

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

Autopilot

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:
              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
                - containerPort: 8431
                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

Standard

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
                - containerPort: 8431
                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 です。
ACCELERATOR_TYPE: ノードプールの作成時に使用した 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 チップ数で割った数のプロダクトにする必要があります。マルチホスト 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）のヘッドレスサービスです。
alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool アノテーションは、すべての Pod が同じスライスにスケジュールされるように Pod アフィニティを構成します。
maxRestarts: 4 は、子 Job が失敗したときに GKE が JobSet を再起動する最大回数を示します。JobSet の再起動回数が定義された最大値に達すると、JobSet は失敗としてマークされます。
parallelism フィールドと completions フィールドは、各ノードプール内のノード数と同じです。
マルチスライスは同期マルチコントローラトレーニングのみをサポートするため、backoff は 0 です。0 に設定する必要があります。Pod が失敗した場合にジョブを失敗させます。
アフィニティセクションの値により、マルチスライスのグループ内で TPU マルチスライスワークロードが 1 つだけ実行されることが保証されます。
containerPort: 8080 は MXLA コーディネーターのポートです。
containerPort: 8431 は、TPU 使用状況の指標をエクスポートするポートです。
securityContext: privileged: true は、ノードに TPU にアクセスする特権があることを示します。GKE バージョン 1.28 以降のノードでは、TPU にアクセスする特権がある必要はありません。詳細については、特権モードなしでコンテナを実行するをご覧ください。

次のようにマニフェストを適用します。
```
kubectl apply -f tpu-multislice.yaml
```

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

kubectl get jobsets

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

NAME            RESTARTS   COMPLETED   AGE
multislice-job                         3s

プロビジョニングされた 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 Standard Pod で hostNetwork を有効にする

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

hostNetworking を有効にするには、Pod 仕様に次の 2 行を追加します。

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

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

ロギング

TPU スライスノードを含む 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>

詳細については、GKE TPU ログを表示するをご覧ください。

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

一般的な 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

次のステップ

TPU スライスでマルチスライスワークロードをオーケストレーションする方法を学習する