このページでは、Google Kubernetes Engine(GKE)の TPU マルチスライス構成について説明します。GKE でマルチスライスを構成する前に、次のコンセプトを理解しておく必要があります。
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 スライス間でのみ実行できます。
始める前に
始める前に、次の作業が完了していることを確認してください。
- Google Kubernetes Engine API を有効にする。 Google Kubernetes Engine API の有効化
- このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。すでに gcloud CLI をインストールしている場合は、
gcloud components update
を実行して最新のバージョンを取得する。
- 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 バージョンに基づくゾーンの名前。- 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
を使用します。
詳細については、GKE での TPU の提供状況をご覧ください。
- TPU v4 の場合は、
CLUSTER_NAME
: クラスタの名前。NODE_ZONE
: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。MACHINE_TYPE
: ノードに使用するマシンのタイプ。使用可能なマシンタイプの詳細については、TPU 構成のマッピングをご覧ください。TPU_TOPOLOGY
: TPU スライスの物理トポロジ。トポロジの形式は、次のように TPU のバージョンによって異なります。- TPU v4 または v5p: トポロジを 3 タプル(
{A}x{B}x{C}
)で定義します。例:4x4x4
。 - TPU v5e: トポロジを 2 タプル(
{A}x{B}
)で定義します。例:2x2
。
詳細については、トポロジをご覧ください。
- TPU v4 または v5p: トポロジを 3 タプル(
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 v4 の場合: トポロジを 3 タプル(
{A}x{B}x{C}
)で定義します。例:4x4x4
。 - TPU v5e の場合: トポロジを 2 タプル(
{A}x{B}
)で定義します。例:2x2
。
- TPU v4 の場合: トポロジを 3 タプル(
必要に応じて、次の変数を使用することもできます。
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 のページに移動します。
クラスタのリストで、変更するクラスタの名前をクリックします。
[add_boxノードプールを追加] をクリックします。
[ノードプールの詳細] セクションで、[ノードのロケーションを指定する] チェックボックスをオンにします。
使用する 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 v4 の場合は、
ナビゲーション パネルで [ノード] をクリックします。
[マシンの構成] セクションで、[TPU] を選択します。
[シリーズ] プルダウン メニューで、次のいずれかを選択します。
- 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 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 コーディネーターのポートです。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 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 |