TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings

By N. Jouppi, George Kurian, Sheng Li, Peter C. Ma, R. Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing, Brian Towles and 4 more. 2023

685 citations recorded by the source corpus.

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 nearly 10x faster overall, which along with OCS flexibility and availability allows a large language model to train at an average of ~60% of peak FLOPS/second. 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 ~2--6x less energy and produce ~20x less CO2e than contemporary DSAs in typical on-premise data centers.

Subfields: Computer Networks and Communications, Hardware and Architecture, Information Systems

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Networks and Communications

1. Yann LeCun
2. Saining Xie
3. Rob Fergus

In Hardware and Architecture

1. Yann LeCun

In Information Systems

1. Yann LeCun
2. Rob Fergus

Citations in the corpus

• SoundStorm: Efficient Parallel Audio Generation (2023, neutral): “When synthesizing dialogues of 30 seconds, we measure a runtime of 2 seconds on a single TPU-v4 (Jouppi et al., 2023).”
• Cambrian-1: A Fully Open, Vision-Centric Exploration of Multimodal LLMs (2024, methodology): “In this work, we train our model on TPU-V4 [53] with FSDP [141] using TorchXLA.”