The hardware that trains Gemini
Tensor Processing Units (TPUs) are Google's custom ASICs designed specifically for neural network matrix multiplication. They're purpose-built for the dominant operation in deep learning, and they do it extraordinarily fast.
| Hardware | Peak TFLOPS (bfloat16) | Memory bandwidth | Used for |
|---|---|---|---|
| NVIDIA A100 | ~312 | 2 TB/s | GPT-scale research |
| Google TPU v4 (per chip) | ~275 | 1.2 TB/s | Gemini/PaLM training |
| TPU v4 Pod (4096 chips) | ~1.1 exaFLOPS | ~4.9 PB/s | LLM pretraining |
Key differences from GPUs you must know:
- No eager execution — all compute must be inside
@tf.functionto be traced into a graph.model.fithandles this automatically. - bfloat16 native — same exponent range as float32, fewer mantissa bits. Set
mixed_precision.set_global_policy('mixed_bfloat16'). - XLA compilation required — all ops run through XLA, which fuses operations for efficiency. Some TF ops without XLA implementations fall back to CPU.
- Large batch sizes — TPUs are most efficient with batch sizes of 1024–8192 per chip.
Free TPU access is available in Google Colab and Kaggle notebooks — great for learning without a cloud bill.