システム アーキテクチャ

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

概要

Tensor Processing Unit(TPU)は、Google 独自に開発されたアプリケーション固有の集積回路(ASIC)であり、機械学習ワークロードの高速化に使用されます。これらの TPU は、機械学習を主導する Google の豊かな経験を生かして新しく設計されました。

Cloud TPU と TensorFlow を使用すると、Google の TPU アクセラレータ ハードウェア上で独自の機械学習ワークロードを実行できます。Cloud TPU は、パフォーマンスと柔軟性を最大化するよう設計されており、研究者、デベロッパー、そして企業が CPU、GPU、TPU を使用する TensorFlow コンピューティング クラスタをビルドする際に役立ちます。高レベルの Tensorflow API を使用すると、Cloud TPU ハードウェア上で複製されたモデルを簡単に実行できます。

TensorFlow アプリケーションは Google Cloud Platform 上のコンテナ、インスタンス、サービスから TPU ノードにアクセスできます。アプリケーションは VPC ネットワークを介して TPU ノードに接続する必要があります。

TPU のバージョン

TPU のバージョンごとに、TPU デバイスのハードウェア特性が異なります。TPU のバージョンにより、TPU コアのアーキテクチャ、各 TPU コアの高帯域幅メモリ(HBM)の量、各 TPU デバイスのコア間の相互接続、デバイス間通信に使用できるネットワーキング インターフェースが決まります。たとえば、TPU バージョンには次の特性があります。

  • TPU v2:
    • TPU コアごとに 8 GiB の HBM
    • TPU コアごとに 1 つの MXU
    • TPU Pod で最大 512 個の TPU コアと合計 4 TiB のメモリ
  • TPU v3:
    • TPU ごとに 16 GiB の HBM
    • TPU コアごとに 2 つの MXU
    • TPU Pod で最大 2,048 個の TPU コアと合計 32 TiB のメモリ

image

各 TPU コアは、スカラー、ベクトル、マトリックス ユニット(MXU)で構成されています。MXU は、TPU チップで大部分の計算処理を行います。各 MXU により、サイクルごとに 16,000 の乗累算演算を処理できます。MXU の入出力は 32 ビットの浮動小数点値ですが、MXU では、低減された bfloat16 の精度で乗算が行われます。bfloat16 は、IEEE の半精度表現よりも優れたトレーニングとモデル精度を提供する、16 ビットの浮動小数点表現です。

TPU デバイス上の各コアは、ユーザー計算(XLA 演算)を個別に実行できます。高帯域幅の相互接続により、TPU デバイスでチップ間の直接通信が可能になります。TPU Pod 構成では、専用の高速ネットワーク インターフェースが複数の TPU デバイスを相互に接続し、機械学習のワークロードに使用できる TPU コアとより大きな TPU メモリプールを提供します。

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 でも入力バウンドになる可能性があります。インフィードの問題を解決する方法については、パイプライン パフォーマンス ガイドをご覧ください。

TPU の異なるバージョンでベンチマークを実行し、TensorBoard ツールでパフォーマンスをモニタリングすることで、TPU v3 でモデルのパフォーマンスが向上するかどうか判断できます。

TPU の構成

Google データセンターでは、TPU v2 と TPU v3 の両方で、次の構成の TPU デバイスを使用できます。

  • 単一デバイスの TPU。専用の高速ネットワークで相互に接続されていない個別の TPU デバイスです。複数の単一デバイスの TPU タイプを組み合わせて、1 つのワークロードにまとめることはできません。
  • TPU Pod。専用の高速ネットワークで相互に接続されていない個別の TPU デバイスです。

単一デバイスの TPU

Google データセンターの単一デバイス TPU 構成とは、他の TPU デバイスへの専用高速ネットワーク接続がない単一の TPU デバイスによる構成のことを意味します。TPU ノードはこの単一のデバイスにのみ接続します。

image

単一デバイス TPU の場合、チップ間の通信にホスト CPU またはホスト ネットワーキング リソースを必要としないように、デバイス上のチップが相互接続されています。

TPU ノードを作成するときには、TPU タイプを指定します。たとえば、v2-8 または v3-8 を指定して、1 つのデバイスで TPU ノードを構成できます。単一デバイス TPU は TPU Pod 構成の一部ではなく、TPU Pod の一部を占めることもありません。TPU ノードで使用可能な単一デバイス TPU 構成については、TPU タイプをご覧ください。

TPU Pod

Google データセンターの TPU Pod 構成では、専用の高速ネットワーク接続を介して複数の TPU デバイスが相互に接続しています。TPU ノードのホストは、すべての TPU デバイスに機械学習ワークロードを分散します。

image

TPU Pod では、デバイス上の TPU チップが相互に接続されています。このため、チップ間の通信にホスト CPU やホスト ネットワーキング リソースは必要ありません。また、TPU Pod 内の TPU デバイスも専用の高速ネットワークで相互に接続しているので、ホスト CPU やホスト ネットワーキング デバイスは必要ありません。

TPU ノードを作成するときに、TPU Pod 全体またはその一部を占有する TPU タイプを指定します。たとえば、v2-512 TPU タイプは v2 TPU Pod 全体を占有し、v2-128 TPU タイプは v2 TPU Pod の 4 分の 1 を占有します。TPU ノードで使用可能な TPU Pod 構成については、TPU タイプをご覧ください。

image

v2 TPU Pod は、最大で 64 台のデバイスから構成され、合計で 512 TPU v2 コアと 4 TiB の TPU メモリを提供します。

image

v3 TPU Pod は最大で 256 台のデバイスから構成され、合計で 2,048 TPU v3 コアと 32 TiB の TPU メモリを提供します。

TPU のバージョンごとのアーキテクチャ上の違いについては、TPU のバージョンをご覧ください。

Cloud TPU API を使用すると、TPU ノードのサイズに関係なく、TPU ノードの管理を自動化できます。その結果、大規模なコンピューティング クラスタへのスケールアップ、ワークロードの実行、ワークロード完了時のクラスタのスケールダウンを簡単に行うことができます。チップに組み込まれたハードウェア サポートにより、さまざまなディープ ラーニング ワークロードのパフォーマンスを効果的にスケーリングできます。実際に、Cloud TPU ソフトウェア スタックにより、TPU Cloud プログラムを生成、実行、提供する複雑さがなくなります。

ソフトウェア アーキテクチャ

アプリケーションを実行すると、TensorFlow は計算グラフを生成し、gRPC を介して TPU ノードに送信します。TPU ノードに選択した TPU タイプによって、ワークロードに使用できるデバイスの数が決まります。TPU ノードは、計算グラフを時間内にコンパイルし、プログラム バイナリを 1 つ以上の TPU デバイスに送信して実行します。多くの場合、モデルへの入力は Cloud Storage に保存されます。TPU ノードは、1 つ以上の TPU デバイスに入力をストリーミングします。

次のブロック図は、ニューラル ネットワーク モデル、TPU Estimator、TensorFlow クライアント、TensorFlow サーバー、XLA コンパイラで構成された、Cloud TPU ソフトウェア アーキテクチャを示しています。

image

TPU Estimator

TPU Estimator は、Estimator 上にビルドされた高レベル API のセットです。これにより、Cloud TPU のモデルのビルドが簡略化され、TPU のパフォーマンスが最大限に向上します。Cloud TPU を使用するニューラル ネットワーク モデルを記述する場合は、TPU Estimator API を使用する必要があります。

TensorFlow クライアント

TPU Estimator により、プログラムは TensorFlow 演算に変換され、変換された演算は、TensorFlow クライアントによって計算グラフに変換されます。計算グラフは、TensorFlow クライアントにより TensorFlow サーバーに送信されます。

TensorFlow サーバー

TensorFlow サーバーは、Cloud TPU サーバー上で実行されます。サーバーで TensorFlow クライアントからの計算グラフを受信すると、サーバーでは次の操作が行われます。

  1. Cloud Storage から入力を読み込む
  2. グラフを Cloud TPU 上で実行可能な部分と CPU 上で実行する必要のある部分に分割する
  3. Cloud TPU 上で実行するサブグラフに対応する XLA 演算を生成する
  4. XLA コンパイラを呼び出す

XLA コンパイラ

XLA は、TensorFlow サーバーによって生成される High Level Optimizer(HLO)演算を入力として使用するジャスト インタイム コンパイラです。XLA により、オンチップ メモリからハードウェア実行単位へのデータのオーケストレーションやチップ間通信など、Cloud TPU 上で実行されるバイナリコードが生成されます。生成されたバイナリは、Cloud TPU サーバーと Cloud TPU との間の PCIe 接続を使用して Cloud TPU に読み込まれた後、実行のために起動されます。

次のステップ

このページは役立ちましたか?評価をお願いいたします。

フィードバックを送信...