简单谈谈Google TPUv6

简单谈谈Google TPUv6

根据Google TPU第六代的数据做了一些性能数据的对比，需要注意的是TPUv6当前应该是一个用于训推一体的单Die的版本,用于训练的V6p双Die版本应该会后期再发布. 需要注意的是在国内外都开始卷大模型推理价格的时候, TPU这样的东西对于提高ROI非常有帮助。快速浏览手Google TPU版本，见表2-21。

表2-21 快速浏览手Google TPU版本

	v4	v5e	v5p	v6e
每个pod的芯片数	4096	256	8960	256
chip bf16 TFLOPS	275	197	459	926
HBM(GB)	32	16	95	32
HBM BW(GB/s)	1228	820	2765	1640
每个芯片的ICI BW (Gb/s)	2400	1600	4800	3200

TPUv6的主要的几个优化点是MXU更大频率更高，浮点能力基本上到H100，HBM带宽和容量翻倍换了HBM3，ICI带宽也升级了。TPU互联拓扑应该和TPUv5e是一致的，如图2-38所示。

 

图2-38 ve5中的TPU主机框架

TPUv5e的物理结构如下，四片一块板上互联，剩余的ICI拉出到OCS光交换机上，如图2-39所示。

图2-39 TPUv5e的物理结构，四片一块板上互联，剩余的ICI拉出到OCS光交换机上

一个比较关键的点是SparseCore的升级，老的SparseCore架构，如图2-40所示。

 

图2-40 稀疏核硬件结构

与GB200直接C2C连接一个Grace不同的是，Google采用了在TPU内置标量核的方法，在相应的性瓦比上会比NV更有优势。TPU ve5有效缩放因子，如图2-41所示。

 

图2-41 TPU ve5有效缩放因子

另外Google这套东西很值得借鉴的是它在TPUv5e开始，e系列负责训推一体，Pod规模不会很大,部署更加灵活，而p系列则负责更大规模的Foundation Model训练任务.。Pathways的调度机制也维持了很好的线性加速比，例如TPUv5e可以做到50000卡基本线性加速。