[論文レビュー] TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings

Norman P. Jouppi, George Thomas Kurian+12 more|arXiv (Cornell University)|Apr 4, 2023

Photonic and Optical Devices被引用数 19

ひとこと要約

TPU v4 は光学的に再構成可能なインターコネクトを optical circuit switches を用いて導入し、埋め込み中心のワークロード向けに SparseCores を追加し、ML パフォーマンスと効率を高めるために 4096 チップへスケールします。

ABSTRACT

In response to innovations in machine learning (ML) models, production workloads changed radically and rapidly. TPU v4 is the fifth Google domain specific architecture (DSA) and its third supercomputer for such ML models. Optical circuit switches (OCSes) dynamically reconfigure its interconnect topology to improve scale, availability, utilization, modularity, deployment, security, power, and performance; users can pick a twisted 3D torus topology if desired. Much cheaper, lower power, and faster than Infiniband, OCSes and underlying optical components are <5% of system cost and <3% of system power. Each TPU v4 includes SparseCores, dataflow processors that accelerate models that rely on embeddings by 5x-7x yet use only 5% of die area and power. Deployed since 2020, TPU v4 outperforms TPU v3 by 2.1x and improves performance/Watt by 2.7x. The TPU v4 supercomputer is 4x larger at 4096 chips and thus ~10x faster overall, which along with OCS flexibility helps large language models. For similar sized systems, it is ~4.3x-4.5x faster than the Graphcore IPU Bow and is 1.2x-1.7x faster and uses 1.3x-1.9x less power than the Nvidia A100. TPU v4s inside the energy-optimized warehouse scale computers of Google Cloud use ~3x less energy and produce ~20x less CO2e than contemporary DSAs in a typical on-premise data center.

研究の動機と目的

• 実用的でスケーラブルなMLスーパーコンピューティング用ハードウェアの必要性を、急速に進化するMLモデルとワークロードに対応して喚起する。
• Optical circuit switches を含む TPU v4 アーキテクチャと SparseCores 埋め込みアクセラレータを説明する。
• 既存の TPU および関連 DSA との比較で性能、電力、エネルギー効率の改善を定量化する。
• 大規模言語モデルとデータセンターのエネルギー消費に関する展開への影響を強調する。

提案手法

• Optical Circuit Switches（OCSes）を活用して、スケール、可用性、利用状況、モジュール性、デプロイ、セキュリティ、電力、パフォーマンスのために相互接続トポロジを動的に再構成する。
• 埋め込み重視モデル向けに最適化されたデータフロー処理ユニットである SparseCores を統合し、5x-7x の加速を達成する一方で die 面積と電力の約5%を使用する。
• TPU v4 を 4倍大きい 4096 チップのスーパーコンピュータとして展開し、 TPU v3 と比較して全体性能を約10x向上させる。
• Graphcore IPU Bow および Nvidia A100 と同等のシステム規模に対する比較的な性能と効率の指標を提供する。
• Google Cloud のウェアハウス規模デプロイと、一般的なオンププレミスのデータセンターにおけるエネルギーと CO2e の影響を評価する。

実験結果

リサーチクエスチョン

• RQ1 OCSes を介した光学的再構成性は、MLワークロードのインターコネクト性能、利用率、信頼性にどのような影響を与えるか？
• RQ2 SparseCores が埋め込み中心のモデルを加速する効果と、得られる面積と電力のフットプリントはどうなるか？
• RQ3 大規模デプロイメントにおける TPU v4 の TPU v3、Graphcore IPU Bow、Nvidia A100 に対する性能、電力、CO2e のメリットは何か？
• RQ4 TPU v4 は 大規模言語モデルとウェアハウス規模データセンターの展開上の考慮にどのような影響を与えるか？

主な発見

• TPU v4 は OCSes を用いて、スケールと利用率を改善するための動的インターコネクト トポロジ（例：ねじれた3Dトーラス）を実現する。
• SparseCores は、埋め込み中心のモデルを約5%のダイ面積と電力で 5x-7x の速度向上を実現する。
• TPU v4 は TPU v3 より 2.1x 高速で、性能/ワットを 2.7x 向上させる。
• The 4096-chip TPU v4 は TPU v3 より全体で約10x、同様の規模の Graphcore IPU Bow より約4.3x-4.5x 高速。
• TPU v4 は Google Cloud のウェアハウス規模デプロイメントで、典型的なオンプレミス DSA と比較してエネルギー消費が約3倍削減、CO2eは約20x削減。