Google releases the seventh-generation Ironwood TPU Developer Training Guide, detailing system-level performance optimization

ME News Report, April 2 (UTC+8), Google officially released a developer training guide for the seventh-generation Ironwood TPU. The guide is designed to help developers fully leverage the Ironwood TPU’s system-level performance for efficient training and deployment of cutting-edge AI models. The Ironwood TPU is a customized AI infrastructure built to meet the compute demands of trillion-parameter models. It constructs a complete system supporting up to 9,216 chips through technologies such as inter-chip connectivity (ICI), optical circuit switches (OCS), data center networks (DCN), and high-bandwidth memory (HBM). The article details multiple key optimization strategies for this hardware, including: using its matrix multiplication unit (MXU), which natively supports FP8 training to improve throughput; adopting the TPU-optimized JAX kernel library Tokamax, which uses “Spill Attention” and “Megablox Grouped Matrix Multiplication” to handle irregular tensors in long-context and mixture-of-experts models; offloading collective communication operations to the fourth-generation SparseCore to hide latency; carefully fine-tuning the allocation of TPU’s on-chip SRAM (VMEM) to reduce memory stalls; and selecting the optimal sharding strategy (such as FSDP, TP, EP) according to model scale, architecture, and sequence length. (Source: InFoQ)