id
flops/alt-tpu-spmd-partitioning
Arithmetic intensity
TPUs use SPMD (Single Program, Multiple Data) where the same program runs on all cores. The XLA compiler automatically partitions tensors across cores using layout constraints. This means the model architecture must be amenable to automatic sharding — manual control is limited compared to GPU tensor parallelism. The compiler inserts collectives (all-reduce, all-gather) automatically, but the sharding strategy is constrained by the SPMD model.
Scales with
num_cores (more cores = more sharding)tensor_dimensions (must be divisible by core count for even sharding)model_architecture (some architectures shard cleanly, others do not)ici_bandwidth (inter-chip interconnect bandwidth bounds collective cost)
Hardware parameters
parametermemory_bandwidtheffectnum_cores (more cores = more sharding) | tensor_dimensions (must be divisible by core count for even sharding) | model_architecture (some architectures shard cleanly, others do not) | ici_bandwidth (inter-chip interconnect bandwidth bounds collective cost)basisunverifiedvalues_forgoogle-tpu-v5e,google-tpu-v5p,google-tpu-v6e
Workarounds
Use FSDP (Fully Sharded Data Parallel) as default — preferred by Google for TPU7xUse Expert Parallelism for MoE models (all-to-all routing)Use Context Parallelism for long sequences (shards sequence dimension)Minimize tensor parallelism — high arithmetic intensity (11.5k on TPU7x) makes TP communication-boundDesign model dimensions to be divisible by core count
notes
The SPMD model is both a constraint and a simplification. Unlike GPU inference where the user manually specifies tensor parallelism degree and placement, TPU's XLA compiler handles sharding automatically. The Ironwood performance guide recommends: 'select the simplest sharding strategy that meets memory constraints, as this minimizes communication overhead.' For MoE, Expert Parallelism is necessary. For long contexts, Context Parallelism is essential. The key insight is that SPMD means the compiler controls layout, so the model architect must work with the compiler rather than against it.
HARDWARE-PARAMETER ROWS ADDED 2026-10-05 (schema authority pass, agent w5-schema). `affected_by_hardware` names PRODUCTS (accelerator record ids); these rows name the hardware QUANTITIES that move this class's cost, which had nowhere to live until flop.hardware_parameters was added in this commit. Every effect below is lifted verbatim from this record's own existing prose and no figure here is new - `magnitude` reproduces numbers the record already quoted. `basis` is 'unverified' where the effect was asserted with nothing concrete behind it, rather than being dressed up. The rows were pulled out of free-text notes; they are not a change to what this record claims.