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

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

このチュートリアルは、Kubernetes コンテナ オーケストレーションを使用して TPU で大規模なモデルのトレーニング、チューニング、推論ワークロードの管理を行う ML エンジニア、プラットフォーム管理者、オペレーターを対象としています。 Google Cloud のコンテンツで使用されている一般的なロールとタスクの例の詳細については、一般的な GKE Enterprise ユーザーロールとタスクをご覧ください。

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

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

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 を実行して最新のバージョンを取得する。

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

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

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

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

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

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_ZONES \
    --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] \
    [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
    [--placement-type=COMPACT]

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

  • POOL_NAME: 新しいノードプールの名前。
  • LOCATION: 使用する TPU バージョンに基づくゾーンの名前。使用可能なロケーションについては、GKE での TPU の可用性をご覧ください。
  • CLUSTER_NAME: クラスタの名前。
  • NODE_ZONES: 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 のチップ数で割った数と同じにする必要があります。

  • --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

  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 トポロジの詳細については、トポロジを選択するの表をご覧ください。

    必要に応じて、次の変数を使用することもできます。

    • 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 バージョンに基づくゾーンの名前を選択します。使用可能なロケーションについては、GKE での TPU の可用性をご覧ください。

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

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

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

    • CT3P: TPU v3 の場合。
    • 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 \
    --zone=ZONE

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

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

  2. 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

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

    • TPU_ACCELERATOR: ノードプールの作成時に使用した TPU アクセラレータのタイプ。たとえば、tpu-v4-podslice や 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 マニフェストを作成します。

    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.10
                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.10
                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-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 にアクセスする特権がある必要はありません。詳細については、特権モードなしでコンテナを実行するをご覧ください。

  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 が決まります。

  5. 省略可: JAX ワークロードを削除します。

    kubectl delete -f tpu-multislice.yaml

追加設定を構成する

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

hostNetwork を使用してネットワーク パフォーマンスを改善する

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

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

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

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

TPU Trillium（v6e）を使用し、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 で hostNetwork を使用せずにネットワーク パフォーマンスを改善する

TPU Trillium を使用しているときに Pod で hostNetworking を使用できない場合は、netdevice モードでマルチネットワーキングを有効にして、ネットワーク パフォーマンスを最適化します。マルチネットワークで netdevice モードの NIC がサポートされていることにより、Kubernetes と GKE Dataplane V2 をバイパスして、VM NIC が Pod に直接渡されます。

ct6e-standard-4t マシンタイプは 2 つの物理 NIC を基盤としています。Kubernetes には、Pod に渡すことのできない vNIC が 1 つ必要です。したがって、Pod が 2 つの vNIC に直接アクセスできるようにするには、各ノードに 3 つの vNIC が必要です。これにより、両方の物理 NIC で最高のパフォーマンスを実現できます。

ct6e-standard-4t の netdevice モードを有効にするには、次の操作を行います。

  1. netdevice モードをサポートする VPC を 2 つ追加作成する
  2. マルチネットワーク機能を備えた GKE クラスタを作成する

  3. 2 つの netdevice ネットワークを構成する。たとえば、次の GKENetworkParamSet オブジェクトと Network オブジェクトを使用できます（SECOND_VPC と THIRD_VPC は前の手順で作成した VPC です）。

    apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: tpu-second
spec:
  vpc: SECOND_VPC
  vpcSubnet: SECOND_VPC_SUBNET
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: tpu-third
spec:
  vpc: THIRD_VPC
  vpcSubnet: SECOND_VPC_SUBNET
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: tpu-second
spec:
  provider: "GKE"
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: tpu-second
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: tpu-third
spec:
  provider: "GKE"
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: tpu-third

  4. Pod を 3 つのネットワークに接続する。たとえば、Pod 仕様で次のアノテーションを使用できます。

    metadata:
  annotations:
    networking.gke.io/default-interface: 'eth0'
    networking.gke.io/interfaces: |
      [
        {"interfaceName":"eth0","network":"default"},
        {"interfaceName":"eth1","network":"tpu-second"},
        {"interfaceName":"eth2","network":"tpu-third"},
      ]

  5. Pod 内のネットワーク sysctl を、初期コンテナまたはアプリケーション コンテナに適用します。たとえば、次の init コンテナを Pod 仕様に追加できます。

    initContainers:
- name: "network-optimization-sysctls"
  image: "busybox"
  securityContext:
    privileged: true
  command:
  - bash
  - -c
  - |
    echo 5000 > /proc/sys/net/ipv4/tcp_rto_min_us
    echo 1 > /proc/sys/net/ipv4/tcp_no_metrics_save
    echo 0 > /proc/sys/net/ipv4/tcp_slow_start_after_idle
    echo 131072 > /proc/sys/net/core/optmem_max
    echo "4096 41943040 314572800" > /proc/sys/net/ipv4/tcp_rmem
ベスト プラクティス:

ネットワーキングのパフォーマンスを向上させるには、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_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

リソースをクリーンアップする

課金を防ぐ最も簡単な方法は、チュートリアル用に作成した Google Cloud プロジェクトを削除することです。また、リソースを個別に削除することもできます。

プロジェクトを削除する

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

リソースを個別に削除する

GKE クラスタを削除します。

```sh
gcloud container clusters delete  CLUSTER_NAME \
   --project=PROJECT_ID  \
   --zone=ZONE
```

次のステップ