Cloud TPU の概要

Tensor Processing Unit(TPU)は、機械学習ワークロード用に Google が設計したハードウェア アクセラレータです。TPU ハードウェアの詳細については、システム アーキテクチャをご覧ください。Cloud TPU は、TPU を Google Cloud でスケーラブルなコンピューティング リソースとして利用できるようにするウェブサービスです。

TPU は、機械学習アルゴリズムでよく見られる大規模な行列演算を実行するために設計されたハードウェアを使用して、モデルを効率的にトレーニングします。TPU にはオンチップの高帯域幅メモリ(HBM)が搭載されており、大きいモデルやバッチサイズを使用できます。TPU は Pod と呼ばれるグループに接続でき、コードをほとんど変更することなくワークロードをスケールアップできます。

単一の TPU で機械学習プロジェクトを開始し、本番環境用に TPU Pod にスケールアウトすることをおすすめします。開始するには、Google の研究チームとエンジニアリング チームが TPU で使用するために最適化した一連のオープンソース リファレンス モデルを利用できます。詳細については、参照モデルをご覧ください。

仕組み

TPU の仕組みを理解するには、ML モデルのトレーニングにおける計算上の課題に、他のアクセラレータがどのように対処しているかを理解することが役立ちます。

CPU の仕組み

CPU は、フォンノイマン アーキテクチャに基づく汎用プロセッサです。つまり、CPU は以下に示すようにソフトウェアとメモリと連動して機能します。

CPU の仕組みを表す図

CPU の最大のメリットは、その柔軟性です。さまざまな種類のアプリケーションに対して、任意の種類のソフトウェアを CPU に読み込むことができます。たとえば、PC での文書処理、ロケット エンジンの制御、銀行取引、さらにはニューラル ネットワークを使用した画像分類にも CPU を使用できます。

CPU はメモリから値を読み込み、計算を行い、すべての計算結果をメモリに保存します。メモリアクセスは、計算速度に比べて遅く、CPU の総スループットを制限することがあります。これはフォンノイマン ボトルネックとも呼ばれます。

GPU の仕組み

スループットを向上させるため、GPU には単一のプロセッサ内に数千の算術論理演算ユニット(ALU)が含まれています。最新の GPU には、通常 2,500~5,000 個の ALU が含まれています。大量のプロセッサが存在するということは、数千の乗算と加算を同時に実行できるということになります。

GPU の仕組みを表す図

ニューラル ネットワークでの行列演算のように大量の並列処理を行うアプリケーションでは、この GPU アーキテクチャが効果を発揮します。実際、ディープ ラーニングの一般的なトレーニング ワークロードでは、GPU によって CPU とは桁違いのスループットを達成できます。

しかし、GPU が汎用プロセッサであることに変わりはなく、数多くのアプリケーションやソフトウェアをサポートしなければなりません。したがって、GPU にも CPU と同じ問題があります。何千もの ALU で行われるすべての計算で、GPU はレジスタや共有メモリにアクセスして、オペランドを読み取り、計算の中間結果保存しなければなりません。

TPU の仕組み

Cloud TPU は、Google がニューラル ネットワークのワークロードに特化して設計した行列プロセッサです。TPU は文書処理、ロケット エンジンの制御、銀行取引といった操作に対応できませんが、ニューラル ネットワークで使用される大規模な行列演算は極めて高速に処理できます。

TPU の主なタスクは、乗算演算と累積演算の組み合わせである行列処理です。TPU には数千の乗算アキュムレータが含まれています。これらは、直接接続して大規模な物理マトリックスを形成しています。この構造は、シストリック アレイ アーキテクチャと呼ばれています。Cloud TPU v3 では、単一プロセッサ上に 128 x 128 ALU のシストリック アレイが 2 つあります。

TPU ホストは、インフィード キューにデータをストリーミングします。TPU は、インフィード キューからデータを読み込んで HBM メモリに保存します。計算が完了すると、TPU は結果をアウトフィード キューに読み込みます。TPU ホストは次に、アウトフィード キューから結果を読み取り、ホストのメモリに保存します。

行列演算を実行するために、TPU は HBM メモリからパラメータを MXU に読み込みます。

TPU がメモリからパラメータを読み込む方法を示す図

続いて、TPU は HBM メモリからデータを読み込みます。乗算が行われるたびに、その結果が次の乗算アキュムレータに渡されます。データとパラメータの乗算結果をすべて合計したものが出力となります。 行列乗算処理中にメモリアクセスは必要ありません。

TPU がメモリからデータを読み込む方法を示す図

その結果、TPU はニューラル ネットワークにおいて、高い演算スループットを達成できます。

TPU の用途

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

CPU

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

GPU

  • CPU 上で少なくとも部分的に実行しなければならない多数のカスタムの TensorFlow / PyTorch / JAX 演算を使用するモデル
  • Cloud TPU で利用できない TensorFlow / PyTorch 演算を使用するモデル
  • 実際のバッチサイズが大きい中~大規模なモデル

TPU

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

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

  • 頻繁な分岐を必要とするか、要素ごとの代数演算を多く含む線形代数プログラム
  • メモリにあまりアクセスしないワークロード
  • 高精度の演算を必要とするワークロード
  • メインのトレーニング ループにカスタム演算が含まれるニューラル ネットワーク ワークロード

VPC Service Controls の統合

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

Cloud TPU のスタートガイド

Google Cloud アカウントの設定 Cloud TPU リソースを使用するには、Google Cloud のアカウントとプロジェクトを作成する必要があります。
Cloud TPU API を有効にする モデルをトレーニングするには、Compute Engine API と Cloud TPU API を有効にする必要があります。
Cloud Storage TPU に Cloud Storage バケットへのアクセスを許可する Cloud TPU では、データセットを保存する Cloud Storage バケットにアクセスする必要があります。
モデルのトレーニング まずは Cloud TPU のクイックスタートチュートリアルをご覧ください。
モデルを分析する TensorBoard などのツールを使用して、モデルを可視化し、主要な指標を追跡します。

次のステップ

Cloud TPU の詳細を調べるには、次のリソースも役立ちます。

Compute Engine を使用したクイックスタート クイックスタートのいずれかで、Cloud TPU を使用してモデルをトレーニングしてみてください。
TPU Colab 無料の Colab を使って Cloud TPU を試してみてください。
Cloud TPU のチュートリアル 参照 ML モデルを使って Cloud TPU を試してみてください。
Cloud TPU システム アーキテクチャ TPU の詳細を確認する。
料金 Cloud TPU では費用効果の高い方法で機械学習ワークロードを処理する仕組みを把握してください。
お問い合わせ 具体的な実装やアプリケーションに関するディスカッションをご希望の場合は、Cloud TPU の営業担当者にお問い合わせください。