Google の第 2 世代カスタム ニューラル ネットワーク数学アクセラレーション チップである TPU2 の潜在的なパフォーマンスに興味があったら、ここに早めのクリスマス プレゼントがあります。
Google のベテラン エンジニアリング エンジニア Jeff Dean 氏と、同じく Google 社員の Chris Ying 氏は、先週カリフォルニア州ロサンゼルスで開催された Neural Information Processing Systems (NIPS) カンファレンスで、このシリコンに関するさらに詳しい情報 [PDF] を明らかにしました。
インターネット広告大手のGoogleはTPU2について口を閉ざしており、チップセットのアーキテクチャや仕様を詳述する資料や論文の公開を今のところ拒否している。分かっているのは、このチップセットがAIソフトウェアに必要な計算を高速化し、汎用プロセッサやGPUの負荷を軽減することで、チップ大手（Intel、NVIDIAなど）によるGoogleへの支配を弱めるために設計されているということだけだ。
カンファレンスの機械学習システム ワークショップでの Dean 氏の講演で、彼は第 2 世代 Tensor Processing Unit v2 デバイスにはそれぞれ次の機能を備えた 4 つのチップが搭載されていることを明らかにしました。
- 2 つのコア、それぞれに 128x128 の混合乗算ユニット (MXU) と 8GB の高帯域幅メモリが搭載され、1 つの 4 チップ デバイスで最大 64GB の HBM が追加されます。
- 600 GB/s のメモリ帯域幅。
- スカラーとベクトル用の 32 ビット浮動小数点精度の演算ユニットと、乗算器用の精度を下げた 32 ビット浮動小数点精度の行列乗算ユニット。
- 最大パフォーマンスは約 45 TFLOPS で、4 チップデバイス 1 つで合計 180 TFLOPS になります。
4つのプロセッサチップを搭載した単一のTPU2デバイス…赤い枠で囲まれた部分は1つのチップを示しています。画像クレジット：Google
各チップにはそれぞれ8GBのHBMを備えた2つのコアが含まれています
TPU2 ASICは、前世代のモデルとは異なり、学習ワークロードだけでなく推論ワークロードも処理できます。64個のTPU2デバイスを接続し、4TBの高帯域幅メモリを使用して最大11.5PFLOPSの演算能力を持つポッドを構成できるように設計されています。
Ying氏はこれを、学習時間を短縮するために必要な膨大な計算能力を提供できる機械学習用スーパーコンピュータと呼んだ。市販のx86およびGPUハードウェアを用いたニューラルネットワークモデルの学習プロセスでは、大規模なデータセットを投入するのに数日、場合によっては数週間かかることもある。研究者は多くの場合、このプロセスを繰り返し、学習データセットを複数回、あるいは複数エポックにわたって実行し、重みとパラメータを調整することで、モデルの最高のパフォーマンスを引き出す。そこでGoogleは、開発時間をわずか数時間に短縮するために、専用のTPUギアを開発した。
TPU2は初期の段階で一定の成功を収めています。Googleは既に、Google Homeデジタルアシスタント向けに人間のような音声を生成するAIシステム「WaveNet」にTPU2を採用しています。TPU2はTensorFlowでプログラミングされており、こちらのサインアップページからハードウェアへのクラウドアクセスをリクエストできます。この製品は店頭で購入することはできません。Googleは社内でこのチップを使用し、自社クラウド経由でアクセスを許可しているからです。
Ying氏は、画像認識モデルResNet-50 v2のトレーニング時間を、様々なカテゴリに分類された数百万枚のラベル付き画像を含む人気のデータベースImageNetを用いて、TPU2の数を増やしながら示した。1台のTPU2で、バッチサイズ256枚の写真、90エポック（データセットを90回実行）のResNet-50 v2をトレーニングし、精度76.6%を達成したのに23時間22分かかった。一方、32台のTPU2（ポッドの半分）で、90エポック、バッチサイズ8,192、精度76.1%を達成したのにはわずか45分しかかからなかった。
コーヒーを一杯飲みに行くのと、モデルの学習が終わるまで数晩寝るのとでは大きな違いがあるとディーン氏は言います。学習のターンアラウンドタイムが短いほど生産性が向上し、研究者は社内のAutoMLプロジェクトのような新しいタイプの研究を試すことができます。
AutoMLチームの研究者たちは、AI設計プロセスを自動化するために機械を活用し、新しいニューラルネットワークアーキテクチャを見つけ出しています。このプロセスでは、様々な構成要素を探索し、新しい戦略を試した上で、最終的に人間の開発者が手作業で構築したものよりもはるかに大規模で複雑なアーキテクチャに落ち着くことになります。
しかし、TPU2の詳細はまだ不明な点が多い。Ying氏のワークショップ講演で聴衆からの質問に対し、同氏は行列乗算の精度や消費電力など、チップの一部の側面については言及できないと述べた。El Reg氏は、TPU2は少数の命令しかサポートしていなかった第1世代TPUよりもはるかに進化しているものの、プログラミングが容易ではないと聞いている。そのため、Googleの幹部はTensorFlowや並列化された行列演算を扱えるエンジニアを探し出し、TPU2搭載プロジェクトに再配置しているという。
現時点では、TPU2 は Google のエッグヘッドとエンジニアのみが利用できますが、「オープンな機械学習研究に取り組んでいるトップクラスの研究者」には 1,000 台を無料で提供する計画があります。®
