GKE の TPU について


このページでは、Cloud TPU の概要を説明し、Google Kubernetes Engine(GKE)での Cloud TPU の使用方法に関する情報の参照先を紹介します。Tensor Processing Unit(TPU)は、Google が独自に開発したアプリケーション固有の集積回路(ASIC)であり、TensorFlowPyTorchJAX などのフレームワークを使用する ML ワークロードを高速化するために使用されます。

GKE で TPU を使用する前に、次の学習プログラムを完了することをおすすめします。

  1. Cloud TPU の概要で、ML アクセラレータの仕組みについて学習する。
  2. Cloud TPU システム アーキテクチャで、現在の TPU バージョンの可用性について学習する。

GKE で Cloud TPU を設定する方法については、次のリソースをご覧ください。

GKE で TPU を使用するメリット

GKE は、TPU VM の作成、構成、削除など、TPU VM のライフサイクル管理全体をサポートしています。GKE は、Spot VM予約済み Cloud TPU の使用もサポートしています。GKE で TPU を使用するメリットは次のとおりです。

  • 一貫性のある運用環境: 1 つのプラットフォームで、すべての ML とその他のワークロードに対応できます。
  • 自動アップグレード: GKE はバージョンの更新を自動化し、運用上のオーバーヘッドを削減します。
  • ロード バランシング: GKE は負荷を分散してレイテンシを減らし、信頼性を向上させます。
  • レスポンシブなスケーリング: GKE は、ワークロードのニーズに合わせて TPU リソースを自動的にスケーリングします。
  • リソース管理: Kubernetes ネイティブのジョブ キューイング システムである Kueue により、キューイング、プリエンプション、優先順位付け、公平な共有を使用して組織内の複数のテナントのリソースを管理できます。

GKE の TPU に関連する用語

このドキュメントでは、TPU に関連する次の用語を使用します。

  • TPU タイプ: Cloud TPU タイプ(v5e など)。
  • TPU スライス: 相互接続された TPU チップのセット。
  • TPU トポロジ: TPU スライス内の TPU チップの数と物理配置。
  • 単一ホストの TPU スライスノード: 1 つ以上の独立した TPU ノード。VM に接続されている TPU が高速の相互接続で接続されることはありません。
  • マルチホストの TPU スライスノード: 相互接続された 2 つ以上の TPU ノード。VM に接続されている TPU は、高速の相互接続で接続されます。マルチホスト TPU ノードには次の特性があります。
    • アトミック: GKE は、相互接続されたすべてのノードを単一のユニットとして扱います。スケーリング オペレーション中、GKE はノードセット全体を 0 にスケーリングし、新しいノードを作成します。グループ内のマシンに障害または停止が発生した場合、GKE はノードセット全体を新しいユニットとして再作成します。
    • 不変: 相互接続されたノードのセットに新しいノードを手動で追加することはできません。

GKE での TPU の仕組み

Kubernetes のリソース管理と優先度では、TPU VM は他の VM タイプと同じように扱われます。TPU チップをリクエストするには、リソース名 google.com/tpu を使用します。

    resources:
        requests:
          google.com/tpu: 4
        limits:
          google.com/tpu: 4

GKE で TPU を使用する場合は、次の TPU 特性を考慮する必要があります。

  • TPU VM は最大 8 個の TPU チップにアクセスできます。
  • TPU スライスには、選択した TPU マシンタイプに応じて一定数の TPU チップが含まれます。
  • リクエストする google.com/tpu の数は、TPU ノードで使用可能なチップの合計数と同じにする必要があります。TPU をリクエストする GKE Pod 内のコンテナは、ノード内のすべての TPU チップを使用する必要があります。そうしないと Deployment が失敗します。GKE では TPU リソースを部分的に使用することはできません。たとえば、次のシナリオをご覧ください。
    • マシンタイプ ct5l-hightpu-8t には、8 個の TPU チップを持つ単一の TPU ノードがあります。ノードには 8 個の TPU チップを必要とする 1 つの GKE Pod をデプロイできますが、4 個の TPU チップを必要とする 2 つの GKE Pod をデプロイすることはできません。
    • 2x4 トポロジのマシンタイプ ct5lp-hightpu-4t には、4 個のチップを持つ 2 つの TPU ノード(合計 8 個の TPU チップ)があります。このノードプール内のどのノードにも、8 個の TPU チップを必要とする GKE Pod はデプロイできませんが、4 個の TPU チップを必要とする 2 個の Pod はデプロイできます。
    • トポロジ 4x4 の TPU v5e には、4 つのノードに 16 個のチップがあります。この構成を選択する GKE Autopilot ワークロードは、レプリカごとに 4 つのチップをリクエストする必要があります(最大 4 つのレプリカ)。
  • Standard クラスタでは、1 つの TPU VM に複数の Kubernetes Pod をスケジュールできますが、TPU チップにアクセスできるのは各 Pod の 1 つのコンテナのみです。
  • 各 Standard クラスタには、kube-dns などの kube-system Pod を作成するために、TPU 以外のノードプールが少なくとも 1 つ必要です。
  • デフォルトでは、TPU ノードには TPU 以外の Pod がスケジュールされないように google.com/tpu taint が設定されています。taint は、TPU リソースが完全に使用されることを保証するものではありません。TPU 以外のノードで TPU を使用しないワークロードを実行できるため、TPU ノードのコンピューティングを TPU を使用するコードに解放できます。
  • GKE は、TPU VM で実行されているコンテナから出力されたログを収集します。詳細については、ロギングをご覧ください。
  • ランタイム パフォーマンスなどの TPU 使用率の指標は Cloud Monitoring で利用できます。詳細については、オブザーバビリティと指標をご覧ください。

TPU 構成を計画する

GKE クラスタで TPU を操作するには、次のパラメータを決定する必要があります。

  • TPU タイプ: TPU タイプによって、価格対性能比、トレーニング スループット、サービス レイテンシなどが異なります。TPU が接続されている VM の CPU とメモリの容量が異なります。
  • トポロジ: TPU タイプのチップのレイアウト。各 TPU タイプは、2D または 3D チップのレイアウトと特定の配置をサポートしています。モデルの並列処理の要件に一致するトポロジを選択してください。
  • TPU の相互接続: TPU のタイプとトポロジによって、高速の相互接続で接続された複数のノードで TPU を取得するかどうかが決まります。これにより、単一ホスト TPU スライスノードとマルチホスト TPU スライスノードのどちらを使用するかが決まります。

    • 大規模モデルの場合は、マルチホスト TPU スライスノードを使用します。
    • 小規模モデルの場合は、単一ホストの TPU スライスノードを使用します。

GKE Autopilot モードの TPU 構成を選択する

Autopilot では、TPU のタイプとトポロジを選択し、Kubernetes マニフェストでそれらを指定します。GKE は、TPU を使用してノードをプロビジョニングし、ワークロードのスケジューリングを管理します。

GKE Autopilot での TPU の可用性

TPU は特定の Google Cloud リージョンで利用できます。GKE ワークロードで TPU タイプを使用するには、そのタイプでサポートされているリージョンにクラスタが存在する必要があります。詳細については、Cloud TPU ドキュメントの TPU のリージョンとゾーンをご覧ください。

Autopilot で TPU タイプを選択する

TPU タイプ vCPU 数 メモリ(GiB) NUMA ノードの数 スライス内の TPU チップの最大数
TPU v5p(プレビュー)
tpu-v5p-slice
208 448 2 6,144
TPU v5e
tpu-v5-lite-podslice
24~224 48~384 1 256
TPU v5e(単一ホストのみ)
tpu-v5-lite-device
24~224 48~384 1~2 8
TPU v4
tpu-v4-podslice
240 407 2 4,096

使用するバージョンを決定するには、Cloud TPU のドキュメントでチップの仕様と料金を確認してください。

Autopilot のトポロジを選択する

トポロジは、TPU スライス内のチップの物理的なレイアウトです。TPU タイプに応じて、トポロジは 2 次元または 3 次元になります。TPU タイプを決定したら、その TPU タイプでサポートされているトポロジを選択します。モデルの並列処理要件はトポロジの決定に役立ちます。トポロジの積は、スライス内のチップ数を表します。例:

  • 2x2x2 は 8 チップのマルチホスト TPU v4 スライスです
  • 2x2 は 4 チップの単一ホスト TPU v5e スライスです。

TPU タイプに選択したトポロジによって、単一ホストとマルチホストのどちらの TPU スライスノードを使用するかが決まります。特定のトポロジが単一ホストとマルチホストの両方のスライスをサポートしている場合、ワークロードがリクエストする TPU チップの数によって、取得するホストタイプが決まります。たとえば、TPU v5e(tpu-v5-lite-podslice)は、単一ホストとマルチホストの両方で 2x4 トポロジをサポートします。ワークロードで 4 個のチップをリクエストすると、4 個のチップを持つマルチホスト ノードが取得されます。ワークロードで 8 個のチップをリクエストすると、8 個のチップを持つ単一ホストノードが取得されます。

次の表に、各 TPU タイプでサポートされているトポロジとチップ数、ノード数、ホストタイプを示します。

TPU タイプ トポロジ スライス内の TPU チップ ノード数 ホストタイプ
TPU v5p(プレビュー)
tpu-v5p-slice
2x2x1 4 1 単一ホスト

64 個を超えるチップのカスタム トポロジがサポートされています。次の条件が適用されます。

  • 64 個を超えるチップの場合、{A}{B}{C} は 4 の倍数にする必要があります。
  • 最大のトポロジは 16x16x24 です。
  • 値は {A}{B}{C} にする必要があります(例: 8x12x16)。
2x2x2 8 2 マルチホスト
2x2x4 16 4 マルチホスト
2x4x4 32 8 マルチホスト
4x4x4 64 16 マルチホスト
{A}x{B}x{C} A*B*C (A*B*C/4)1 マルチホスト
TPU v5e
tpu-v5-lite-podslice
1×1 1 1 単一ホスト カスタム トポロジはサポートされていません。
2x2 4 1
2x4 8 1
2x4 2 1 マルチホスト
4x4 16 4
4x8 32 8
8x8 64 16
8x16 128 32
16x16 256 64
TPU v5e(単一ホストのみ)
tpu-v5-lite-device
1×1 1 1 単一ホスト カスタム トポロジはサポートされていません。
2x2 4 1
2x4 8 1
TPU v4
tpu-v4-podslice
2x2x1 4 1 単一ホスト

64 個を超えるチップのカスタム トポロジがサポートされています。次の条件が適用されます。

  • 64 個を超えるチップの場合、{A}{B}{C} は 4 の倍数にする必要があります。
  • 最大のトポロジは 12x16x16 です。
  • 値は {A}{B}{C} にする必要があります(例: 8x12x16)。
2x2x2 8 2 マルチホスト
2x2x4 16 4 マルチホスト
2x4x4 32 8 マルチホスト
4x4x4 64 16 マルチホスト
{A}x{B}x{C} A*B*C (A*B*C/4)1 マルチホスト
  1. トポロジの積を 4 で割った値から計算されます。

TPU のタイプとトポロジを選択したら、ワークロード マニフェストでそれらを指定します。手順については、GKE Autopilot に TPU ワークロードをデプロイするをご覧ください。

GKE Standard モードの TPU 構成を選択する

以下のセクションでは、GKE で TPU ワークロードを計画して設定する際に評価する TPU の特性について説明します。使用可能なバージョン、マシンタイプ、有効なトポロジ、チップ数の詳細については、このドキュメントの TPU 構成のマッピングをご覧ください。

GKE Standard モードでの TPU の可用性

次の表に、マシンタイプとバージョンに応じた TPU の可用性を示します。

TPU バージョン 次で始まるマシンタイプ 最小 GKE バージョン 可用性 ゾーン
TPU v4 ct4p- 1.26.1-gke.1500 一般提供 us-central2-b
TPU v5e ct5l- 1.27.2-gke.2100 一般提供 europe-west4-b
us-central1-a
TPU v5e ct5lp- 1.27.2-gke.2100 一般提供 europe-west4-a1
us-central1-a1
us-east1-c
us-east5-b1
us-west1-c
us-west4-a
us-west4-b1
TPU v5p ct5p- 1.28.3-gke.1024000 プレビュー us-east1-d
us-east5-a
us-east5-c
  1. europe-west4-aus-central1-aus-east5-bus-west4-b のいずれかのゾーンで、マシンタイプが ct5lp- で始まる TPU v5e を作成する場合、単一ホストの TPU v5e ノードプールはサポートされません。これらのゾーンのいずれかで TPU v5e ノードプールを作成する場合は、トポロジが 2x4 以上のマシンタイプ ct5lp-hightpu-4t のみがサポートされます。us-central1 または europe-west4 で単一ホストの TPU v5e を作成するには、それぞれ us-central1-a または europe-west4-b ゾーンを選択し、ct5l- で始まるマシンタイプ(ct5l-hightpu-1tct5l-hightpu-4tct5l-hightpu-8t など)を使用します。us-west4 リージョンに単一ホストの TPU v5e を作成するには、ゾーン us-west4-a を選択し、ct5lp- で始まるマシンタイプ(ct5lp-hightpu-1t など)を使用します。ct5l- で始まるマシンタイプには、ct5lp- で始まるマシンタイプとは異なる割り当てが必要です。

以下のセクションでは、GKE で TPU ワークロードを計画して設定する際に評価する TPU の特性について説明します。使用可能なバージョン、マシンタイプ、有効なトポロジ、チップ数の詳細については、このドキュメントの TPU 構成のマッピングをご覧ください。

マシンタイプ

TPU リソースをサポートするマシンタイプは、TPU のバージョンとノードあたりのチップ数(ct<version>-hightpu-<node-chip-count>t など)を含む命名規則に従います。たとえば、マシンタイプ ct5lp-hightpu-1t は TPU v5e をサポートし、合計で 1 つの TPU チップを使用します。

トポロジ

トポロジでは、TPU スライス内の TPU の物理的な配置を定義します。GKE は、TPU のバージョンに応じて、2 次元または 3 次元のトポロジで TPU スライスをプロビジョニングします。トポロジは、各ディメンションの TPU チップの数として指定します。

  • マルチホストの TPU スライスのノードプールにスケジュールされる TPU v4 と v5p の場合、トポロジを 3 タプル({A}x{B}x{C})で定義します(例: 4x4x4)。{A}x{B}x{C} の積で、ノードプール内のチップ数が定義されます。たとえば、2x2x22x2x42x4x4 などのトポロジ フォームを使用して、64 チップ未満の小さなトポロジを定義できます。64 チップを超えるトポロジを使用する場合、{A}、{B}、{C} に割り当てる値は、次の条件を満たしている必要があります。

    • {A}、{B}、{C} は 4 の倍数です。
    • v4 でサポートされている最大のトポロジは 12x16x16 で、v5p は 16x16x24 です。
    • 割り当てられた値は、A ≤ B ≤ C の関係を維持する必要があります。たとえば、4x4x8 や、8x8x8 です。

TPU 構成のマッピング

次の表を使用して、ユースケースで使用する TPU のマシンタイプとトポロジを定義します。

  • 小規模なモデルのトレーニングまたは推論の場合は、単一ホストの TPU スライス ノードプールで TPU v4 または TPU v5e を使用します。
  • 大規模なモデルのトレーニングや推論を行うには、マルチホストの TPU スライス ノードプールで TPU v4 または TPU v5e を使用します。
TPU バージョン マシンタイプ トポロジ TPU チップの数 VM 数 ノードプールのタイプ
TPU v4 ct4p-hightpu-4t 2x2x1 4 1 単一ホスト
2x2x2 8 2 マルチホスト
2x2x4 16 4 マルチホスト
2x4x4 32 8 マルチホスト
{A}x{B}x{C} A*B*C (A*B*C/4)1 マルチホスト
TPU v5p ct5p-hightpu-4t 2x2x1 4 1 単一ホスト
2x2x2 8 2 マルチホスト
2x2x4 16 4 マルチホスト
2x4x4 32 8 マルチホスト
{A}x{B}x{C} A*B*C (A*B*C/4)1 マルチホスト
TPU v5e ct5l-hightpu-1t 1×1 1 1 単一ホスト
ct5l-hightpu-4t 2x2 4 1 単一ホスト
ct5l-hightpu-8t 2x4 8 1 単一ホスト
ct5lp-hightpu-1t 1×1 1 1 単一ホスト
ct5lp-hightpu-4t 2x2 4 1 単一ホスト
ct5lp-hightpu-8t 2x4 8 1 単一ホスト
ct5lp-hightpu-4t 2x4 8 2 マルチホスト
4x4 16 4 マルチホスト
4x8 32 8 マルチホスト
8x8 64 16 マルチホスト
8x16 128 32 マルチホスト
16x16 256 64 マルチホスト
  1. トポロジの積を 4 で割った値から計算されます。

TPU v5e の特性

TPU v5e マシンの技術特性は次のとおりです。

マシンタイプ vCPU 数 メモリ(GB) NUMA ノードの数 プリエンプトされる可能性
ct5l-hightpu-1t 24 48 1
ct5l-hightpu-4t 112 192 1
ct5l-hightpu-8t 224 384 2
ct5lp-hightpu-1t 24 48 1
ct5lp-hightpu-4t 112 192 1
ct5lp-hightpu-8t 224 384 1

TPU v4 と v5p の特性

TPU v4p マシンと v5p マシンには、次の技術的特性があります。

マシンタイプ vCPU 数 メモリ(GB) NUMA ノードの数
ct4p-hightpu-4t 240 407 2
ct5p-hightpu-4t 208 448 2

TPU の予約

TPU の予約は、コミットメントを購入した場合に利用できます。任意の TPU 予約を GKE で使用できます。

TPU ノードプールを作成する場合は、--reservation フラグと --reservation-affinity=specific フラグを使用して、予約済み TPU インスタンスを使用します。

GKE での TPU の自動スケーリング

GKE は、ML ワークロードを高速化するために Tensor Processing Unit(TPU)をサポートしています。単一ホストの TPU スライス ノードプールとマルチホストの TPU スライス ノードプールはどちらも、自動スケーリングと自動プロビジョニングをサポートしています。

GKE クラスタで --enable-autoprovisioning フラグを指定すると、GKE は、TPU のバージョンとトポロジが保留中のワークロードの要件を満たしている単一ホストまたはマルチホストの TPU スライス ノードプールを作成または削除します。

--enable-autoscaling を使用すると、GKE はタイプに基づいてノードプールを次のようにスケーリングします。

  • 単一ホストの TPU スライス ノードプール: GKE は、既存のノードプールで TPU ノードを追加または削除します。ノードプールには、0 からノードプールの最大サイズまでの任意の数の TPU ノードが含まれます。最大サイズは、--max-nodes フラグと --total-max-nodes フラグによって決まります。ノードプールがスケーリングされると、ノードプール内のすべての TPU ノードのマシンタイプとトポロジは同じになります。単一ホストの TPU スライス ノードプールを作成する方法については、ノードプールを作成するをご覧ください。

  • マルチホスト TPU スライス ノードプール: GKE は、ノードプールを 0 から TPU トポロジを満たすために必要なノード数までアトミックにスケールアップします。たとえば、マシンタイプが ct5lp-hightpu-4t でトポロジが 16x16 の TPU ノードプールの場合、ノードプールには 64 個のノードが含まれます。GKE オートスケーラーは、このノードプールのノード数が 0 または 64 になるように調整します。スケールダウンすると、GKE はスケジュールされたすべての Pod を強制排除し、ノードプール全体が 0 になるようにドレインします。マルチホスト TPU スライス ノードプールの作成方法については、ノードプールを作成するをご覧ください。

制限事項

  • 容量予約の場合は、特定の予約を使用する必要があります。
  • GKE の費用割り当てと使用状況測定には、予約済みの TPU v4 の使用量や費用に関するデータは含まれません。
  • TPU v5p と v5e は、us-east5 の riptide / イメージ ストリーミングをサポートしていません。
  • TPU v5p 自動スケーリングは、1.29.2-gke.1035000 以降または 1.28.7-gke.1020000 以降のコントロール プレーンを実行する GKE クラスタでサポートされています。

ワークロードのスケジューリングに関する考慮事項

TPU には、Kubernetes で特別なワークロードのスケジューリングと管理が必要になる独自の特性があります。以降のセクションでは、スケジューリングのベスト プラクティスについて説明します。

CPU

GKE は各 TPU Pod を独自のノードに配置するため、このセクションの説明は Autopilot クラスタには当てはまりません。

TPU VM のオンボード CPU でワークロードをスケジューリングするには、GKE Pod が google.com/tpu taint を許容できることを確認します。ワークロードを特定のノードにデプロイする場合は、ノードセレクタを使用します。

Kubernetes のリソース管理と優先度では、TPU VM は他の VM タイプと同じように扱われます。TPU のスケジューリングを必要とする Pod を同じノード上の他の Pod よりも優先するには、その TPU Pod の最大 CPU またはメモリをリクエストします。優先度の低い Pod では、次のことを行います。

  1. CPU リクエストとメモリ リクエストを少なく設定して、ノードに TPU ワークロード用に十分なリソースを確保します。詳細については、Kubernetes がリソースのリクエストと上限を適用する方法をご覧ください。
  2. CPU の上限を設定しない(無制限)と、Pod がバーストして、未使用のサイクルがすべて使用される可能性があります。
  3. ノードの安定性を維持しながら、Pod が未使用メモリを使用できるように、メモリの上限を高く設定します。

Kubernetes Pod が TPU をリクエストしていても、CPU とメモリをリクエストしない場合、この Kubernetes Pod はベスト エフォート Pod となりますが、CPU とメモリが必要と見なされる保証はありません。このような状態が保証されるのは、CPU とメモリを明示的にリクエストする Pod のみです。詳しくは、Pod とコンテナのリソース管理をご覧ください。

詳細については、Kubernetes のベスト プラクティス: リソースのリクエストと上限をご覧ください。

ワークロードの中断を減らす

ML モデルのトレーニングに TPU を使用していて、ワークロードが中断された場合、最後のチェックポイント以降に実行された作業はすべて失われます。ワークロードが中断される可能性を低くするには、次の操作を行います。

  • この Job に他のすべての Job よりも高い優先度を設定します。リソースが不足している場合、GKE スケジューラは優先度の低い Job をプリエンプトして、優先度の高い Job をスケジュールします。また、優先度の高いワークロードは、クラスタで使用可能な合計リソースの範囲内であれば、必要なすべてのリソースを受け取ることもできます。詳細については、Pod の優先度とプリエンプションをご覧ください。
  • メンテナンスの除外を構成する: メンテナンスの除外は、自動メンテナンスの実行が禁止される定期的な時間枠です。詳細については、メンテナンスの除外をご覧ください。
  • Autopilot で拡張ランタイム Pod を使用する: スケールダウンまたはノードのアップグレードのため、GKE が Pod を終了するまでの最大 7 日間の猶予期間は、拡張ランタイム Pod を使用します。

ノード メンテナンスによる中断に対処する

TPU VM をホストするノードを含むすべての GKE ノードは、メンテナンス イベントや、ノードのシャットダウンの原因となる中断の影響を受けます。コントロール プレーンがバージョン 1.29.1-gke.1425000 以降を実行している GKE クラスタで、実行中のワークロードに対する影響を軽減できます。GKE は、強制終了の 5 分前に SIGTERM シグナルをノードに送信して、シャットダウンの直前であることをノードに通知します。ワークロードで Orbax を使用して MaxText、Pax、JAX などの ML フレームワークを利用している場合、ワークロードは SIGTERM シグナルをキャプチャして、チェックポイント プロセスを開始できます。

最大通知時間内に ML ワークロードが正常に終了するように GKE を構成できます。Pod マニフェストで、spec.terminationGracePeriodSeconds フィールドを 300 秒(5 分)に設定します。GKE は、これらの Pod を正常に終了し、トレーニング状態の保存など、定義した終了アクションを実行するために最善を尽くします。GKE では、PodDisruptionBudget または terminationGracePeriodSeconds の設定が最大 5 分間適用されます。

spec.terminationGracePeriodSeconds フィールドは、プリエンプティブでない VM で発生するメンテナンス イベントと断片化イベントによる中断のみを処理します。GKE は、ハードウェア障害などの偶発的な障害を処理しません。

詳細については、TPU ノードの正常終了を構成するをご覧ください。

TPU 使用率を最大にする

TPU を最大限に活用するには、異なる優先度の Job を組み合わせてスケジュールし、TPU の稼働時間が最大になるようにキューに入れます。Job レベルのスケジューリングとプリエンプションが必要な場合は、Job をキューにオーケストレートする Kubernetes のアドオンを使用する必要があります。その場合は Kueue の使用をおすすめします。

次のステップ