システム アーキテクチャ

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

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

1 つの TPU ボードは、4 つの TPU チップで構成されています。各チップには、2 つの TPU コアが含まれています。TPU コアで使用可能なリソースは、バージョンによって異なります。各 TPU コアには、スカラー、ベクトル、行列乗算ユニット(MXU)があります。TPU チップでは、処理能力の大半が MXU により提供されます。各 MXU では、サイクルごとに 16,000 の積和演算を、bfloat16 に下げた精度で実行できます。bfloat16 は、16 ビットの浮動小数点表現であり、IEEE の半精度表現よりも優れた精度のトレーニングとモデルを実現します。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 および TPU v3 のアーキテクチャの詳細とパフォーマンス特性については、ディープ ニューラル ネットワークをトレーニングするための一領域に特化したスーパーコンピュータをご覧ください。

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 バージョンは v1 と v2 です。

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 Node システム アーキテクチャは、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 に似た形で動作します。

画像

次のステップ