Google TPU v4 + v5 (v5e + v5p), in depth: megacore, the optically switched torus, SparseCores, and moving v4 jobs to v5p and v5e

The v4 chip as software sees it

A v4 chip contains two TensorCores. Each TensorCore has four matrix-multiply units (MXUs), a vector unit and a scalar unit. The chip carries 32 GiB of HBM2 at about 1,200 GB/s and peaks at 275 teraflops in bf16 or int8. Four chips share one host VM, so the smallest v4 slice, v4-8, is one host with four chips.

That name is the first trap. v4 accelerator types count TensorCores, not chips: v4-8 is four chips, v4-128 is 64 chips on 16 hosts. v5p keeps the same convention, while v5e names count chips (v5litepod-8 is eight chips). Any capacity script that divides a name by a constant needs to know which generation it is looking at.

The second idea is megacore. On v4 (and again on v5p) the two TensorCores on a chip are presented to the framework as one logical device with one HBM address space; XLA splits each operation across the two cores. In JAX this means jax.device_count() on a v4-8 reports 4, not 8. Your mesh is laid out over chips, and per-device memory is the full 32 GiB, which is why a v4 chip feels like one large accelerator rather than two small ones.

Optical circuit switches and the reconfigurable torus

The v4 pod holds 4,096 chips, but the job never sees a fixed 4,096-chip machine. Chips are built into cubes of 4x4x4 (64 chips). Links inside a cube are electrical. Every cube face connects to a bank of optical circuit switches (the paper describes 48 of them), which are arrays of steerable mirrors: they do not inspect packets, they physically point one fibre at another. When the scheduler allocates a slice it picks free cubes anywhere in the pod and programs the switches to wire their faces together, closing the wraparound links that turn a mesh into a torus.

Three operational consequences follow. First, availability: a cube with a failed host can be left out and a spare patched in, so a large slice does not need a contiguous, perfect block of hardware. Second, shape: Google documents full 3D torus shapes such as 4x4x8, 4x8x8 and 12x12x24, plus optional twisted torus variants that skew the wraparound links to raise bisection bandwidth for all-to-all-heavy traffic. Third, cost model: the paper reports the optical switching at under 5% of system cost and under 3% of system power, which is why Google kept the approach for v5p.

For software, the topology determines what collectives cost. An all-reduce over one torus axis uses the links in both directions around a ring; a mesh axis laid out along a short dimension finishes sooner than one stretched across a long dimension. That is why JAX meshes on v4 are built from the physical device grid rather than a flat list.

SparseCore and embedding-heavy models

Recommendation models spend much of their step in embedding lookups: gather a few rows from tables that can be hundreds of gigabytes, sum or average them, and later scatter gradients back. That work is memory-bound and irregular, which is exactly what a systolic array is bad at. v4 added SparseCores, small dataflow processors beside the TensorCores, dedicated to these gathers and scatters over embeddings sharded across the whole slice. The ISCA paper credits them with a 5x to 7x speed-up on embedding-heavy models for about 5% of die area and power.

You do not program SparseCores directly. In TensorFlow the entry point is the TPU embedding API (tf.tpu.experimental.embedding.TPUEmbedding), which places tables and routes lookups to them; other frameworks reach them through their own embedding libraries, and support varies by release, so check your framework's current documentation. v5p's specification lists four SparseCores per chip. The v5e page lists none, which matters if you planned to move a large recommender to the cheaper chip.

What changed in v5e and v5p

v5e and v5p split v4's design in two directions. v5e cuts the chip down for cost per operation: one TensorCore, 16 GB of HBM at roughly 0.8 TB/s, 197 bf16 teraflops and 393 int8 tera-ops, eight chips per host, and a 2D torus of at most 256 chips. v5p scales it up: two TensorCores and four SparseCores per chip, 95 GiB of HBM at 2,765 GB/s, 459 bf16 teraflops, four chips per host, a 3D torus, and an 8,960-chip pod whose largest single training slice Google lists as 6,144 chips.

Read the table as a set of migration questions rather than a ranking. Going to v5p, per-chip memory triples, so the sharding you chose on v4 to fit optimizer state is probably over-sharded. Going to v5e, per-chip memory halves, the torus loses a dimension, and there is no SparseCore, so a v4 layout may not fit at all.

Worked example: moving a v4 job to v5p and v5e

Suppose a team trains a 7-billion-parameter decoder on v4-128 (64 chips, 16 hosts) with fully sharded data parallelism, and serves it on v4 as well. They want training on v5p and serving on v5e.

Training to v5p-128. Because both generations count TensorCores, v5p-128 is also 64 chips on 16 hosts, so the process count and per-host data loader do not change. Mixed-precision Adam holds about 16 bytes per parameter (bf16 weights and gradients, fp32 master copy and two moments), about 112 GB here. Sharded over 64 chips that is under 2 GB per chip on either generation, so memory is not the constraint; the freed HBM on v5p is better spent on a larger per-chip batch or less activation rematerialisation. The code change is to stop hard-coding the mesh:

import jax
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
from jax.experimental import mesh_utils

n = jax.device_count()            # 64 on v4-128 and on v5p-128 (megacore: one device per chip)
model_axis = 4                     # tensor-parallel width; keep it inside a host
devices = mesh_utils.create_device_mesh((n // model_axis, model_axis))  # follows the physical torus
mesh = Mesh(devices, axis_names=("data", "model"))

param_spec = NamedSharding(mesh, P("data", "model"))   # FSDP over data, TP over model
batch_spec = NamedSharding(mesh, P("data", None))
print(jax.devices()[0].device_kind, n, mesh.shape)

Checkpoints written on v4 restore on v5p as long as the restore call is given the new target shardings, which lets the checkpoint library reshard on load instead of assuming the saved layout. Expect a recompilation on the first step: XLA executables are specific to the chip, so the compilation cache from v4 is useless and should be keyed by device_kind.

Serving to v5e. The bf16 weights alone are 14 GB against 16 GB of HBM per chip, which leaves almost nothing for the KV cache. Two workable options: shard the model four ways over a v5litepod-4 so each chip holds about 3.5 GB of weights and the rest is cache, or quantise weights to int8 (about 7 GB) and use v5e's int8 throughput. Evaluate the quantised model's quality before committing. Provisioning uses the runtime images Google lists per generation:

# v4 training slice (TensorCore-counted name, Ubuntu base runtime)
gcloud compute tpus tpu-vm create train-v4 --zone=$ZONE \
  --accelerator-type=v4-128 --version=tpu-ubuntu2204-base

# v5p replacement: same name size, same 64 chips
gcloud compute tpus tpu-vm create train-v5p --zone=$ZONE \
  --accelerator-type=v5p-128 --version=v2-alpha-tpuv5

# v5e serving slice: chip-counted name
gcloud compute tpus tpu-vm create serve-v5e --zone=$ZONE \
  --accelerator-type=v5litepod-4 --version=v2-alpha-tpuv5-lite

One training step and where it stalls

Whichever generation you run, the step has the same shape, and knowing it tells you where to look when throughput drops. Each host reads its shard of the input and places it on its local chips. XLA executes the compiled step: all-gathers of parameter shards along the data axis, matrix multiplies on the MXUs, reduce-scatters of gradients back along the same axis, and an optimizer update on the local shard. On v4 and v5p, embedding lookups run concurrently on SparseCores. Nothing leaves the slice except checkpoint writes and logs.

The bottleneck moves with the generation. On v5e, the 2D torus and smaller HBM make collectives and memory traffic proportionally more expensive, so tensor-parallel axes should stay inside a host. On v5p, faster compute exposes the input pipeline: a loader that kept up with v4 can starve v5p, and the profiler shows it as idle time between steps rather than slow steps.

Failure modes

• Capacity math off by 2x. A script that treats v4-128 and v5litepod-128 as the same size requests half or double the chips intended. Normalise everything to chips.
• Recompile storms. Shapes that vary per batch trigger fresh compilations on any TPU; after a migration the cold cache makes this visible all at once. Pad to a few fixed buckets.
• Embedding models stall on v5e. Without SparseCores, large-table lookups fall back to TensorCore gathers and HBM traffic. Benchmark before moving a recommender.
• Slice unavailable. Large slices depend on free cubes; queue for capacity instead of retrying creates in a loop, and checkpoint often enough that preemption or maintenance costs minutes.
• Mesh crosses the slow axis. A tensor-parallel axis that spans hosts on v5e turns every layer into a cross-host collective. Keep it within the eight chips of a host.
• Quantisation regressions. int8 serving that looks fine on perplexity can fail on long-context or code tasks; gate it on task evaluations.

Trade-offs

Prices change by region and commitment, so compare current published rates per chip-hour against measured step time on each candidate; cost per trained token or per served request is the number that decides.

What to do next

1. Convert every accelerator name in your scripts to a chip count, noting which generations count TensorCores.
2. Replace hard-coded device meshes with mesh_utils.create_device_mesh built from jax.device_count().
3. Restore one v4 checkpoint on the target generation with explicit target shardings and compare loss on a fixed batch.
4. Profile ten steps on the new slice and check for host-side idle time before tuning anything on the chip.
5. For embedding-heavy models, benchmark on v5p and v5e separately; do not assume SparseCore behaviour carries over.
6. Read how the systolic array works and fully sharded data parallelism before redesigning the sharding plan.