谷歌发布第七代Ironwood TPU开发者训练指南，详解系统级性能优化

ME News 消息，4 月 2 日（UTC+8），谷歌官方近日发布了面向第七代Ironwood TPU的开发者训练指南。该指南旨在帮助开发者充分利用Ironwood TPU的系统级性能，以高效训练和部署前沿AI模型。Ironwood TPU是为满足万亿参数模型算力需求设计的定制化AI基础设施，其通过芯片间互连（ICI）、光路交换机（OCS）、数据中心网络（DCN）及聚合高带宽内存（HBM）等技术，构建了支持多达9,216芯片的完整系统。文中详细介绍了针对该硬件的多项关键优化策略，包括：利用其矩阵乘法单元（MXU）原生支持FP8训练以提升吞吐量；采用专为TPU优化的JAX内核库Tokamax，通过"飞溅注意力"和"Megablox分组矩阵乘法"处理长上下文和混合专家模型中的不规则张量；利用第四代稀疏核心（SparseCore）卸载集体通信操作以隐藏延迟；精细调优TPU快速片上SRAM（VMEM）的分配以减少内存停滞；以及根据模型规模、架构和序列长度选择最佳分片策略（如FSDP、TP、EP）。（来源：InFoQ）