Googleの機械学習マシン「第2世代TPU」の構造を公開された写真から推測する

Googleが先日開催されたGoogle I/O 2017で、機械学習用に最適化された第2世代の「TensorFlow Processing Unit(TPU)」を発表しました。しかし、Googleは第2世代TPU(TPU2)の詳細なスペックを明らかにしていないため、TPU2の構造はよく分かりません。そこで、クラウドコンピューティングやデータ分析を行うTIRIAS ResearchのチーフアナリストのPaul Teich氏が、公開された写真からスペックを推測しています。

Under The Hood Of Google’s TPU2 Machine Learning Clusters
https://www.nextplatform.com/2017/05/22/hood-googles-tpu2-machine-learning-clusters/

第1世代TPU(TPU1)が16ビット演算(FP16)で23TOPSだったのに対して、TPU2は浮動小数点演算で45TFLOPSの能力を持ちます。TPU1は、PCI-Express(x8)コネクタでマザーボードに接続するアクセラレーターでした。TPU1にはスケジュール機能がなく、仮想化も不可。プロセッサボードに接続して、単純なマトリックスコプロセッサとして機能します。ちなみにGoogleは、1枚のマザーボード上にTPU1を何枚まで挿して使えるのかについて明らかにしていません。

これに対してGoogle I/O 2017で発表されたTPU2ボード。なお、Googleは内製マシンとしてTPU2を活用しており、TPU2が搭載されたボードやサーバーが外部に販売されることはありません。ユーザーは、TensorFlow Research Cloud(TRC)プログラムを通じて、TPU2のマシンパワーを借りることができるだけです。

Teich氏がGoogleの公開したポッド群の写真を見て最初に気づいたのは、「TPU2は2本のケーブルで1つのプロセッサボードに接続している」ということ。「TPU2が異なる2つのプロセッサボードにつながっている可能性はないわけではないが、インストール作業やプログラミング、スケジューリングを複雑にしないために、TPU2とプロセッサボードは1：1で接続しているだろう」と考えています。

TPU2のスタンプの画像。Aが「CPUラック」、BとCが「TPU2ラック」、Dが「CPUラック」、青色部分がソリッドボックス。システムは4つのラックに収まるスタンプでデザインされており、4ラック単位で1つの「pod(ポッド)」を構成。4つのスタンプには1024個のTPU2チップが含まれています。

GoogleはTPU2スタンプに関して異なる3つの写真を公開していますが、ラックの構成と配線はすべて同じように見えるとのこと。ケーブルの色分けは比較を容易にしています。

1枚のボードにはヒートシンク(A)の下に4つのTPUチップが搭載されています。1基のTPU2につき2本伸びるBlueLink(25GB/s)ケーブル(B)、ボードあたり2つあるOmni-Path Architecture(OPA)ケーブル用コネクタ(C)、電源コネクタ(D)、ネットワークスイッチ(E)という構成だろうとTeich氏は推測しています。

フロントパネルには見たことがない正方形の形をした4つのコネクタ。GoogleがOpenCAPIのメンバーであり、OpenPowerFoundationの創設メンバーであることから、これらはBlueLinkだと考えるのが理にかなっているとのこと。ちなみにIBMのBlueLinkでは、サブリンクあたり25GB/sを8方向出して200GB/sの信号レーン(合計16レーン)が定義されているそうです。BlueLinkないしOPAの銅線ケーブルの下に、色分けされたステッカーがあることから、Teich氏はカラーコードを用いてTPU2スタンプを大規模に展開する計画がありそうだと指摘しています。

TPU2スタンプのラック構成は左右対称になっています。なお、下図のDの画像は、対称なのが分かるように左右反転させて比較しているとのこと。Googleのラックユニットの高さは業界標準の1.75インチより小さな1.5インチだとのこと。

同じくCも左右反転の画像。「Googleの公開した画像からはネットワーク・トポロジーを正確に判断することはできないが、何らかのハイパーメッシュネットワークが構築されているように見える」とTeich氏は述べています。

1.5インチサーバーユニットのフォームファクタに合致するCPUボードは、Intel Xeonデュアルソケット対応のマザーボードです。ただし、AMDのEPYC対応サーバーやARMチップを搭載するQualcommのCentriq対応サーバーがシングルソケット仕様であることから、TPU2が使うデュアルソケットマザーボードの構成は、将来的に変化する可能性もありそうだとのこと。

また、各CPUボード1枚に対して1枚のTPU2ボードを接続することで、OPAケーブルを介して25GB/sの総帯域幅に到達させるとTeich氏は推測しています。一般的なディープラーニング用マシンにおいて、GPUアクセラレーターとCPUとの接続比は4:1から6:1であるのに対して、TPU2：Xeon CPUは2：1という低い値であるのが特徴的です。これは、TPU1がGPUよりもFPU(浮動小数点ユニット)に親和的な設計だった思想を引き継いだことから来るのではないかとのこと。

BからEまでの他のサーバーボードに比べると圧倒的に大きなヒートシンクを持つTPU2ボード(A)。ほぼ、ラック幅2つ分で3インチ(約7.6センチメートル)程度で、ヒートシンクの大きさはそれぞれ消費電力が200Wを超えることを物語っているとのこと。

Teich氏はTPU2システムの消費電力についても考察しています。Googleによると、TRCでは1000個(正確には1024個)のTPU2チップが使われているとのこと。そこから、フルパワーで演算処理するTRCの消費電力は0.5メガワット(50万W)だと見積もっています。4つのスタンプからなるTPU2ユニットのイニシャルコストは導入時の1回でありたいしたことはなくとも、0.5メガワットのランニングコストはGoogleレベルの企業でもバカにできないコストになるだろうとTeich氏は述べています。仮に今後、1年の間、GoogleがTPU1を併用するならば、それはGoogleがTPU2の電力コストを深刻に考えていることの証になるかもしれないとのことです。