TPU (Tensor Processing Unit) Deep Dive
a year ago
- #AI Hardware
- #TPU
- TPU是谷歌研发的专用集成电路芯片，专为极致矩阵乘法吞吐量和能效比而设计
- TPU技术萌芽于2006年，2013年起因神经网络计算需求激增进入高速发展阶段
- 谷歌绝大多数AI服务（包括Gemini和Veo等模型的训练与推理）均由TPU提供算力支持
- 单个TPUv4芯片集成两个张量核心，共享CMEM和HBM两种内存单元
- TPU采用脉动阵列实现高效矩阵运算和卷积计算，但在稀疏矩阵处理上存在局限
- TPU依赖XLA编译器进行提前编译优化，以此实现内存访问和能耗效率的最大化
- TPU架构设计核心在于减少内存操作，从而实现能耗降低与性能提升的双重目标
- TPU具备弹性扩展能力，支持从单芯片到数千芯片组成的多pod系统等多种配置
- TPU机柜采用3D环面拓扑结构，配合光路交换技术(OCS)实现高效灵活的内部通信
- TPU计算单元可配置为立方体/雪茄形等多种拓扑结构，以适应不同并行计算需求
- 多pod级TPU系统通过数据中心网络(DCN)实现跨pod通信，支撑PaLM等超大规模模型训练