Model development best practices
TPUs are powerful accelerators, but getting the most out of them takes a few design principles. The goal is to keep the TPU compute-bound - spending most of its time doing math on-chip, not waiting for data to move or for the graph to recompile.
The MXU is the point
At the heart of every Cloud TPU is the Matrix Multiplier Unit (MXU), a systolic array built to do large matrix operations extremely efficiently - a super-fast calculator for the math that powers deep learning.
If your program is dominated by non-matrix work - element-wise additions, reshaping tensors, concatenating them - you are not using the MXU's power. It is like owning a sports car and only driving it in stop-and-go traffic. Prioritize core computations that are matrix multiplications.
Principle 1: Layout for efficiency
The XLA compiler transforms your code to run on the MXU, including a crucial step called tiling - breaking large matrix multiplications into smaller, manageable blocks. It leans on fixed hardware characteristics to tile optimally:
- The MXU is a 256x256 systolic array (TPUs before Trillium used 128x128).
- The memory subsystem prefers tensor dimensions that are multiples of 8.
- Conducive layouts - certain tensor data arrangements - tile naturally.
- Other layouts force reshape operations before tiling is possible.
Reshape operations demand more memory access than actual processing, so they slow computations down. Design your model to minimize complex reshaping, especially right before core matrix multiplications.
Principle 2: Fixed shapes for predictable performance
XLA compiles your ML graph just in time (JIT) for the first batch of data, and that compilation is heavily optimized.
If later batches have different shapes than the first, the model will not run correctly - recompiling the whole graph on every shape change is far too slow for the speeds TPUs target. Any model with dynamic tensor shapes (a tensor whose size changes step to step) is generally not well-suited to TPUs. Use fixed tensor shapes for consistent execution.
Principle 3: Avoiding unnecessary padding
A high-performing TPU program is one whose dense computations break cleanly into 128x128 chunks that perfectly fill the MXU. When a computation cannot fill an entire MXU (for example, a dimension that is not a multiple of 128), XLA automatically pads tensors with zeros.
- Underutilization - computing on "empty" zero data wastes processing power that could be doing real work.
- Increased memory - padding grows the on-chip storage a tensor needs; in extreme cases it can trigger an out-of-memory error.
Padding is managed at three levels:
| Level | How to handle padding |
|---|---|
| Automatic | XLA handles padding for you - it inserts the zeros so the computation can proceed. |
| Detection | Use the op_profile tool to measure exactly how much padding your model is performing. |
| Prevention | Pick TPU-friendly tensor dimensions - multiples of 8, and ideally multiples of 128 for matrix operations. |
Recap
- Prioritize matrix operations - keep core computations on matrix multiplications so the MXU stays busy.
- Stay compute-bound - the TPU should spend its time calculating, not waiting on memory or I/O.
- Smart dimensions - choose sizes that are multiples of 8, ideally 128 for matrix computations, to minimize padding.
- Optimize layouts - minimize costly reshape operations, especially before major computations.
- Fixed tensor shapes - avoid dynamic shapes; keep dimensions consistent across batches.
- Avoid costly recompilations - caused by frequent shape changes, control statements that depend on tensor values, or premature host (CPU) data access.