GKE で TPU を計画する

このページでは、Google Kubernetes Engine（GKE）で Tensor Processing Unit（TPU）の使用量を計画し、TPU の構成ミス、非可用性エラー、割り当て不足による中断リスクを軽減する方法について説明します。

GKE で TPU を使用する前に、GKE での TPU の定義と用語を理解しておいてください。

TPU 構成を計画する

GKE クラスタで TPU を使用するには、構成を計画する必要があります。次の手順をおすすめします。

  1. GKE のオペレーション モードを選択する: GKE Autopilot または Standard クラスタの TPU でワークロードを実行します。

    ベスト プラクティス:

    フルマネージドの Kubernetes エクスペリエンスを実現するには、Autopilot クラスタを使用します。

  2. TPU のバージョンを選択する: TPU のタイプによって、価格性能比、トレーニング スループット、サービス レイテンシなどが異なります。TPU のタイプは、使用可能な CPU とメモリの容量に影響します。

  3. TPU の可用性を確認する: TPU は特定の Google Cloud リージョンで利用できます。GKE ワークロードで TPU タイプを使用するには、そのタイプでサポートされているリージョンにクラスタが存在する必要があります。

  4. TPU トポロジを選択する: TPU スライス内の TPU の物理的な配置。モデルの並列処理の要件に一致するトポロジを選択してください。

このページの参照表を使用して、ノードプールが単一ホスト TPU スライスノードか、マルチホスト TPU スライスノードかを特定します。

GKE のオペレーション モードを選択する

TPU は、クラスタで使用可能な GKE オペレーション モードで使用できます。

  • Autopilot モード（推奨）: 基盤となるインフラストラクチャ（ノード構成、自動スケーリング、自動アップグレード、ベースライン セキュリティ構成、ベースライン ネットワーキング構成など）は GKE によって管理されます。Autopilot では、TPU のタイプとトポロジを選択して Kubernetes マニフェストで指定します。GKE は、TPU を使用してノードをプロビジョニングし、ワークロードのスケジューリングを管理します。
  • Standard モード: 個々のノードの構成を含め、基盤となるインフラストラクチャをユーザーが管理します。

ワークロードに最適な GKE のオペレーション モードを選択するには、GKE のオペレーション モードを選択するをご覧ください。

TPU のバージョンを選択する

TPU スライスの VM の技術特性は次のとおりです。

Autopilot

TPU バージョン マシンタイプ vCPU 数 メモリ（GiB） NUMA ノードの数 TPU スライスノード内の TPU チップの最大数
TPU Trillium（v6e） （プレビュー）
 tpu-v6e-slice 44 ～ 180 176 ～ 1,440 1～2 256
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
TPU v3（単一ホストのみ）
 tpu-v3-device 96 340 2 8
TPU v3
 tpu-v3-slice 48 340 1 256

スタンダード

TPU バージョン マシンタイプ vCPU 数 メモリ（GiB） NUMA ノードの数 プリエンプトされる可能性
TPU Trillium（v6e） （プレビュー）
 ct6e-standard-1t 44 448 2 高い
TPU v6e（プレビュー）
 ct6e-standard-4t 180 720 1 中
TPU v6e（プレビュー）
 ct6e-standard-8t 180 1440 2 低
TPU v5p
 ct5p-hightpu-4t 208 448 2
TPU v5e
 ct5l-hightpu-1t 24 48 1 高
TPU v5e
 ct5l-hightpu-4t 112 192 1 中
TPU v5e
 ct5l-hightpu-8t 224 384 2 低
TPU v5e
 ct5lp-hightpu-1t 24 48 1 高
TPU v5e
 ct5lp-hightpu-4t 112 192 1 中
TPU v5e
 ct5lp-hightpu-8t 224 384 1 低
TPU v4
 ct4p-hightpu-4t 240 407 2
TPU v3（単一ホストのみ）
 ct3-hightpu-4t 96 340 2 8
TPU v3
 ct3p-hightpu-4t 48 340 1 256

モデルに基づいて使用に適したマシンタイプを評価する際は、次の構成を検討してください。

  • ct5l- マシンタイプは、小規模から中規模のモデルの提供に適していますが、大規模モデルには適していません。ct5l- マシンタイプは単一ホストであるため、複数のホスト間の高速相互接続リンクはありません。
  • マルチホスト ct5lp- マシンタイプは、大規模なモデルの提供やトレーニングに適しています。マルチホスト ct5lp- マシンは高速リンクで相互接続されています。

使用する TPU 構成を決定するには、Cloud TPU の料金に関するドキュメントで TPU の仕様と料金を確認してください。

制限事項

使用する TPU を選択する際に、次の制限事項を考慮してください。

  • TPU v6e はプレビュー版で、次のバージョンで使用できます。
    • バージョン 1.31.1-gke.1846000 以降の Standard クラスタ。
    • バージョン 1.31.2-gke.1115000 以降の Autopilot クラスタ。
  • TPU v6e では、ct6e-standard-8t で SMT を 2 に設定することはできません。
  • GKE の費用割り当てと使用状況測定には、予約済みの TPU v4 の使用量や費用に関するデータは含まれません。
  • TPU v5p と v5e は、us-east5 の riptide / イメージ ストリーミングをサポートしていません。
  • TPU v5p 自動スケーリングは、バージョン 1.29.2-gke.1035000 または 1.28.7-gke.1020000 以降のコントロール プレーンを実行する GKE クラスタでサポートされています。
  • 容量予約では、特定の予約を使用します。

GKE で TPU の可用性を検証する

TPU は特定の Google Cloud リージョンで利用できます。GKE クラスタで TPU タイプを使用するには、そのタイプでサポートされているリージョンにクラスタが存在する必要があります。

Autopilot

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

TPU バージョン cloud.google.com/gke-tpu-accelerator 最小 GKE バージョン 可用性 ゾーン
TPU v6e
 tpu-v6e-slice 1.31.2-gke.1384000 プレビュー us-east5-b
europe-west4-a
us-east1-d
asia-northeast1-b
us-south1-a
TPU v5e tpu-v5-lite-device 1.27.2-gke.2100 一般提供 europe-west4-b
us-central1-a
TPU v5e tpu-v5-lite-podslice 1.27.2-gke.2100 一般提供 europe-west4-a
us-central1-a
us-east1-c
us-east5-b
us-west1-c
us-west4-a
us-west4-b
TPU v5p tpu-v5p-slice 1.28.3-gke.1024000 一般提供 us-east1-d
us-east5-a
us-east5-c
TPU v4 tpu-v4-podslice 1.26.1-gke.1500 一般提供 us-central2-b
TPU v3 tpu-v3-slice 1.31.1-gke.1146000 一般提供 us-east1-d
europe-west4-a
TPU v3 tpu-v3-device 1.31.0-gke.1500 一般提供 us-east1-d
europe-west4-a
us-central1-a
us-central1-b
us-central1-f

スタンダード

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

TPU バージョン 次で始まるマシンタイプ 最小 GKE バージョン 可用性 ゾーン
TPU v6e ct6e- 1.31.2-gke.1115000 プレビュー us-east5-b
europe-west4-a
us-east1-d
asia-northeast1-b
us-south1-a
TPU v5e ct5l- 1.27.2-gke.2100 一般提供 europe-west4-b
us-central1-a
TPU v5e ct5lp- 1.27.2-gke.2100 一般提供 europe-west4-a
us-central1-a
us-east1-c
us-east5-b
us-west1-c
us-west4-a
us-west4-b
TPU v5p ct5p- 1.28.3-gke.1024000 一般提供 us-east1-d
us-east5-a
us-east5-c
TPU v4 ct4p- 1.26.1-gke.1500 一般提供 us-central2-b
TPU v3 ct3p- 1.31.1-gke.1146000 一般提供 us-east1-d
europe-west4-a
TPU v3 ct3- 1.31.0-gke.1500 一般提供 us-east1-d
europe-west4-a
us-central1-a
us-central1-b
us-central1-f

TPU を構成する際は、次の注意事項を考慮してください。

  • 特定のゾーン（europe-west4-a、us-east5-b、us-west4-b）では、マシンタイプが ct5lp- で始まり、ct5l- では始まらない単一ホストの TPU v5e ノードプールを作成できます。これらのゾーンでは、トポロジが 2x4 以上の ct5lp-hightpu-4t を使用できます。
  • us-west4 リージョンに単一ホストの TPU v5e を作成するには、ゾーン us-west4-a を選択し、ct5lp- で始まるマシンタイプ（ct5lp-hightpu-1t など）を使用します。
  • 上記の表に記載されている他のリージョンに単一ホストの TPU v5e を作成するには、ct5l- で始まるマシンタイプ（ct5l-hightpu-1t、ct5l-hightpu-4t、ct5l-hightpu-8t など）を使用します。
  • ct5l- で始まるマシンタイプには、ct5lp- で始まるマシンタイプとは異なる割り当てが必要です。

トポロジを選択する

TPU バージョンを決定したら、その TPU タイプでサポートされているトポロジを選択します。TPU タイプに応じて、トポロジは 2 次元または 3 次元になります。モデルの並列処理要件はトポロジの決定に役立ちます。トポロジ内の各サイズの積を計算することで、スライス内の TPU チップの数を特定できます。例:

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

特定のトポロジが単一ホストとマルチホストの両方の TPU スライスノードをサポートしている場合、ホストタイプは、ワークロードがリクエストする TPU チップの数によって決まります。

たとえば、TPU v5e（tpu-v5-lite-podslice）は、単一ホストとマルチホストの両方で 2x4 トポロジをサポートします。設定に応じて次のようになります。

  • ワークロードで 4 個のチップをリクエストすると、4 個の TPU チップを持つマルチホスト ノードが提供されます。
  • ワークロードで 8 個のチップをリクエストすると、8 個の TPU チップを持つ単一ホストノードが提供されます。

次の表を使用して、ユースケースの TPU マシンタイプとトポロジを選択してください。

  • 小規模なモデルのトレーニングまたは推論の場合は、単一ホストの TPU スライス ノードプールで TPU v4 または TPU v5e を使用します。

  • 大規模なモデルのトレーニングや推論を行うには、マルチホストの TPU スライス ノードプールで TPU v4 または TPU v5e を使用します。

Autopilot

TPU バージョン マシンタイプ トポロジ スライス内の TPU チップの数 ノード数 ノードプールのタイプ
TPU v6e（プレビュー） tpu-v6e-slice 1×1 1 1 単一ホスト
2x2 4 4 単一ホスト
2x4 8 8 単一ホスト
4x4 16 4 マルチホスト
4x8 32 8 マルチホスト
8x8 64 16 マルチホスト
8x16 128 32 マルチホスト
16x16 256 64 マルチホスト
TPU v5p tpu-v5p-slice 2x2x1 4 1 単一ホスト
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-podslice2 1×1 1 1 単一ホスト
2x2 4 1
2x4 8 1
2x4 8 2 マルチホスト
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-podslice2 2x2x1 4 1 単一ホスト
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 v3 tpu-v3-slice 4x4 16 2 マルチホスト
4x8 32 4 マルチホスト
8x8 64 8 マルチホスト
8x16 128 16 マルチホスト
16x16 256 32 マルチホスト
TPU v3 tpu-v3-device 2x2 4 1 単一ホスト

  1. トポロジの積を 4 で割った値から計算されます。↩

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

    • 64 個を超えるチップの場合、{A}、{B}、{C} は 4 の倍数にする必要があります。
    • 最大のトポロジは 16x16x24 です。
    • 値は {A}≤{B}≤{C} にする必要があります（例: 8x12x16）。

  2. カスタム トポロジはサポートされていません。

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

Standard

TPU バージョン マシンタイプ トポロジ TPU チップの数 VM 数 ノードプールのタイプ
TPU v6e（プレビュー） ct6e-standard-1t 1×1 1 1 単一ホスト
ct6e-standard-8t 2x4 8 1 単一ホスト
ct6e-standard-4t 2x2 4 1 単一ホスト
2x4 8 2 マルチホスト
4x4 16 4 マルチホスト
4x8 32 8 マルチホスト
8x8 64 16 マルチホスト
8x16 128 32 マルチホスト
16x16 256 64 マルチホスト
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 マルチホスト
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 v3 ct3-hightpu-4t 2x2 4 1 単一ホスト
TPU v3 ct3p-hightpu-4t 4x4 16 4 マルチホスト
4x8 32 8 マルチホスト
8x8 64 16 マルチホスト
8x16 128 32 マルチホスト
16x16 256 64 マルチホスト
16x32 512 128 マルチホスト
32x32 1024 256 マルチホスト

  1. トポロジの積を 4 で割った値から計算されます。↩

高度な構成

以降のセクションでは、高度な TPU 構成のスケジューリングのベスト プラクティスについて説明します。

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 スライス ノードプールの作成方法については、ノードプールを作成するをご覧ください。

Standard クラスタの CPU

GKE は各 TPU スライスを独自のノードに配置するため、このセクションの説明は Autopilot クラスタには当てはまりません。詳細については、Autopilot モードでの TPU の仕組みをご覧ください。

Standard クラスタの場合は、次のスケジューリングのベスト プラクティスを検討してください。

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

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

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

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

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

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

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

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

これらの推奨事項は中断を最小限に抑えるために役立ちますが、中断を完全に防ぐものではありません。たとえば、ハードウェア障害やデフラグメンテーションによるプリエンプションは引き続き発生する可能性があります。同様に、GKE メンテナンスの除外を設定しても、Compute Engine のメンテナンス イベントは実行されます。

ベスト プラクティス:

チェックポイントを頻繁に保存し、トレーニング スクリプトにコードを追加して、再開時に最後のチェックポイントから開始できるようにします。

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

TPU をホストする GKE ノードは、メンテナンス イベントなど、ノード シャットダウンの原因となる可能性のある中断の影響を受けます。コントロール プレーンがバージョン 1.29.1-gke.1425000 以降を実行している GKE クラスタでは、ワークロードを正常に停止するように GKE を構成することで、ワークロードの中断を軽減できます。

AI / ML ワークロードを実行している GKE ノードで発生する可能性のある中断イベントを理解、構成、モニタリングするには、GPU と TPU の GKE ノードの中断を管理するをご覧ください。

TPU 使用率を最大にする

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

ベスト プラクティス:

Kueue を使用して、Job をキューにオーケストレートします。

次のステップ