このページでは、費用対効果の高い大規模なトレーニングのために Cloud TPU マルチスライス構成を使用して、Google Kubernetes Engine（GKE）にワークロードをデプロイする方法について説明します。
このチュートリアルは、Kubernetes コンテナ オーケストレーションを使用して TPU で大規模なモデルのトレーニング、チューニング、推論ワークロードの管理を行う ML エンジニア、プラットフォーム管理者、オペレーターを対象としています。 Google Cloud のコンテンツで使用されている一般的なロールとタスクの例の詳細については、一般的な GKE ユーザーのロールとタスクをご覧ください。
GKE でマルチスライスを構成する前に、次のコンセプトを理解しておいてください。
TPU マルチスライスとは
TPU マルチスライスは、2 つ以上の Cloud TPU スライスがデータセンター ネットワーク（DCN）を介して通信する TPU スライス内の VM のアーキテクチャ組織です。マルチスライスは、フルスタックで費用対効果に優れた大規模なトレーニングを可能にします。最大数万の TPU チップまでほぼ線形にスケールアップできます。マルチスライス構成では、GKE は複数の TPU スライスにマルチスライス ワークロードをデプロイします。スライス内での 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 では、このドキュメントのコマンドを実行できない場合があります。
- Standard クラスタまたは、マルチスライスをサポートするバージョンを実行する Autopilot クラスタを作成します。サポートされているバージョンについては、GKE でのマルチスライスの可用性をご覧ください。
- GKE の Cloud TPU 用にプロジェクトに十分な割り当てがあることを確認します。
- v0.2.3 以降の JobSet をインストールします。
GKE の実装方法を選択する
TPU を使用するクラスタは、Google Cloud CLI または Accelerated Processing Kit（XPK）を使用して作成できます。
- gcloud CLI を使用して、GKE クラスタ インスタンスを手動で作成し、既存の GKE 本番環境を正確にカスタマイズまたは拡張します。
- XPK を使用すると、GKE クラスタをすばやく作成し、概念実証とテスト用のワークロードを実行できます。詳細については、XPK README をご覧ください。
GKE ノードプールを手動で作成する
このセクションでは、マルチスライスでワークロードを実行するように Standard モードのノードプールを準備する方法について説明します。GKE Autopilot モードを使用している場合は、マルチスライス ワークロードを実行するセクションに進みます。バージョン 1.29.2-gke.1521000 以降を実行する Autopilot クラスタでは、デフォルトで TPU が有効になります。
このセクションでは、次の手順について説明します。
- 3 つのマルチホスト TPU スライス ノードプールを作成します。
- ノードプールのステータスを確認する
- マルチスライス ワークロードを実行する
TPU スライス ノードプールを作成する
複数のマルチホストの TPU スライス ノードプールを作成できます。このガイドでは、マルチスライス ワークロードを実行する 3 つのマルチホスト TPU スライス ノードプールを作成します。マルチホスト TPU スライス ノードプールを作成する手順は、Ironwood（TPU7x）を使用するか、以前の TPU バージョンを使用するかによって異なります。
Ironwood（TPU7x）
Google Cloud CLI または Terraform を使用して、Ironwood（TPU7x）バージョンのマルチホスト TPU スライス ノードプールを作成できます。
gcloud
Ironwood（TPU7x）でマルチホスト TPU スライス ノードプールを作成するには、まずワークロード ポリシーを作成する必要があります。
ワークロード ポリシーを作成します。
gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \ --type=HIGH_THROUGHPUT \ --accelerator-topology=TPU_TOPOLOGY \ --project=PROJECT_ID \ --region=REGION次のように置き換えます。
WORKLOAD_POLICY_NAME: ワークロード ポリシーの名前。TPU_TOPOLOGY: TPU Ironwood（TPU7x）トポロジ。 例:2x2x2サポートされているすべての Ironwood（TPU7x）トポロジについては、トポロジのセクションをご覧ください。PROJECT_ID: 実際の Google Cloud プロジェクト ID。REGION: ワークロード ポリシーのリージョン。 ワークロード ポリシーはリージョン リソースであり、同じトポロジを共有するノードプール間で再利用できます。
ワークロード ポリシーを使用してノードプールを作成します。
gcloud container node-pools create NODE_POOL_NAME \ --cluster=CLUSTER_NAME \ --machine-type=tpu7x-standard-4t \ --placement-policy=WORKLOAD_POLICY_NAME \ --location=CONTROL_PLANE_LOCATION \ --node-locations=NODE_ZONE \ --project=PROJECT_ID \ --reservation=RESERVATION_NAME \ --reservation-affinity=specific次のように置き換えます。
NODE_POOL_NAME: 新しいノードプールの名前。CLUSTER_NAME: GKE クラスタの名前。WORKLOAD_POLICY_NAME: 作成したワークロード ポリシーの名前。CONTROL_PLANE_LOCATION: クラスタのコントロール プレーンの Compute Engine のロケーション。リージョン クラスタの場合はリージョン、ゾーンクラスタの場合はゾーンを指定します。NODE_ZONE: 使用する TPU バージョンに基づくゾーンの名前。省略可:
使用可能なロケーションについては、GKE での TPU の可用性をご覧ください。us-central1-ai1aなどの AI ゾーンを使用します。AI ゾーンは、 Google Cloud リージョン内の AI/ML ワークロード向けに最適化された特別なロケーションです。PROJECT_ID: 実際の Google Cloud プロジェクト ID。RESERVATION_NAME: 使用する予約の名前。
このコマンドでは、
--tpu-topologyフラグが--placement-policyフラグに置き換えられています。
Terraform
googleプロバイダのバージョン 4.84.0 以降を使用していることを確認します。ワークロード ポリシーを作成します。
resource "google_compute_resource_policy" { name = "WORKLOAD_POLICY_NAME" region = CLUSTER_LOCATION workload_policy { type = "HIGH_THROUGHPUT" accelerator_topology = "TPU_TOPOLOGY" } }次のように置き換えます。
WORKLOAD_POLICY_NAME: ワークロード ポリシーの名前。CLUSTER_LOCATION: クラスタのコンピューティングのロケーション。Kubernetes コントロール プレーンの信頼性を高めるため、リージョン クラスタを使用することをおすすめします。ゾーンクラスタを使用することもできます。詳細については、TPU のバージョンとトポロジの選択をご覧ください。TPU_TOPOLOGY: TPU Ironwood（TPU7x）トポロジ。 例:2x2x2サポートされているすべての Ironwood（TPU7x）トポロジについては、TPU を計画するをご覧ください。
google_compute_resource_policyリファレンスの詳細については、Terraform プロバイダをご覧ください。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] } flex_start = false spot = true } placement_policy { policy_name = WORKLOAD_POLICY_NAME } }次のように置き換えます。
NODE_POOL_RESOURCE_NAME: Terraform テンプレートのノードプール リソースの名前。PROJECT_ID: プロジェクト ID。CLUSTER_NAME: ノードプールを追加する既存のクラスタの名前。POOL_NAME: 作成するノードプールの名前。NODE_ZONES: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。省略可:
us-central1-ai1aなどの AI ゾーンを使用します。AI ゾーンは、 Google Cloud リージョン内の AI/ML ワークロード向けに最適化された特別なロケーションです。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 のチップ数で割った数と同じにする必要があります。たとえば、TPU_TOPOLOGYが2x2x2の場合、積は 8 です。tpu7x-standard-4tの各 VM には 4 つのチップがあるため、ノード数は 2 です。
spot: TPU スライスノードに Spot VM を使用するノードプール。 この設定は、ノードプールの作成後に変更することはできません。詳細については、Spot VM をご覧ください。flex_start: Flex Start 消費オプションを使用するノードプール。spotが有効になっている場合、この設定をtrueに設定することはできません。
その他の TPU バージョン
Google Cloud CLI、Terraform、または Google Cloud コンソールを使用して、バージョン v3、v4、v5p、v5e、Trillium（v6e）でマルチホスト TPU スライス ノードプールを作成できます。
gcloud
gcloud container node-pools create POOL_NAME \
--location=CONTROL_PLANE_LOCATION \
--cluster=CLUSTER_NAME \
--node-locations=NODE_ZONE \
--machine-type=MACHINE_TYPE \
--tpu-topology=TPU_TOPOLOGY \
[--num-nodes=NUM_NODES] \
[--spot \]
[--flex-start \]
[--enable-autoscaling \
--max-nodes MAX_NODES]
[--reservation-affinity=specific \
--reservation=RESERVATION_NAME] \
[--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
[--placement-type=COMPACT]
以下を置き換えます。
POOL_NAME: 新しいノードプールの名前。CONTROL_PLANE_LOCATION: 使用する TPU バージョンに基づくゾーンの名前。使用可能なロケーションについては、GKE での TPU の可用性をご覧ください。CLUSTER_NAME: クラスタの名前。NODE_ZONES: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。省略可:
us-central1-ai1aなどの AI ゾーンを使用します。AI ゾーンは、 Google Cloud リージョン内の AI/ML ワークロード向けに最適化された特別なロケーションです。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 をご覧ください。--flex-start: Flex Start VM を使用するようにノードプールを設定します。Flex Start VM は、GKE バージョン 1.33.0-gke.1712000 以降でサポートされている Flex Start 使用オプションを使用して作成されます。--enable-autoscaling: 自動スケーリングが有効なノードプールを作成します。GKE がマルチホスト TPU スライス ノードプールをスケーリングすると、ノードプールがゼロから最大サイズまでアトミックにスケールアップされます。MAX_NODES: ノードグループの最大サイズ。--enable-autoscalingが指定されている場合、--max-nodesフラグは必須です。TPU_TOPOLOGY（{A}x{B}x{C}）で定義された値の積を、各 VM のチップ数で割った数と同じにする必要があります。
--node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY: マルチホスト TPU スライス ノードプールがコレクションであることを GKE に伝えます。このフラグは、次の条件に該当する場合に使用します。- ノードプールが推論ワークロードを実行する。
- ノードプールが TPU Trillium を使用する。
- Spot VM が、コレクションのスケジュール設定をサポートしていない。
コレクションのスケジュール設定の管理の詳細については、マルチホスト TPU スライスでコレクションのスケジュール設定を管理するをご覧ください。
--placement-type=COMPACT: コンパクト プレースメントを有効にしてノードプールを作成します。このオプションは、--tpu-topologyフラグとともに使用する必要があります。詳細については、コンパクト プレースメント ポリシーを作成すると TPU トポロジをご覧ください。
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] } flex_start = false 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 つ以上のゾーンのカンマ区切りリスト。省略可:
us-central1-ai1aなどの AI ゾーンを使用します。AI ゾーンは、 Google Cloud リージョン内の AI/ML ワークロード向けに最適化された特別なロケーションです。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 をご覧ください。flex_start: Flex Start 使用オプションを使用するようにノードプールを設定します。spotが有効になっている場合、trueに設定することはできません。
コンソール
TPU を使用してノードプールを作成するには:
Google Cloud コンソールで [Google Kubernetes Engine] ページに移動します。
クラスタのリストで、変更するクラスタの名前をクリックします。
[ノード] タブをクリックします。
add_box [ユーザー管理のノードプールを作成] をクリックします。
[ノードプールの詳細] セクションで、[ノードのロケーションを指定する] チェックボックスをオンにします。
使用する TPU バージョンに基づくゾーンの名前を選択します。使用可能なロケーションについては、GKE での TPU の可用性をご覧ください。
ナビゲーション パネルで [ノード] をクリックします。
[マシンの構成] セクションで、[TPU] を選択します。
[シリーズ] プルダウン メニューで、次のいずれかを選択します。
- CT3: TPU v3、単一ホストデバイス
- CT3P: TPU v3、マルチホスト Pod スライス
- CT4P: TPU v4
- CT5LP: TPU v5e
- CT5P: TPU v5p
- CT6E: TPU Trillium（v6e）
[マシンタイプ] プルダウン メニューで、ノードに使用するマシンの名前を選択します。TPU バージョンを選択するの表で、マルチホスト TPU スライス ノードプールを作成するマシンタイプと TPU トポロジを定義する方法を確認します。
[TPU トポロジ] プルダウン メニューで、TPU スライスの物理トポロジを選択します。
[変更が必要です] ダイアログで [変更を適用] をクリックします。
[ブートディスクの種類] が [標準永続ディスク] と [SSD 永続ディスク] のいずれかであることを確認します。
必要に応じて、[Spot VM 上にノードを作成する] チェックボックスをオンにして、ノードプール内のノードで Spot VM を使用します。
[作成] をクリックします。
ノードプールのステータスを確認する
kubectlを使用してクラスタにアクセスできるように、認証情報を取得します。gcloud container clusters get-credentials CLUSTER_NAME \ --project=PROJECT_ID \ --location=CONTROL_PLANE_LOCATION次のように置き換えます。
CLUSTER_NAME: クラスタの名前。PROJECT_ID: プロジェクト ID。CONTROL_PLANE_LOCATION: クラスタのコントロール プレーンの Compute Engine におけるロケーション。リージョン クラスタの場合はリージョン、ゾーンクラスタの場合はゾーンを指定します。
Cloud Shell で
kubectlを使用して、TPU スライスノードを確認します。kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=ACCELERATOR_TYPE \ -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY次のように置き換えます。
ACCELERATOR_TYPE: TPU スライス ノードプールのアクセラレータ タイプ。この値は、ノードプールの作成時に使用したマシンタイプにマッピングする必要があります。たとえば、ct5lp-hightpu-4tマシンタイプでノードプールを作成した場合は、tpu-v5-lite-podsliceを使用します。使用するアクセラレータ タイプについては、GKE で TPU を計画するの表をご覧ください。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: 2 # This value defines a "Multislice" setup, by creating 2 independent TPU slices. template: spec: parallelism: 4 completions: 4 backoffLimit: 0 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 4x4 containers: - name: jax-tpu image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest ports: - containerPort: 8471 - containerPort: 8080 - containerPort: 8431 command: - bash - -c - | python -c 'import jax; print("Global device count:", jax.device_count())' sleep 60 resources: limits: google.com/tpu: 4 # This container requires 4 TPU chipsこのマニフェストでは、
4x4トポロジでtpu-v5-lite-podsliceを選択します。この構成は、次のことを意味します。- トポロジには合計 16 個のチップ（
4*4=16）があります。 tpu-v5-lite-podsliceは、VM あたり 4 つのチップを持つct5lp-hightpu-4tマシンにデプロイされます。つまり、4x4トポロジを使用するには、4 つの VM（16/4=4）が必要です。したがって、parallelismフィールドとcompletionsフィールドは4に設定する必要があります。- Pod 内の各コンテナは 4 つの TPU チップをリクエストします。これは、VM あたりのチップ数に対応します。これは、マニフェストで
google.com/tpu: 4行を使用して指定します。 replicasフィールドは2に設定され、2 つのスライスが作成されます。
これらの変数の可能な値の詳細については、GKE で TPU を計画するの表をご覧ください。
標準
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: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest ports: - containerPort: 8471 - containerPort: 8080 - containerPort: 8431 securityContext: privileged: true command: - bash - -c - | 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 スライス ノードプールのアクセラレータ タイプ。この値は、ノードプールの作成時に使用したマシンタイプにマッピングする必要があります。たとえば、ct5lp-hightpu-4tマシンタイプでノードプールを作成した場合は、tpu-v5-lite-podsliceを使用します。使用するアクセラレータ タイプについては、GKE で TPU を計画するの表をご覧ください。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: VM あたりのチップ数。
- トポロジには合計 16 個のチップ（
次のようにマニフェストを適用します。
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 2m30smultislice-jobJobSet は、Pod をスケジュール設定、作成、完了するまで実行します。Pod 名の形式は<jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>です。jobsetName接頭辞により、Pod が属する JobSet が決まります。省略可: JAX ワークロードを削除します。
kubectl delete -f tpu-multislice.yaml
XPK を使用して GKE リソースを作成する
- XPK README ファイルの手順を完了して、XPK をインストールします。
使用可能な TPU があるリージョンに Standard モードで GKE クラスタを作成します。次のコマンドは、指定された TPU リソースを含むノードプールを使用して GKE クラスタを作成します。
xpk cluster create \ --cluster=$CLUSTER_NAME \ --project=$PROJECT_ID \ --zone=$ZONE \ --tpu-type=ACCELERATOR_TYPE \ --num-slices=NUM_SLICES \ --reservation=$RESERVATION_NAME次のように置き換えます。
CLUSTER_NAME: クラスタの名前。PROJECT_ID: 実際の Google Cloud プロジェクト ID。ZONE: 使用する TPU バージョンに基づくゾーンの名前。使用可能なロケーションについては、GKE での TPU の可用性をご覧ください。予約済み容量を使用するには、容量を予約したゾーンを使用していることを確認してください。ACCELERATOR_TYPE: TPU のバージョンとタイプ。たとえば、Ironwood（TPU7x）にはtpu7x-standard-4tを使用します。NUM_SLICES: クラスタで使用する TPU スライスの数。マルチスライスの場合、1 より大きい値にする必要があります。RESERVATION_NAME: 予約の名前。
出力の最後の 2 行は次のようになります。
See your GKE Cluster here: https://console.cloud.google.com/kubernetes/clusters/details/<location>/CLUSTER_NAME/details?project=PROJECT_ID接続されている TPU デバイスの数を表示するワークロードを実行します。
xpk workload create \ --cluster=$CLUSTER_NAME \ --project=$PROJECT_ID \ --zone=$ZONE \ --tpu-type=ACCELERATOR_TYPE \ --num-slices=NUM_SLICES \ --docker-image=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest \ --workload=WORKLOAD_NAME \ --command="python -c 'import jax; print(\"TPU cores:\", jax.device_count())'"次のように置き換えます。
CLUSTER_NAME: クラスタの名前。PROJECT_ID: 実際の Google Cloud プロジェクト ID。ZONE: 選択したマシンのノードプールのゾーン。ACCELERATOR_TYPE: 選択した TPU タイプ。bash の予測入力を使用して、使用可能なオプションを特定します。NUM_SLICES: クラスタで構成された TPU スライスの数。これは、クラスタの作成時に使用された値と一致する必要があります。WORKLOAD_NAME: ワークロードの名前。
追加設定の構成
以降のセクションでは、マルチスライスに適用できる追加の構成について説明します。
hostNetwork を使用してネットワーク パフォーマンスを改善する
TPU スライス間のネットワーク パフォーマンスを改善するには、hostNetworking を有効にすることをおすすめします。Pod 仕様で hostNetwork: true を使用して、すべての Kubernetes ネットワーク スタックをスキップし、Kubernetes Pod が VM 間通信にホスト ネットワークを直接使用できるようにします。
hostNetworking を有効にするには、Pod 仕様に次の 2 行を追加します。
hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet
hostNetwork でワーカーノードの検出に podHostnames を引き続き使用するには、dnsPolicy: ClusterFirstWithHostNet を設定します。これは、トレーニングを自動再開するジョブを実行しているときに、同じチェックポイントを再読み込みするのに同じ名前である必要がある場合に重要です。
TPU Trillium（v6e）または Ironwood（TPU7x）（プレビュー）を使用し、Pod が hostNetworking を使用している場合は、次の DaemonSet をインストールして、ノードで /proc/sys/net/ipv4/tcp_rmem を調整します。
kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/51bf3dcab6ff658cf62cc32867f96860bf58dfdc/scripts/network-setup/v6e-increase-rmem.yaml
TPU Trillium または Ironwood（TPU7x）で hostNetwork を使用せずにネットワーク パフォーマンスを改善する
TPU Trillium または Ironwood（TPU7x）（プレビュー）を使用しているときに Pod で hostNetworking を使用できない場合は、GKE DRANET を実装して、ネットワーク パフォーマンスを最適化するためにネットワーク デバイスを割り当てます。GKE DRANET 機能を使用すると、Kubernetes と GKE Dataplane V2 をバイパスして、Pod の高性能ネットワーキング リソースをリクエストできます。
ct6e-standard-4t マシンタイプは 2 つの物理 NIC を基盤としています。Kubernetes には、Pod に渡すことのできない vNIC が 1 つ必要です。したがって、Pod が 2 つの vNIC に直接アクセスできるようにするには、各ノードに 3 つの vNIC が必要です。これにより、両方の物理 NIC で最高のパフォーマンスを実現できます。
DRANET 機能を有効にして ct6e-standard-4t にネットワーク デバイスを割り当てるには、次の操作を行います。
次のコマンドを使用して、追加の VPC を 2 つ作成します。
gcloud compute networks subnets create ADDITIONAL_SUBNET_1 \ --project=PROJECT_ID \ --network=ADDITIONAL_VPC_1 \ --location=CONTROL_PLANE_LOCATIONgcloud compute networks subnets create ADDITIONAL_SUBNET_2 \ --project=PROJECT_ID \ --network=ADDITIONAL_VPC_2 \ --location=CONTROL_PLANE_LOCATION作成したネットワークと GKE DRANET を有効にしてノードプールを作成します。
cloud.google.com/gke-networking-dra-driver=trueラベルを追加します。gcloud container node-pools create NODE_POOL_NAME \ --location=CONTROL_PLANE_LOCATION \ --cluster=CLUSTER_NAME \ --machine-type=MACHINE_TYPE \ --num-nodes=NUM_NODES \ --additional-node-network network=ADDITIONAL_VPC_1,subnetwork=ADDITIONAL_SUBNET_1 \ --additional-node-network network=ADDITIONAL_VPC_2,subnetwork=ADDITIONAL_SUBNET_2 \ --node-labels=cloud.google.com/gke-networking-dra-driver=trueノードプールのステータスを確認したら、
ResourceClaimTemplateを適用して、追加のネットワーク デバイスを含むワークロードを作成します。apiVersion: resource.k8s.io/v1 kind: ResourceClaimTemplate metadata: name: two-netdev spec: spec: devices: requests: - name: req-netdev exactly: deviceClassName: netdev.google.com allocationMode: ExactCount count: 2 --- apiVersion: v1 kind: Pod metadata: name: sample-netdevice-pod-1 spec: containers: - name: sample-netdevice-pod image: busybox command: ["sleep", "infinity"] resources: claims: - name: netdev limits: google.com/tpu: 4 restartPolicy: Always resourceClaims: - name: netdev resourceClaimTemplateName: two-netdevインターフェースが Pod で使用可能であることを確認します。
kubectl exec -it pods/sample-netdevice-pod-1 -- ip a
ネットワーキングのパフォーマンスを向上させるには、eth0 インターフェースではなく、eth1 インターフェースと eth2 インターフェースを使用します。これを行うには、ワークロード仕様に export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" を追加します。
ロギングを有効にする
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_latencieskubernetes.io/container/multislice/network/collective_end_to_end_latencieskubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencieskubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies
TPU スライスとマルチスライス
次の表は、TPU スライスとマルチスライスのアーキテクチャ組織の差を示しています。
| TPU スライス | マルチスライス | |
|---|---|---|
| 相互接続 | ワークロードは単一の TPU スライスで実行されます。スライス内のすべての TPU チップは ICI で接続されます。 | ワークロードは複数の TPU スライスで実行されます。スライス内の通信は ICI 上で行われます。スライス間の通信は DCN 上で行われます。 |
| サポートされているノードプール | 単一ホスト TPU スライスとマルチホスト TPU スライス | マルチホスト TPU スライスのグループ |
| 推奨されるワークロード タイプ | IndexedJob または JobSet | JobSet |
リソースをクリーンアップする
課金を防ぐ最も簡単な方法は、チュートリアル用に作成した Google Cloud プロジェクトを削除することです。また、リソースを個別に削除することもできます。
プロジェクトを削除する
Google Cloud プロジェクトを削除します。
gcloud projects delete PROJECT_ID
リソースを個別に削除する
GKE クラスタを削除します。
```sh
gcloud container clusters delete CLUSTER_NAME \
--project=PROJECT_ID \
--location=CONTROL_PLANE_LOCATION
```
XPK で作成されたリソースを削除する
XPK を使用してクラスタを作成した場合は、追加料金が発生しないようにクラスタを削除します。
xpk cluster delete \
--cluster=$CLUSTER_NAME \
--project=$PROJECT_ID \
--zone=$ZONE
次のように置き換えます。
CLUSTER_NAME: クラスタの名前。PROJECT_ID: 実際の Google Cloud プロジェクト ID。ZONE: クラスタが作成されたゾーン。