Cloud Tensor Processing Unit（TPU）

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

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

TPU の利点

Cloud TPU リソースにより、機械学習アプリケーションで頻繁に使用される線形代数計算のパフォーマンスが向上します。TPU では、大規模で複雑なニューラル ネットワーク モデルをトレーニングする際に、一定の精度に到達するまでの時間が最小になります。他のハードウェア プラットフォームで以前はトレーニングに数週間かかっていたモデルが、TPU では数時間で収束する可能性があります。

利用状況

Cloud TPU のリージョンとゾーンには、Cloud TPU が利用可能なロケーションが示されています。

Cloud TPU プログラミング モデル

Cloud TPU では、密ベクトルおよび行列計算が非常に高速に実行されます。Cloud TPU とホストメモリの間のデータ転送は計算速度に比べて遅く、PCIe バスの速度は Cloud TPU Interconnect とオンチップの高帯域幅メモリ（HBM）のどちらよりもはるかに低速です。モデルを部分的にコンパイルすることにより、実行場所がホストとデバイス間を行き来するため、多くの場合、TPU はアイドル状態になり、データが PCIe バス経由で受信されるまで待機します。この状況を改善するために、Cloud TPU のプログラミング モデルは、トレーニングの多くを TPU で実行するように設計されています（理想的にはトレーニング ループ全体）。

TPU プログラミング モデルの主な特徴は次のとおりです。

  • すべてのモデル パラメータは、オンチップの高帯域幅メモリに保持されます。
  • Cloud TPU で計算を開始するコストは、ループで多くのトレーニング ステップを実行することによって平均化されます。
  • 入力トレーニング データは、Cloud TPU の「インフィード」キューにストリーミングされます。Cloud TPU 上で実行されるプログラムによって、各トレーニング ステップ中にこれらのキューからバッチが取得されます。
  • ホストマシン（Cloud TPU デバイスに接続された CPU）上で稼働する TensorFlow サーバーは、データを取得して事前に処理してから、Cloud TPU ハードウェアへの「インフィード」を行います。
  • データの並列処理: Cloud TPU 上のコアでは、それぞれの HBM に存在する同じプログラムが同期的に実行されます。各ニューラル ネットワーク ステップの終了時にすべてのコアにわたってリダクション演算が実施されます。

TPU の用途

Cloud TPU は特定のワークロード用に最適化されています。状況によっては、Compute Engine インスタンスで GPU または CPU を使用して機械学習ワークロードを実行する場合があります。一般に、次のガイドラインに基づいてワークロードに最適なハードウェアを決定できます。

  • CPU

    • 最大限の柔軟性を必要とする迅速なプロトタイピング
    • トレーニングに時間がかからない単純なモデル
    • 実際のバッチサイズが小さい小規模なモデル
    • C++ で記述されたカスタム TensorFlow 演算が多くを占めるモデル
    • ホストシステムの使用可能な I/O またはネットワーク帯域幅によって制限が課せられるモデル

  • GPU

    • ソースが存在しないモデルまたはソースを変更するのが煩雑すぎるモデル
    • CPU 上で少なくとも部分的に実行しなければならない多数のカスタム TensorFlow 演算を使用するモデル
    • Cloud TPU で利用できない TensorFlow 演算を使用するモデル（利用可能な TensorFlow 演算のリストをご覧ください）
    • 実際のバッチサイズが大きい中〜大規模なモデル

  • TPU

    • 行列計算が多くを占めるモデル
    • メインのトレーニング ループ内にカスタム TensorFlow 演算がないモデル
    • トレーニングに数週間または数か月かかるモデル
    • 実際のバッチサイズが非常に大きい非常に大規模なモデル

Cloud TPU は、以下のワークロードには適していません。

  • 頻繁な分岐を必要とするか、要素ごとの代数が多くを占める線形代数プログラム。TPU は高速かつ大規模な行列乗算を行うよう最適化されているため、行列乗算が多くを占めないワークロードは、他のプラットフォームと比較して TPU で適切に実行される見込みがありません。
  • メモリにあまりアクセスしないワークロードは、TPU で使用できない可能性があります。
  • 高精度の演算を必要とするワークロード。たとえば、倍精度演算は TPU には適していません。
  • C++ で記述されたカスタム TensorFlow 演算を含むニューラル ネットワーク ワークロード。具体的には、メインのトレーニング ループの本体にあるカスタム演算は、TPU には適していません。

ニューラル ネットワーク ワークロードは、トレーニング ループ全体の複数の反復を TPU で実行できる必要があります。これは TPU 自体の基本要件ではありませんが、これは効率のために必要とされる TPU ソフトウェア エコシステムの現在の制約の 1 つです。

従来のトレーニングとの相違点

一般的な TensorFlow トレーニング グラフは、以下のようなさまざまな機能を提供する複数の重なり合うサブグラフで構成されます。

  • トレーニング データを読み込む I/O オペレーション。
  • 入力前処理ステージ（多くの場合キューを介して接続されます）。
  • モデル変数。
  • それらの変数用の初期化コード。
  • モデル自体。
  • 損失関数。
  • 勾配コード（通常は自動的に生成されます）。
  • トレーニングをモニタリングするためのサマリー オペレーション。
  • チェックポインティング用の保存 / 復元オペレーション。

Cloud TPU では、TensorFlow プログラムは XLA ジャストインタイム コンパイラによってコンパイルされます。Cloud TPU でのトレーニングで、ハードウェア上でコンパイルおよび実行できるコードは、モデルの密な部分、損失および勾配サブグラフに対応するコードだけです。TensorFlow プログラムの他のすべての部分は、ホストマシン（Cloud TPU サーバー）上で、通常の分散 TensorFlow セッションの一部として実行されます。これは、通常は、トレーニング データ、前処理コード（たとえば、圧縮イメージのデコード、ランダムなサンプリング / 切り抜き、トレーニング ミニバッチの構築）、グラフのすべてのメンテナンス部分（チェックポイントの保存 / 復元など）を読み取る I/O オペレーションで構成されます。

モデル開発のベスト プラクティス

単一の Cloud TPU チップには 2 つのコアがあり、そのそれぞれに、密な行列乗算と畳み込みが多くを占めるプログラムを高速化するように設計された複数のマトリックス ユニット（MXU）が含まれています（システム アーキテクチャをご覧ください）。実行時間のかなりの部分を行列乗算に費やすプログラムは、一般に Cloud TPU に適しています。計算の多くを非行列演算（加算、再形成、連結など）が占めるプログラムは、MXU 使用率が高くならない可能性があります。次に、Cloud TPU に適したモデルを選択して作成する際に役立つガイドラインを示します。

1 つの Cloud TPU 端末は 4 つのチップで構成され、それぞれのチップに 2 つの TPU コアがあります。したがって、Cloud TPU を効率的に利用するには、プログラムで 8 つのコアそれぞれを使用する必要があります。

レイアウト

XLA コンパイラでは、マトリックス ユニット（MXU）で効率的に計算を実行するために、行列乗算をより小さなブロックにタイリングするなどのコード変換が行われます。MXU ハードウェアの構造、128x128 のシストリック アレイ、8 の倍数の次元を基本とする TPU のメモリ サブシステムの設計が、タイリングの効率化のために XLA コンパイラによって使用されます。その結果、特定のレイアウトはタイリングに役立ちますが、他のレイアウトではタイリングの前に再形成が必要になります。再形成オペレーションは、多くの場合、Cloud TPU でメモリバウンドです。

形状

XLA コンパイラによって、最初のバッチに合わせて TensorFlow グラフがコンパイルされます。後続のバッチに異なる形状がある場合、モデルは機能しません（形状が変化するたびにグラフを再コンパイルするのでは遅すぎます）。したがって、実行時に変化する動的形状を持つテンソルがあるモデルは、TPU にはあまり適していません。

パディング

高パフォーマンスの Cloud TPU プログラムは、密な計算を 128x128 のチャンクに簡単にタイリングできるプログラムです。MXU 全体が行列計算で占有されない場合、コンパイラによってテンソルはゼロでパディングされます。パディングには次の 2 つの欠点があります。

  • ゼロでパディングされたテンソルでは、TPU コアの利用が不十分になります。
  • パディングによって、テンソルに必要なオンチップ メモリ ストレージ量が増加し、極端な場合にはメモリ不足エラーにつながる可能性があります。

パディングは必要に応じて XLA コンパイラによって自動的に行われますが、op_profile ツールを使用して、行われるパディングの量を決定できます。TPU に適したテンソル次元を選択することで、パディングを回避できます。

次元

適切なテンソル次元を選択すると、TPU ハードウェア、特に MXU から最大のパフォーマンスを引き出すために役立ちます。XLA コンパイラは、バッチサイズまたは特徴ディメンションのいずれかを使用して MXU を最大限に利用しようとします。したがって、これらのいずれかは 128 の倍数である必要があります。そうでない場合、コンパイラによっていずれかが 128 までパディングされます。理想的には、バッチサイズとフィーチャ次元は 8 の倍数である必要があり、これによりメモリ サブシステムから高パフォーマンスを引き出すことができます。

演算

利用可能な TensorFlow 演算のリストをご覧ください。

VPC Service Controls の統合

Cloud TPU のVPC Service Controlsを使用して、Cloud TPU リソースのセキュリティ境界を定義し、その境界をまたがるデータの移動を制御できます。VPC Service Controls の詳細については、VPC Service Controls の概要をご覧ください。VPC Service Controls とともに Cloud TPU を使用する際の制限事項については、サポートされているプロダクトと制限事項をご覧ください。

Edge TPU

クラウドでトレーニングされた機械学習モデルは、「エッジ」（モノのインターネット（IoT）のエッジで動作するデバイス）で推論を実行する必要があります。こうしたデバイスには、リアルタイムのデータを収集し、インテリジェントな意思決定を行った後で、アクションしたり、他のデバイスやクラウドに情報を伝えるセンサーや他のスマート デバイスがあります。

このようなデバイスは限られた電力（電池を含む）で動作する必要があるため、Google は Edge TPU コプロセッサを低電力デバイス上で ML 推論を高速化するように設計しました。個々の Edge TPU は 1 秒間に 4 兆回のオペレーション（4 TOPS）を実行できますが、2 ワットの電力しか消費しません。言い換えると、1 ワットあたり 2 TOPS です。たとえば、Edge TPU では電力消費を抑えつつ、MobileNet V2 などの最先端のモバイル ビジョンモデルを約 400 フレーム / 秒で実行できます。

この低消費電力 ML アクセラレータは、Cloud TPU と Cloud IoT を強化し、エンドツーエンド（クラウドツーエッジ、ハードウェア + ソフトウェア）のインフラストラクチャを提供します。これにより、AI ベースのソリューションの提供が簡単になります。

Edge TPU は、シングルボード コンピュータ、システム オンモジュール、PCIe/M.2 カード、サーフェス マウント モジュールなど、さまざまなフォーム ファクタのプロトタイピング デバイスと製品版デバイスで利用可能です。Edge TPU とすべての利用可能な製品の詳細については、coral.ai をご覧ください。

次のステップ