Google TPU v5e

What the v5e is

Google TPU v5e is a machine-learning accelerator chip from Google's tensor processing unit (TPU) line. TPUs are chips Google designs for the matrix-heavy arithmetic that dominates neural networks. Wikipedia lists the v5e as introduced in 2023 ^[2], and Google positions it for both training and serving ^[1] — that is, both for building models and for running them for users.

What is on the chip

Each v5e chip is built around one TensorCore. That TensorCore contains four matrix multiplication units, which carry the dense linear algebra at the heart of a neural network layer, plus one vector unit for elementwise and other non-matrix operations and one scalar unit for the control and address arithmetic around them ^[1]. The point of the split is to keep the matrix units busy multiplying while everything that is not a matrix multiply goes to the other two units.

Compute and memory

Google quotes 197 TFLOP of peak bfloat16 compute and 393 TOPs of peak INT8 compute for a single chip ^[1]. bfloat16 is a 16-bit floating-point format that keeps the exponent range of a 32-bit float, which is why it is widely used in training; INT8 is an 8-bit integer format used mostly for inference. The INT8 number coming out at roughly twice the bfloat16 number is the usual pattern when the same multiplier hardware performs twice as many narrower operations per cycle.

Memory per chip is 16 GB of high bandwidth memory (HBM) with 800 GiBps of bandwidth in Google's documentation ^[1]; Wikipedia lists the same 16 GB capacity with 819 GB/s of bandwidth ^[2]. Capacity decides how large a model, or how large a batch, fits on one chip, while bandwidth decides how quickly weights and activations can be fed to the arithmetic units.

Hosts and the links between chips

Chips do not work alone. Each host carries 512 GiB of DRAM and up to eight v5e chips ^[1]. Chips on a host, and chips across hosts, communicate through the inter-chip interconnect (ICI): each chip has four ICI ports and 400 GBps of bidirectional chip-to-chip bandwidth ^[1]. Because a single model layer is usually split across many chips, that link carries the gradients and activations of distributed training, and it is often what limits how far a job can scale.

Pods

Google groups v5e chips into pods of 256 chips wired on a two-dimensional torus, a mesh in which each chip connects to its neighbours and the edges wrap around, so that every chip has the same number of links ^[1]. A full pod reaches 50.63 PFLOPs of bfloat16 peak compute ^[1]. Within a pod, all-reduce bandwidth is 51.2 TB/s and bisection bandwidth is 1.6 TB/s ^[1]. All-reduce is the collective operation in which every chip ends up holding the sum of all chips' gradients — the synchronisation step of a data-parallel training run — while bisection bandwidth measures how much traffic can cross the middle of the network at once, which matters when many chips must exchange data at the same time.

Training and serving

Google positions the v5e for both training and serving, with training jobs tuned for throughput and serving jobs tuned for latency ^[1]. The distinction is practical. A training job pushes as many examples through the chip as possible and cares about total work per second, so large batches and high utilisation are good. A serving job answers one request at a time and cares about how fast a single answer comes back, so small batches and short response times matter more. One chip has to be configurable for both, which is why capacity, bandwidth and interconnect are quoted together as the things that decide which workloads it suits.