TPU Clusters

A TPU is Google's accelerator for the matrix arithmetic that dominates neural network training and inference. A TPU cluster connects multiple chips so that a workload can distribute model state, examples, and computation. A pod or a slice of one is a common unit; larger systems can span slices. This entry uses TPU v4 as a concrete architectural example, not as a description of the newest available generation.

The vocabulary is small and worth fixing, because published descriptions often blur it. A TensorCore is the compute unit. A chip contains a fixed number of them: Google's documentation states two TensorCores per chip for both v3 and v4. A host is the CPU machine that drives some chips. A pod is the full interconnected set — 1,024 chips for v3, 4,096 for v4.

The arithmetic that decides the shape of a run

Capacity can force distribution before speed enters the decision: 500 billion bf16 parameters occupy about 931 GiB, while one v4 chip has 32 GiB. That is at least 30 chips of aggregate HBM for weights alone, before gradients, optimizer state, activations and supported topology constraints. Aggregate free memory is only a lower bound: each shard must fit where it is placed, and a supported chip topology must be allocated.

The distribution strategy determines what travels between chips. Data parallelism keeps a model replica on each worker and gives each a different part of the batch. Workers combine gradients so their updates remain synchronized; this reduces time per global batch but does not reduce the memory occupied by each full replica. Sharding weights, gradients, or optimizer state addresses a different memory constraint and introduces the communication needed to assemble or use those pieces.

Tensor parallelism divides an operation, such as a large matrix multiplication, among chips. Its partial results must be combined or exchanged before dependent operations proceed. Pipeline parallelism places successive groups of layers on different workers and passes activations and gradients between stages. Splitting a batch into microbatches lets stages overlap, but filling and draining the pipeline leaves idle periods. These methods can be combined; the placement must match the memory needs and frequency of communication.

A simple fixed-work example explains the scaling limit. For a replicated model with 2 MB of gradient data, an ideal ring all-reduce sends 2(N − 1)/N times that payload per participant over a complete reduce-scatter and all-gather, where N is the number of participants. At 8 participants that is 3.5 MB sent per participant; at 64 it is 3.9375 MB. Received traffic is additional if counting both directions. Meanwhile each participant’s share of a fixed global compute workload falls by eight when going from 8 to 64. Communication therefore becomes harder to hide behind useful computation. This is an illustrative collective algorithm, not a claim that a TPU runtime uses a ring for every reduction.

This is what the TPU v4 interconnect is built to address. Jouppi and colleagues describe the system as optically reconfigurable, reporting that the optical circuit switches and their components are “<5% of system cost and <3% of system power”. The interconnect is not an accessory to the accelerators; at pod scale it is the part that determines whether they can be kept busy.

Batch growth and serial work

A different scaling experiment keeps the local batch fixed. With 32 examples per replica, 8 replicas produce a global batch of 256 and 64 replicas produce 2,048. Covering two million examples takes 7,813 or 977 updates respectively if the last partial batch is retained. This changes the amount of work per update as well as the update count; it is not the fixed-global-work experiment above. Large batches can require changes to learning-rate schedules, warmup or other training settings, and their quality and throughput must be tested; more replicas do not by themselves preserve the optimization trajectory.

A separate fixed-work bound is Amdahl’s law: with serial fraction f and N ideal workers, speedup is 1 / (f + (1 − f)/N). If one percent of fixed work remains serial and the rest scales ideally, 4,096 workers give about 97.6 times the baseline speed. Data loading, checkpointing and reductions can contribute overhead, but their actual scaling must be measured; the assumed serial fraction is illustrative.

Reading the vendor numbers

Google publishes 275 teraflops of bf16 per v4 chip and 1.1 exaflops for a full 4,096-chip pod. Those are peak figures: the arithmetic rate if every unit were busy every cycle, rather than an achieved end-to-end training rate. They are useful for sizing and for comparing generations, and they are not a prediction of how long a run will take. The gap between peak and achieved throughput is where most of the engineering in large-scale training lives.

Comparative claims deserve the same care. It is common to read that TPUs are simply faster and more efficient than GPUs. The measured version is narrower and more useful: the v4 paper reports the system as 1.2×–1.7× faster than an NVIDIA A100 while using 1.3×–1.9× less power, on the workloads and system sizes it benchmarked. Those are real advantages with stated bounds, and they do not transfer automatically to a different model, a different generation of either chip, or a workload dominated by anything other than dense matrix arithmetic.

Keeping the distributed run productive

Hosts must feed batches, coordinate work, and move checkpoints while accelerators execute compiled computations. Slow input preparation can leave chips idle even when the interconnect is healthy. Shape changes or poorly supported operations can also change compilation or execution behavior. Measure input wait, compute, collective communication, and checkpoint time separately before adding chips. Arithmetic intensity helps explain device-level limits, while step-time measurements expose distributed overhead.

A synchronized step may wait for its slowest participant, and a failed participant can interrupt the job. Recovery requires a consistent checkpoint of the state needed to resume, such as parameters, optimizer state, step count, and the data-position or sampling state appropriate to the training procedure. Choose checkpoint frequency against lost work and checkpoint cost, and test restore on the intended allocation. More chips can reduce useful compute time while making coordination and recovery a larger part of the job.

Compare configurations by time and cost to a stated quality target, with global batch, token budget, precision, and stopping rule recorded. Throughput at a larger batch is useful evidence, but it does not establish the same learning trajectory. A larger slice is justified when memory fit or measured progress benefits outweigh the added communication, allocation, and operating cost.

References

• Google Cloud documentation, TPU v4 (two TensorCores per chip, 32 GiB HBM, 275 teraflops bf16, 4,096-chip pod, 1.1 exaflops). cloud.google.com
• Google Cloud documentation, TPU v3 (two TensorCores per chip, 1,024-chip pod, 126 petaflops bf16). cloud.google.com
• Jouppi, N. P., et al. (2023). TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings. ISCA 2023, arXiv:2304.01433. arxiv.org