TPU v4

This document describes the architecture and supported configurations of Cloud TPU v4.

System architecture

Each TPU v4 chip contains two TensorCores. Each TensorCore has four matrix-multiply units (MXUs), a vector unit, and a scalar unit. The following table shows the key specifications for a v4 TPU Pod.

Specification Values
Peak compute per chip 275 teraflops (bf16 or int8)
HBM2 capacity and bandwidth 32 GiB, 1200 GBps
Measured min/mean/max power 90/170/192 W
TPU Pod size 4096 chips
Interconnect topology 3D mesh
Peak compute per Pod 1.1 exaflops (bf16 or int8)
All-reduce bandwidth per Pod 1.1 PB/s
Bisection bandwidth per Pod 24 TB/s

The following diagram illustrates a TPU v4 chip.

Diagram of a TPU v4 chip

For more information about architectural details and performance characteristics for TPU v4, see TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings.

3D mesh and 3D torus

v4 TPUs have a direct connection to the nearest neighboring chips in 3 dimensions, resulting in a 3D mesh of networking connections. The connections can be configured as a 3D torus on slices where the topology, AxBxC, is either 2A=B=C or 2A=2B=C, where each dimension is a multiple of 4. For example, 4x4x8, 4x8x8, or 12x12x24. In general, the performance of a 3D torus configuration will be better than a 3D mesh configuration. For more information, see Twisted torus topologies.

Performance benefits of TPU v4 over v3

This section shows a memory-efficient way to run a sample training script on TPU v4, as well as the performance improvements for TPU v4 compared to TPU v3.

Memory system

Non Uniform Memory Access (NUMA) is a computer memory architecture for machines that have multiple CPUs. Each CPU has direct access to a block of high-speed memory. A CPU and its memory is called a NUMA node. NUMA nodes are connected to other NUMA nodes that are directly adjacent to each other. A CPU from one NUMA node can access memory in another NUMA node, but this access is slower than accessing memory within a NUMA node.

Software running on a multi-CPU machine can place data needed by a CPU within its NUMA node, increasing memory throughput. For more information about NUMA, see Non Uniform Memory Access on Wikipedia.

You can take advantage of NUMA-locality benefits by binding your training script to NUMA Node 0.

To enable NUMA node binding:

  1. Create a TPU v4.

  2. Connect to the TPU using SSH.

  3. Install the numactl command line tool. numactl lets you run processes with a specific NUMA scheduling or memory placement policy.

     $ sudo apt-get update
 $ sudo apt-get install numactl

  4. Bind your script code to NUMA Node 0. Replace your-training-script with the path to your training script. 

     $ numactl --cpunodebind