[论文解读] 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 引入光学可重构互连，使用光学电路开关；新增用于嵌入式工作负载的 SparseCores；并扩展到 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 超级计算硬件。
• 描述 TPU v4 架构，包括光学电路开关和 SparseCores 嵌入式加速器。
• 量化相对于早期 TPU 及相关 DSA 的性能、功耗与能效提升。
• 突出大型语言模型部署及数据中心能耗足迹的影响。

提出的方法

• 利用 Optical Circuit Switches（OCSes）动态重新配置互连拓扑，以提升规模、可用性、利用率、模块化、部署、安全、功率和性能。
• 集成 SparseCores，这些数据流处理器针对嵌入密集型模型进行优化，在使用约 5% 的芯片面积和功耗的同时实现 5x-7x 的加速。
• 将 TPU v4 部署为一个4倍规模的4096芯片超算，相较于 TPU v3 可实现高达约10x 的整体性能提升。
• 给出与 Graphcore IPU Bow 和 Nvidia A100 在相似系统规模下的性能与效率对比指标。
• 评估在 Google 云端仓库级部署中与典型本地数据中心相比的能源与 CO2e 影响。

实验结果

研究问题

• RQ1 通过 OCSes 的光学可重构性如何影响 ML 工作负载的互连性能、利用率和可靠性？
• RQ2 SparseCores 在加速嵌入为主的模型方面的有效性如何，所带来的人线面积和功耗足迹是多少？
• RQ3 在大规模部署中，TPU v4 相对于 TPU v3、Graphcore IPU Bow 和 Nvidia A100 的性能、功耗和 CO2e 有何提升？
• RQ4 TPU v4 如何影响大型语言模型部署和仓库级数据中心的部署考虑？

主要发现

• 搭载 OCSes 的 TPU v4 实现了动态互连拓扑选择（例如扭曲的 3D torus），以提升规模和利用率。
• SparseCores 在嵌入密集型模型上实现 5x-7x 的加速，使用约 5% 的芯片面积和功耗。
• TPU v4 比 TPU v3 快 2.1x，性能/功耗提升 2.7x。
• 4096 芯片的 TPU v4 在整体上比 TPU v3 快约 10x，在相同规模下比 Graphcore IPU Bow 快约 4.3x-4.5x。
• 在 Google Cloud 仓库规模部署中，TPU v4 的能耗更低，CO2e 排放也更少，与典型的本地数据中心 DSA 相比大约少 3 倍能源和大约少 20 倍 CO2e。