システム アーキテクチャ

このドキュメントでは、Cloud TPU システムのハードウェア コンポーネントとソフトウェア コンポーネントのアーキテクチャについて説明します。

Tensor Processing Unit(TPU)は、Google 独自に開発された特定用途向け集積回路(ASIC)であり、機械学習(ML)ワークロードの高速化に使用されます。TPU は、機械学習を主導する Google の豊かな経験を生かして新しく設計されました。

1 つの TPU ボードには 4 つの TPU チップが含まれています。各チップには 2 つの TPU コアが含まれています。TPU コアで使用可能なリソースはバージョンによって異なります。各 TPU コアには、スカラー、ベクトル、行列乗算ユニット(MXU)があります。MXU は、TPU チップで大部分の計算処理を行います。各 MXU により、サイクルごとに 16,000 の乗累算演算を bfloat16 の精度で実行できます。bfloat16 は、IEEE の半精度表現よりも優れた精度のトレーニングとモデルを提供する、16 ビットの浮動小数点表現です。TPU ボード上の各コアは、ユーザー計算を個別に実行できます。高帯域幅の相互接続により、チップ間で直接通信できます。技術的な詳細については、ディープ ニューラル ネットワークをトレーニングするためのドメイン固有のスーパーコンピュータをご覧ください。

TPU は、TPU Pod にスケールアウトされるように設計されています。TPU Pod は、最大 2,048 個の TPU コアを備えたスーパーコンピュータで、処理負荷を複数の TPU ボードに分散できます。TPU Pod 構成では、専用の高速ネットワーク インターフェースが複数の TPU デバイスを相互に接続し、機械学習ワークロードに使用できる TPU コアとより大きな TPU メモリプールを提供します。

各 TPU ボードは、TPU にデータを読み込んで前処理を行う用途向けなど、高性能の CPU ベースのホストマシンに接続されます。

Cloud TPU

Cloud TPU は、Google Cloud Platform(GCP)を介して TPU にアクセスできるサービスです。Cloud TPU を使用して、Google の TPU アクセラレータ ハードウェアで機械学習ワークロードを実行できます。Cloud TPU は、パフォーマンスと柔軟性を最大化するように設計され、研究者、デベロッパー、企業が ML ワークロードをトレーニングするのに役立ちます。

TPU の管理は、Cloud TPU API を使用することで自動化できます。その結果、大規模なコンピューティング クラスタへのスケールアップ、ワークロードの実行、ワークロード完了時のクラスタのスケールダウンが容易になります。チップに組み込まれたハードウェアのサポートにより、幅広いディープ ラーニング ワークロードのパフォーマンスを効果的に線形スケーリングできます。

TPU の構成

Cloud TPU では、さまざまな TPU 構成にアクセスできます。

  • TPU Pod
  • TPU スライス
  • TPU ボード

TPU Pod

TPU Pod では、TPU Pod 内のすべての TPU チップが高速相互接続を介して直接接続され、これにより CPU ホストを経由する通信遅延がなくなります。チップは 2 次元トーラスの形に接続され、各チップは 4 つの隣接するチップと直接通信します。このアーキテクチャは、all-reduce のような ML ワークロードにおける一般的な通信パターンで特に高いパフォーマンスを発揮します。次の表に、TPU v2 と v3 のアーキテクチャの詳細とパフォーマンス特性を示します。詳細については、ディープ ニューラル ネットワークをトレーニングするためのドメイン固有のスーパーコンピュータをご覧ください。

機能 TPU v2 TPU v3
ネットワーク リンク x Gbits / 秒 / チップ 4 x 496 4 x 656
全二分帯域幅テラビット / 全 Pod 15.9 42

TPU スライス

TPU スライスは、TPU Pod の一部分です。Pod 全体のリソースは必要ない場合に、Pod の一部を使用できます。さまざまなスライス構成が使用できます。詳細については、Cloud TPU の Cloud TPU 料金のセクションをご覧ください。

シングル TPU ボード

シングルボード TPU 構成は、4 個の TPU チップ(8 個の TPU コア)を備えたスタンドアロン ボードで、他の TPU ボードへのネットワーク接続はありません。シングルボード TPU は TPU Pod の構成の一部ではなく、TPU Pod の一部分を占有することもありません。利用可能なシングルボード TPU 構成については、TPU タイプのページをご覧ください。

TPU のバージョン

TPU のバージョンにより、各 TPU コアのアーキテクチャ、各 TPU コアの高帯域幅メモリ(HBM)の量、各 TPU ボードのコア間の相互接続、デバイス間通信に使用できるネットワーキング インターフェースが決まります。利用可能な TPU バージョンには、v2 と v3 があります。

TPU v2

TPU v2 ボードには、それぞれ 2 つのコアを内蔵した 4 つの TPU チップが搭載されています。8 GiB の HBM が TPU コアのそれぞれにあり、1 つのコアには 1 つの MXU があります。TPU v2 Pod は、最大 512 個の TPU コアと 4 TiB のメモリを備えます。

画像

TPU v3

TPU v3 ボードには、それぞれ 2 つのコアを内蔵した 4 つの TPU チップが搭載されています。16 GiB の HBM が TPU コアのそれぞれにあり、1 つのコアには 2 つの MXU があります。TPU v3 Pod は、最大 2,048 個の TPU コアと 32 TiB のメモリを備えます。

画像

TPU v2 に対する TPU v3 のパフォーマンス上の利点

TPU v3 構成では、コアあたりの FLOPS とメモリ容量が増えているため、以下の方法でモデルのパフォーマンスを改善できます。

  • 計算依存型モデルの場合、TPU v3 構成を使用するコアごとのパフォーマンスが大幅に向上します。メモリ依存型モデルの場合、TPU v2 構成で TPU v3 構成と同等のパフォーマンスの向上は達成できません。

  • データが TPU v2 構成のメモリに収まらない場合、TPU v3 を使用することで、パフォーマンスが向上し、中間値の再計算(再実体化)も減少します。

  • TPU v3 構成では、TPU v2 構成に収まらないバッチサイズの新しいモデルを実行できます。たとえば、TPU v3 を使用すると、さらに深い ResNets を扱え、RetinaNet でより大きなイメージを処理できます。

TPU v2 で入力依存型(「インフィード」)に近いモデルは、トレーニング ステップで入力を待機するため、Cloud TPU v3 でも入力依存型になる可能性があります。インフィードの問題を解決する方法については、パイプライン パフォーマンス ガイドをご覧ください。

Cloud TPU VM アーキテクチャ

各 TPU ボードは、ホストマシン(TPU ホスト)に物理的に接続されます。

画像

TPU Pod では、TPU ボードごとに TPU ホストがあります。

画像

TPU ホスト(および TPU ボード)の操作方法は、使用する TPU VM アーキテクチャ(TPU ノードまたは TPU VM)によって異なります。

TPU ノード

TPU ノードは、TPU の本来の扱い方です。これには、gRPC を介して TPU ホストと通信する追加のユーザー VM が必要です。TPU ホストに直接アクセスすることはできません。

画像

TPU VM

TPU VM を使用する場合は、TPU ホストで実行されている Google Compute Engine VM に直接 SSH 接続します。マシンへの root アクセス権を取得すれば、必要なコードを実行できます。TPU コンパイラとランタイムから、デバッグログとエラー メッセージに直接アクセスできます。TPU VM は、TPU ノードでは不可能な新しいユースケースをサポートしています。たとえば、入力パイプラインでカスタム演算を実行できます。また、ローカル ストレージを使用することもできます。

画像

ユーザー VM がないため、コードと TPU VM の間にネットワーク、Virtual Private Cloud、ファイアウォールが不要で、入力パイプラインのパフォーマンスが向上します。さらに、TPU VM では、ユーザー VM に対する費用がかからないため、より安価です。

サポートされているフレームワーク

JAX、PyTorch、TensorFlow などのフレームワークでは、すべての TPU VM で提供される libtpu という共有ライブラリを介して TPU にアクセスします。このライブラリには、TPU プログラムのコンパイルに使用される XLA コンパイラ、コンパイルされたプログラムの実行に使用される TPU ランタイム、TPU への低レベルアクセスにランタイムで使用される TPU ドライバが含まれます。

画像

TensorFlow

TPU VM

TPU VM を使用すると、ユーザー VM 上で実行する Python コードではなく、TPU ホストで直接実行できます。

画像

TensorFlow と Cloud TPU の詳細については、Cloud TPU での TensorFlow モデルの実行をご覧ください。

TPU ノード

Cloud TPU ノードシステムのアーキテクチャは、もともと TensorFlow とその分散プログラミング モデル用に構築されています。TPU ホストはユーザーがアクセスできず、TensorFlow サーバーのヘッドレス コピーのみを実行します。Python や TensorFlow グラフとして表されていないユーザーコードは実行されません。ユーザーコードは、ネットワークを介して TPU ホストと通信する別のリモート VM で実行されます。

画像

PyTorch

TPU VM

TPU VM では、PyTorch コードは TPU ホスト上で直接実行されます。

画像

PyTorch と Cloud TPU の詳細については、Cloud TPU での PyTorch モデルの実行をご覧ください。

TPU ノード

PyTorch は、XRT と呼ばれるライブラリを使用して Cloud TPU ノード アーキテクチャ上で実行され、XLA グラフとランタイム命令を TensorFlow gRPC 接続を介して送信し、TensorFlow サーバーで実行できます。各 TPU ホストには、ユーザー VM が必要です。

画像

JAX

TPU VM

TPU VM では、コードを TPU ホストで直接実行できるため、ユーザー VM は必要ありません。

画像

Cloud TPU で JAX を実行する詳細については、JAX クイックスタートをご覧ください。

TPU ノード

Cloud TPU ノード上の JAX は、ホスト VM ごとに個別のユーザー VM が必要という点で PyTorch に似た形で動作します。

画像

次のステップ