Google TPU v4

What TPU v4 is

TPU v4 is a machine-learning chip and the supercomputer Google built around it. It has been deployed since 2020, and the paper Google published at ISCA 2023, by Norman Jouppi, David Patterson and colleagues, calls it the company's fifth domain-specific architecture and its third supercomputer built for machine-learning models^[1]. "Domain-specific" means the hardware is designed around one kind of work — the matrix arithmetic of neural networks — rather than being a general-purpose processor.

The chip

Each TPU v4 chip carries two TensorCores, and each TensorCore holds four matrix-multiply units: the blocks that compute the sums of products at the heart of a neural-network layer^[2]. Peak throughput is 275 teraflops in bf16 or int8, where bf16 is a 16-bit floating-point format and int8 an 8-bit integer format, both cheaper than full 32-bit arithmetic^[2]. Each chip has 32 GiB of HBM2 memory running at 1,200 GB/s, and measured power ranges from 90 W at minimum through 170 W mean to 192 W maximum^[2]. It is made on a 7 nm process, down from 16 nm for TPU v3^[2].

Pods and the optical switch

Chips are grouped into pods of 4,096, which together deliver 1.1 exaflops in bf16 or int8^[2]. The paper reports that this supercomputer is 4x larger than TPU v3's and nearly 10x faster overall, with each chip 2.1x faster than TPU v3 and 2.7x better performance per watt^[2].

The unusual part is how the chips are wired. Optical circuit switches dynamically rewire the links between 4×4×4 cubes of chips, so a user can choose a topology such as a twisted 3D torus instead of accepting one fixed network^[1]. Google says the optics account for under 5% of system cost and under 3% of system power^[1]. In the Gemini 1.0 report of December 2023, Google described each 4,096-chip SuperPod as having a dedicated optical switch able to reconfigure 4×4×4 cubes into a 3D torus in about 10 seconds^[1].

SparseCores for embeddings

TPU v4 also adds SparseCores, dataflow processors aimed at embedding-heavy models such as recommenders. Google reports that they make such models 5–7x faster while taking about 5% of die area and power^[1]. Embeddings — large tables mapping identifiers to vectors — are memory-bound lookups rather than dense arithmetic, which is why a separate block helps.

How it compared

For similar-sized systems, Google reported TPU v4 as 1.2–1.7x faster than the Nvidia A100 and using 1.3–1.9x less power; this is Google's own comparison rather than an independent benchmark^[1].

What ran on it

PaLM, from April 2022, was trained on 6,144 TPU v4 chips across two pods using the Pathways system, which the TPU v4 paper describes as the largest TPU configuration described to date^[3]. The Gemini 1.0 report of December 2023 says the Gemini models were trained on TPUv5e and TPUv4, with Gemini Ultra using a large fleet of TPUv4 accelerators across multiple data centers^[1].

Why it matters

TPU v4 shows two directions for machine-learning hardware. One is specialization: separate blocks for dense matrix math and for sparse embedding lookups, in place of one general unit. The other is reconfigurable networking: if the interconnect can be reshaped in software, a single machine can serve workloads with different communication patterns. The Gemini 1.0 report's account of the optical switch shows the design carrying forward into later systems, where the same pods were used to train models at a much larger scale^[1].