Cloud Tensor Processing Unit(TPU)

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

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

TPU の利点

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

可用性

Cloud TPU は、次のゾーンで使用できます。

  • 米国(US)
  • ヨーロッパ(EU)
    • europe-west4-a
  • アジア太平洋(APAC)
    • asia-east1-c

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

Cloud TPU は、密ベクトルおよび行列計算の実行が高速です。Cloud TPU とホストメモリ間のデータ転送は計算速度に比べて遅く、PCIe バスの速度は Cloud TPU Interconnect とオンチップの高帯域幅メモリ(HBM)のどちらよりもはるかに低速です。これは、モデルの部分的なコンパイル(実行がホストと端末の間で「行ったり来たり」する)では、データが PCIe バスで到達するのを待機して端末はほとんどの時間アイドル状態であるため、端末が非効率的な方法で使用されることを意味します。この状況を軽減するために、Cloud TPU のプログラミング モデルは、トレーニングの多くを TPU で実行するように設計されています(理想的にはトレーニング ループ全体)。

以下は、TPUEstimator によって実装されるプログラミング モデルの主な特徴です。

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

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

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

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

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

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

Cloud TPU ハードウェアには、密な行列乗算と畳み込みが多くを占めるプログラムを高速化するように設計されたマトリックス ユニット(MXU)が含まれています。実行時間のかなりの部分で行列乗算が実行されないプログラムは、Cloud TPU にあまり適していない可能性があります。一例として、計算の多くを非行列演算(加算、再形成、連結など)が占めるプログラムは、MXU 使用率が高くならない可能性があります。次に、Cloud TPU に適したモデルを選択して構築するために役立つガイドラインを示します。

レプリカ

1 つの Cloud TPU 端末は 4 つのチップで構成され、それぞれのチップに 2 つの TPU コアがあります。したがって、Cloud TPU を効率的に利用するには、プログラムで 8 つのコアそれぞれを使用する必要があります。 TPUEstimator によって、複製された計算を構築して実行するためのグラフ演算子が提供されます。各レプリカは本質的には各コアで実行されるトレーニング グラフのコピーであり、これによりバッチサイズ全体の 1/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 演算のリストをご覧ください。

TPU の用途

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

  • CPU:

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

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

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

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

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

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

次のステップ

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

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