GKE の TPU について

このページでは、Google Kubernetes Engine（GKE）での Cloud TPU について概要を説明します。Tensor Processing Unit（TPU）は、Google が独自に開発したアプリケーション固有の集積回路（ASIC）であり、TensorFlow、PyTorch、JAX などのフレームワークを使用する ML ワークロードを高速化するために使用されます。

GKE で TPU を使用する前に、Cloud TPU の概要で ML アクセラレータの仕組みについて学習することをおすすめします。

このページでは、用語、TPU のメリット、ワークロードのスケジューリングに関する考慮事項など、Google Kubernetes Engine（GKE）での Cloud TPU の基本について説明します。

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

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

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

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

GKE の TPU に関連する用語

このページでは、TPU に関連する次の用語を使用します。

  • TPU タイプ: Cloud TPU タイプ（v5e など）。
  • TPU スライスノード: 相互接続された 1 つ以上の TPU チップを備えた単一の VM で表される Kubernetes ノード。
  • TPU スライス ノードプール: クラスタ内の Kubernetes ノードのグループ。すべて同じ TPU 構成です。
  • TPU トポロジ: TPU スライス内の TPU チップの数と物理的配置。
  • アトミック: GKE は、相互接続されたすべてのノードを単一のユニットとして扱います。スケーリング オペレーション中、GKE はノードセット全体を 0 にスケーリングし、新しいノードを作成します。グループ内のマシンに障害または停止が発生した場合、GKE はノードセット全体を新しいユニットとして再作成します。
  • 不変: 相互接続されたノードのセットに新しいノードを手動で追加することはできません。ただし、目的の TPU トポロジで新しいノードプールを作成し、新しいノードプールでワークロードをスケジュールできます。

TPU スライス ノードプールのタイプ

GKE は、次の 2 種類の TPU ノードプールをサポートしています。

TPU のタイプとトポロジによって、TPU スライスノードをマルチホストまたは単一ホストのいずれとして設定できるかが決まります。次のように設定することをおすすめします。

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

マルチホスト TPU スライス ノードプール

マルチホスト TPU スライスのノードプールは、相互接続された 2 つ以上の TPU VM を含むノードプールです。各 VM には TPU デバイスが接続されています。マルチホスト TPU スライスの TPU は、高速相互接続（ICI）を介して接続されます。マルチホスト TPU スライスのノードプールを作成した後に、ノードを追加することはできません。たとえば、v4-32 ノードプールを作成してから、後でこのノードプールに Kubernetes ノード（TPU VM）を追加することはできません。GKE クラスタに TPU スライスを追加するには、新しいノードプールを作成する必要があります。

マルチホスト TPU スライスのノードプール内の VM は、単一のアトミック単位として扱われます。GKE がスライスに 1 つのノードをデプロイできない場合、TPU スライスノードにノードはデプロイされません。

マルチホスト TPU スライス内のノードを修復する必要がある場合、GKE は TPU スライス内のすべての VM をシャットダウンし、ワークロード内のすべての Kubernetes Pod を強制的に排除します。TPU スライスのすべての VM が稼働したら、新しい TPU スライスの VM で Kubernetes Pod をスケジュールできます。

次の図は、v5litepod-16（v5e）マルチホスト TPU スライスを示しています。この TPU スライスには 4 つの VM があります。TPU スライスの各 VM には、高速相互接続（ICI）で接続された 4 つの TPU v5e チップがあり、各 TPU v5e チップには 1 つの TensorCore があります。

マルチホスト TPU スライスの図

次の図は、1 つの TPU v5litepod-16（v5e）TPU スライス（トポロジ: 4x4）と 1 つの TPU v5litepod-8（v5e）スライス（トポロジ: 2x4）を含む GKE クラスタを示しています。

TPU v5e Pod の図

単一ホスト TPU スライス ノードプール

単一ホストスライスのノードプールは、1 つ以上の独立した TPU VM を含むノードプールです。各 VM には TPU デバイスが接続されています。単一ホストスライスのノードプール内の VM はデータセンター ネットワーク（DCN）を介して通信できますが、VM に接続された TPU は相互接続されません。

次の図は、7 つの v4-8 マシンを含む単一ホストの TPU スライスの例を示しています。

単一ホストスライスのノードプールの図

GKE の TPU の特性

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} の積で、ノードプール内のチップ数が定義されます。たとえば、2x2x2、2x2x4、2x4x4 などのトポロジ フォームを使用して、64 個未満の TPU チップを持つ小さなトポロジを定義できます。64 個を超える TPU チップを持つ大きなトポロジを使用する場合、{A}、{B}、{C} に割り当てる値は、次の条件を満たしている必要があります。

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

マシンタイプ

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

特権モード

特権モードは、securityContext の他のセキュリティ設定の多くをオーバーライドします。TPU にアクセスするには、GKE ノードで実行されているコンテナが次の要件を満たす必要があります。

  • バージョン 1.28 以前では、特権モードを有効にする必要があります。
  • バージョン 1.28 以降では、特権モードは必要ありません。

GKE での TPU の仕組み

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

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

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

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

次のステップ