GKE Autopilot に TPU ワークロードをデプロイする

このページでは、Google Kubernetes Engine（GKE）Autopilot クラスタに Cloud TPU アクセラレータ（TPU）を使用してデプロイすることで、ML ワークロードを高速化する方法について説明します。次のコンセプトに精通している必要があります。

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

Autopilot での TPU の仕組み

Autopilot ワークロードで TPU を使用するには、ワークロード マニフェストで TPU バージョンと、その TPU バージョンでサポートされているトポロジをリクエストします。次に、Kubernetes の resources.requests フィールドと resources.limits フィールドを使用して、使用するワークロードの TPU チップ数を指定します。ワークロードをデプロイすると、GKE はリクエストされた TPU 構成を持つノードをプロビジョニングし、ノードに Pod をスケジュールします。GKE は各ワークロードを独自のノードに配置するため、各 Pod は中断のリスクを最小限に抑えながら、ノードのすべてのリソースにアクセスできます。

Autopilot の TPU は、次の機能に対応しています。

  1. Spot Pod
  2. 特定の容量予約
  3. 拡張ランタイム Pod

TPU 構成を計画する

TPU をリクエストする前に、ワークロードの CPU とメモリの要件に基づいて、必要な構成を決定します。次のことを決める必要があります。

  • TPU バージョン: 特定の Cloud TPU バージョン（v5e など）。
  • 選択した TPU バージョンのトポロジ: TPU の配置と数。

選択した TPU バージョンとトポロジによって、GKE がノードを単一ホストまたはマルチホストのスライスとしてプロビジョニングするかどうかが決まります。単一ホストのスライスでは、各ノードは他の TPU スライスノードとは独立しています。マルチホスト スライスでは、GKE は相互接続された VM を持つノードのグループを TPU スライスに作成します。マルチホスト スライスはアトミックです。つまり、GKE は相互接続されたノードグループ全体を単一のユニットとしてスケーリングします。

使用可能な TPU バージョン、対応するトポロジ、CPU とメモリの容量、スライスの種類については、Autopilot の TPU 構成を選択するをご覧ください。

料金

料金については、Autopilot の料金をご覧ください。

始める前に

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

  • Google Kubernetes Engine API を有効にする。
    • Google Kubernetes Engine API の有効化
  • このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。すでに gcloud CLI をインストールしている場合は、gcloud components update を実行して最新のバージョンを取得する。
  • GKE バージョン 1.29.2-gke.1521000 以降を実行している Autopilot クラスタがあることを確認します。
  • 予約済みの TPU を使用するには、特定の容量予約がすでにあることを確認してください。手順については、予約済みのゾーンリソースの使用をご覧ください。

TPU の割り当てがあることを確認する

TPU スライスノードを作成するには、既存の容量予約を使用する場合を除き、TPU の割り当てが必要です。予約済みの TPU を使用している場合は、このセクションをスキップしてください。

GKE で TPU スライスノードを作成するには、Cloud TPU API の割り当て（tpu.googleapis.com）ではなく、Compute Engine API の割り当て（compute.googleapis.com）が必要です。割り当ての名前は、通常の Autopilot Pod と Spot Pod で異なります。

TPU 用の Compute Engine API の割り当ての上限と現在の使用量を確認するには、次の操作を行います。

  1. Google Cloud コンソールで [割り当て] ページに移動します。

    [割り当て] に移動

  2. [ フィルタ] ボックスで次の操作を行います。

    1. [サービス] プロパティを選択し、「Compute Engine API」と入力して Enter キーを押します。

    2. [タイプ] プロパティを選択し、[Quota] を選択します。

    3. [名前] プロパティを選択して、次のように必要な TPU タイプに基づいて割り当て名を入力します。

      • TPU v5p（tpu-v5p-slice）: TPU v5p チップ
      • TPU v5e（tpu-v5-lite-podslice）: TPU v5 Lite Pod スライスチップ
      • TPU v5e（tpu-v5-lite-device）: TPU v5 Lite デバイスチップ
      • TPU v4（tpu-v4-podslice）: TPU v4 Pod スライスチップ

      Spot Pod の場合は、対応する Preemptible 割り当てを選択します。

    4. [項目（ロケーションなど）] プロパティを選択し、「region:」に続けて、GKE で TPU を作成するリージョンの名前を入力します。たとえば、ゾーン us-west4-a で TPU スライスノードを作成する場合は、「region:us-west4」と入力します。TPU の割り当てはリージョン単位であるため、同じリージョン内のすべてのゾーンで同じ TPU の割り当てが使用されます。

入力したフィルタに一致する割り当てがない場合、プロジェクトには目的のリージョンで指定した割り当てのいずれも付与されていないため、TPU 割り当ての増加をリクエストする必要があります。

TPU アプリケーションを準備する

TPU ワークロードには、次の準備要件があります。

  1. JAX、PyTorch、TensorFlow などのフレームワークは、libtpu 共有ライブラリを使用して TPU VM にアクセスします。libtpu には、XLA コンパイラ、TPU ランタイム ソフトウェア、TPU ドライバが含まれています。PyTorch と JAX の各リリースには、特定の libtpu.so バージョンが必要です。GKE で TPU を使用するには、次のバージョンを使用してください。
    TPU タイプ libtpu.so のバージョン
    TPU v5e
    tpu-v5-lite-podslice
    tpu-v5-lite-device
    TPU v5p
    tpu-v5p-slice
    • 推奨される jax[tpu] バージョン: 0.4.19 以降。
    • 推奨される torchxla[tpuvm] バージョン: 毎晩更新されるバージョンの 2023 年 10 月 23 日付のビルドを使用することをおすすめします。
    TPU v4
    tpu-v4-podslice
  2. TPU リソースをリクエストするコンテナに、次の環境変数を設定します。
    • TPU_WORKER_ID: 各 Pod の一意の整数。この ID は、TPU スライス内の一意のワーカー ID を示します。このフィールドでサポートされる値の範囲は、0 から Pod 数から 1 を引いた値までです。
    • TPU_WORKER_HOSTNAMES: スライス内で相互に通信する必要がある TPU VM ホスト名または IP アドレスのカンマ区切りのリスト。スライス内の TPU VM ごとにホスト名または IP アドレスが必要です。IP アドレスまたはホスト名のリストは、TPU_WORKER_ID によって順序付けされ、ゼロのインデックスが付けられます。

      • GKE は、completionMode: Indexed、subdomain、parallelism > 1 で Job が作成され、google.com/tpu プロパティをリクエストしたときに、変更用 Webhook を使用してこれらの環境変数を自動的に挿入します。また、Service をバックアップする Pod に DNS レコードが追加されるように、ヘッドレス Service が追加されます。

ワークロードの準備が完了したら、TPU を使用する Job を実行できます。

ワークロードで TPU をリクエストする

このセクションでは、Autopilot で TPU をリクエストする Job の作成方法について説明します。TPU を必要とするワークロードでは、次のことを指定する必要があります。

  • TPU のバージョンとトポロジのノードセレクタ
  • ワークロード内のコンテナの TPU チップの数

サポートされている TPU バージョン、トポロジ、スライス内の TPU チップとノードの数については、Autopilot TPU 構成を選択するをご覧ください。

ワークロード内の TPU リクエストに関する考慮事項

TPU を使用できるのは、Pod 内の 1 つのコンテナのみです。コンテナがリクエストする TPU チップの数は、スライス内のノードに接続している TPU チップ数と同じにする必要があります。たとえば、2x4 トポロジで TPU v5e（tpu-v5-lite-podslice）をリクエストする場合は、次のいずれかのリクエストを実行できます。

  • 4 チップ。これにより、それぞれ 4 個の TPU チップを持つ 2 個のマルチホスト ノードが作成されます。
  • 8 チップ。8 個の TPU チップを持つ単一ホストノードが 1 つ作成されます。

費用対効果を最大にするためのベスト プラクティスは、リクエストするスライス内のすべての TPU を常に使用することです。4 個の TPU チップを持つ 2 個のノードのマルチホスト スライスをリクエストする場合は、両方のノードで実行され、スライスの 8 個の TPU チップをすべて消費するワークロードをデプロイする必要があります。

TPU をリクエストするワークロードを作成する

TPU をリクエストするジョブを次の手順で作成します。マルチホスト TPU スライスで実行されるワークロードがある場合は、名前でワークロードを選択するヘッドレス Service も作成する必要があります。このヘッドレス Service は、ワークロード内の Pod を指すように Kubernetes DNS 構成を更新して、マルチホスト スライス内の異なるノード上の Pod が相互に通信できるようにします。

  1. 次のマニフェストを tpu-autopilot.yaml として保存します。

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS

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

    • TPU_TYPE: 使用する TPU タイプ（tpu-v4-podslice など）。GKE でサポートされている値にする必要があります。
    • TOPOLOGY: スライス内の TPU チップの配置（2x2x4 など）。選択した TPU タイプでサポートされているトポロジである必要があります。
    • NUMBER_OF_CHIPS: 使用するコンテナの TPU チップの数。limits と requests の値は同じにする必要があります。

  2. Job をデプロイします。

    kubectl create -f tpu-autopilot.yaml

この Job を作成すると、GKE は自動的に次の処理を行います。

  1. Pod を実行するノードをプロビジョニングします。指定した TPU タイプ、トポロジ、リソース リクエストに応じて、これらのノードは単一ホストスライスまたはマルチホスト スライスのいずれかになります。
  2. Pod に taint を追加してノードに toleration を追加し、他のワークロードが TPU ワークロードと同じノードで実行されないようにします。

例: マルチホスト スライスで TPU チップの合計数を表示する

次のワークロードは、マルチホスト TPU スライス内のノード全体の TPU チップ数を返します。マルチホスト スライスを作成する場合、ワークロードに次のパラメータを設定します。

  • TPU バージョン: TPU v4
  • トポロジ: 2x2x4

このバージョンとトポロジの選択により、マルチホスト スライスが作成されます。

  1. 次のマニフェストを available-chips-multihost.yaml として保存します。 
    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
  2. マニフェストをデプロイします。 
    kubectl create -f available-chips-multihost.yaml

    GKE は、4 つの VM（マルチホスト TPU スライス）を使用して TPU v4 スライスを実行します。スライスには、相互接続された 16 個の TPU チップがあります。

  3. Job によって 4 つの Pod が作成されたことを確認します。 
    kubectl get pods

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

    NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
  4. いずれかの Pod のログを取得します。 
    kubectl logs POD_NAME

    POD_NAME は、作成した Pod の名前に置き換えます。例: tpu-job-podslice-0-5cd8r

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

    TPU cores: 16

例: 単一ノードで TPU チップを表示する

次のワークロードは、特定のノードに関連付けられている TPU チップの数を表示する静的 Pod です。単一ホストノードを作成するには、ワークロードに次のパラメータを設定します。

  • TPU バージョン: TPU v5e
  • トポロジ: 2x4

このバージョンとトポロジを選択すると、単一ホストのスライスが作成されます。

  1. 次のマニフェストを available-chips-singlehost.yaml として保存します。 
    apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8
  2. マニフェストをデプロイします。 
    kubectl create -f available-chips-singlehost.yaml

    GKE は、TPU v5e を使用する 8 つの単一ホスト TPU スライスを含むノードをプロビジョニングします。各 TPU ノードには 8 個の TPU チップがあります（単一ホストの TPU スライス）。

  3. Pod のログを取得します。 
    kubectl logs tpu-job-jax-v5

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

    Total TPU chips: 8

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

ダッシュボード

Google Cloud コンソールの [Kubernetes クラスタ] ページの [オブザーバビリティ] タブに、TPU オブザーバビリティ指標が表示されます。詳細については、GKE のオブザーバビリティ指標をご覧ください。

TPU ダッシュボードは、GKE クラスタでシステム指標が有効になっている場合にのみ表示されます。

ランタイム指標

GKE バージョン 1.27.4-gke.900 以降、JAX バージョン 0.4.14 以降を使用し containerPort: 8431 を指定する TPU ワークロードでは、TPU 使用率の指標を GKE システム指標としてエクスポートします。Cloud Monitoring では、TPU ワークロードのランタイム パフォーマンスをモニタリングするために、次の指標を使用できます。

  • デューティ サイクル: 過去のサンプリング期間（60 秒）において、TensorCore が TPU チップでアクティブに処理していた時間の割合。割合が大きいほど、TPU 使用率が高くなります。
  • メモリ使用量: 割り当てられたアクセラレータ メモリの量（バイト単位）。60 秒ごとにサンプリングされます。
  • メモリの総容量: アクセラレータの総メモリ（バイト単位）。60 秒ごとにサンプリングされます。

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

Kubernetes コンテナ:

  • kubernetes.io/container/accelerator/duty_cycle
  • kubernetes.io/container/accelerator/memory_used
  • kubernetes.io/container/accelerator/memory_total

Kubernetes ノード:

  • kubernetes.io/node/accelerator/duty_cycle
  • kubernetes.io/node/accelerator/memory_used
  • kubernetes.io/node/accelerator/memory_total

ホスト指標

GKE バージョン 1.28.1-gke.1066000 以降では、TPU スライスの VM は TPU 使用率の指標を GKE システム指標としてエクスポートします。Cloud Monitoring では、次の指標を使用して TPU ホストのパフォーマンスをモニタリングできます。

  • TensorCore の使用率: 使用されている TensorCore の現在の割合。TensorCore の値は、マトリックス乗算ユニット（MXU）およびベクトル単位の合計と等しくなります。TensorCore の使用率の値は、過去のサンプル期間（60 秒）に実行された TensorCore オペレーションを、同じ期間にサポートされている TensorCore オペレーションの数で割ったものです。値が大きいほど、使用率が高いことを意味します。
  • メモリ帯域幅の使用率: 現在使用されているアクセラレータ メモリ帯域幅の割合。サンプル期間（60 秒）で使用されたメモリ帯域幅を、同じサンプル期間でサポートされる最大帯域幅で割って計算されます。

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

Kubernetes コンテナ:

  • kubernetes.io/container/accelerator/tensorcore_utilization
  • kubernetes.io/container/accelerator/memory_bandwidth_utilization

Kubernetes ノード:

  • kubernetes.io/container/node/tensorcore_utilization
  • kubernetes.io/container/node/memory_bandwidth_utilization

詳細については、Kubernetes の指標と GKE のシステム指標をご覧ください。

ロギング

TPU VM を含む GKE ノード上で実行されているコンテナによって出力されたログは、GKE ロギング エージェントが収集し、Logging に送信して、Logging に表示します。

Autopilot での TPU ワークロードに関する推奨事項

次の推奨事項により、TPU ワークロードの効率が向上する可能性があります。

  • スケールダウンまたはノードのアップグレードのため、GKE が Pod を終了するまでの最大 7 日間の猶予期間は、拡張ランタイム Pod を使用します。拡張ランタイム Pod でメンテナンスの時間枠と除外を使用すると、ノードの自動アップグレードをさらに遅らせることができます。
  • 容量予約を使用すると、ワークロードが、可用性のキューに配置されることなく、リクエストされた TPU を受け取ることができます。