GKE の TPU について


このページでは、Google Kubernetes Engine(GKE)での Cloud TPU について概要を説明します。Tensor Processing Unit(TPU)は、Google が独自に開発したアプリケーション固有の集積回路(ASIC)であり、TensorFlowPyTorchJAX などのフレームワークを使用する 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 により、キューイング、プリエンプション、優先順位付け、公平な共有を使用して組織内の複数のテナントのリソースを管理できます。

TPU Trillium(v6e)を使用するメリット

TPU Trillium(v6e)は、Cloud TPU の最新の生成 AI アクセラレータです。すべての技術インターフェース(API、ログ)と GKE ドキュメント全体で、TPU Trillium(v6e)は v6e と呼ばれ、Google の第 6 世代の TPU を表します。

  • TPU v6e では、TPU v5e と比較してチップあたりのコンピューティング パフォーマンスが向上しています。
  • TPU v6e では、高帯域幅メモリ(HBM)の容量と帯域幅が向上し、チップ間相互接続(ICI)の帯域幅も TPU v5e よりも向上しています。
  • TPU v6e には、第 3 世代の SparseCore が搭載されています。これは、高度なランキングやレコメンデーションのワークロードでよくある極めて大規模なエンベディングを処理するための、特別なアクセラレータです。
  • TPU v6e のエネルギー効率は TPU v5e より 67% 以上向上しています。
  • TPU v6e は、高帯域幅かつ低レイテンシの単一の TPU スライスで最大 256 個の TPU までスケールアップできます。

TPU Trillium のメリットの詳細については、TPU v6e の発表に関するブログ投稿をご覧ください。TPU の設定を開始するには、GKE で TPU を計画するをご覧ください。

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} の積で、ノードプール内のチップ数が定義されます。たとえば、2x2x22x2x42x4x4 などのトポロジ フォームを使用して、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 で利用できます。詳細については、オブザーバビリティと指標をご覧ください。

次のステップ