Google第四代TPU芯片发布


在9天前的2021 Google I/O大会上,Google CEO Sundar Pichai正式发布了第四代TPU(Tensor Processing Unit)芯片。消息一出,各大科技媒体纷纷报道。不过让人感到遗憾的是,Pichai只用了4页PPT来介绍TPUv4。除了展示TPUv4板卡和介绍其带来的性能提升之外,对各种设计细节只字未提。当然,这非常符合Google一贯的风格。对于数据中心硬件设备,Google一直将其作为差异化竞争的手段,因此有着非常高的保密等级。不过,从Google的只言片语和过去在不同场合留下的零星信息,我们还是能够大致地完成这幅拼图。接下来,就让我们睁大眼睛,开启这次的侦探之旅。

在户外举行的2021 Google I/O大会上,Pichai用下面的这张大图向业界公布了TPUv4板卡的真容。和上一代TPUv3相比,一张TPU板卡上仍然放置了4个TPU芯片。从复用TPUv3液冷机箱和提升芯片散热能力两个方面出发,TPUv4板卡继续使用冷板式液冷方案。在TPUv4板卡的右侧,取消了在TPUv3板卡上使用的ExaMAX高速连接器。这说明TPUv4板卡之间的互连将完全依靠左侧的IO端口,在服务器机箱内部将不再有与板卡相连的小背板。

Google I/O大会上对TPUv4的介绍图

TPUv4板卡左侧的IO端口从上一代的OSFP+CDFP的组合简化为4个CDFP,每个CDFP端口对外提供400G网络带宽(16×25Gbps)。这4个CDFP对应4个TPUv4芯片,表明每个TPUv4芯片上都具有网络接口。Pichai在演讲中也提到了TPUv4拥有更大的网络带宽,这可以让每个TPUv4芯片更快地从外部内存中读取数据。

400G CDFP CageCable


在这里,我们附上Google的TPUv2和TPUv3产品的板卡和集群Pod的图片和性能数据,对于不熟悉上两代产品的读者,可以通过反复的对比,来理解前面所谈到的TPUv4板卡上的更改点,以及后面要做的进一步分析。

Google TPUv2TPUv3产品


Pichai给出的第二张PPT是关于TPUv4 Pod,在最新的TPU计算集群(Pod)里总共使用了4096块TPUv4芯片。由于一块TPUv4板卡上放置了4块TPUv4芯片,因此整个Pod里总共有1024张TPUv4板卡,这是TPUv3 Pod里TPU板卡数量的4倍。

使用4096个TPU芯片的TPUv4 Pod

TPUv3 Pod将机柜数量从TPUv2 Pod的4个扩展到8个,在每个机柜的槽位上放置了2张TPUv3板卡,从而容纳了1024块TPUv3芯片。现在TPUv4 Pod可以容纳的芯片数量增加了4倍,在机柜高度不变的情况下(机柜高度变更涉及到整个数据中心机房的基础设施,牵一发而动全身,因此不会轻易更改),只能通过增加机柜每个槽位可以放置的TPU板卡数量,并增加集群里机柜的数量来实现。因此,估计Google在每个槽位上放置了4块TPUv4板卡,并将集群机柜数量从上一代的8个扩展到16个,来实现拥有4096个TPUv4芯片的新的集群。

TPUv2和TPUv3 Pod与TPUv4 Pod参数对比

扩展TPU Pod集群的规模也很容易理解。在经历了三代产品的迭代之后,TPU Pod不仅在Google内部使用,还通过Google Cloud的方式开放给全球各地的AI开发人员。只要是使用Tensor Flow AI框架开发的AI模型,都可以很容易地在TPU Pod上进行训练和推理。

Pichai展示的第三张PPT宣称TPUv4 Pod具有ExaFlops(1e18)的峰值浮点计算性能。要理解这一点,需要回到2020年11月MLPerf Training 0.7上发布的TPUv4 Pod和TPUv3 Pod对比数据。

具有ExaFlops计算性能的TPUv4 Pod

从MLPerf发布的对比结果中可以看到,针对不同的AI模型,TPUv4 Pod的训练效率是TPUv3 Pod训练效率的2.2~3.7倍。做一个简单的平均计算,得到平均效率提升为2.7倍。训练效率的提升来源于TPUv4芯片浮点计算能力的提升,考虑到TPUv3板卡的峰值计算能力是420TFlops,因此TPUv4板卡的峰值计算能力为1134TFlops左右(420TFlops × 2.7)。对于具有1024块TPUv4板卡的TPUv4 Pod而言,其峰值算力将达到1.16ExaFlops左右(1134TFlops × 1024)。正是基于这一BF16浮点数据的计算性能,Pichai才宣称TPUv4 Pod是全球首个计算能力达到E级别的AI计算集群。

TPUv4 PodTPUv3 Pod MLPerf Training对比结果

这里需要澄清一下,TPUv4 Pod更多的是针对AI模型进行加速计算。常规的E级别超算需要处理更多样的计算负载,不仅要进行BF16浮点数计算,还需要进行更高精度的FP64浮点数计算(如天气预报、地质勘探、航空航天等)。因此,真正意义上的E级别超算时代的到来还需要等待一段时间。

对于TPUv4芯片的信息挖掘是否到这里就结束了呢?非也!从TPUv4 Pod的算力能力上,我们还可以推理出更多的信息。

单个TPUv4芯片的BF16浮点算力显然比TPUv3芯片提升了2倍以上,这2倍多的性能提升是从哪里来的呢?首先,Google必然采用了更新的芯片工艺制程。TPUv3芯片使用的是16/12nm工艺,而TPUv4芯片肯定使用了目前已经非常成熟的7nm工艺制程。通过晶体管的微缩,Google才能在差不多大小的芯片面积上将矩阵乘法计算单元(Matrix Multiply Unit)增加一倍。在从TPUv2升级到TPUv3的时候,Google也是通过加倍MMU模块来实现性能提升的。

Google TPUv3芯片框图

TPUv3芯片上单个TPU Core只带有16GB HBM内存,在增加了矩阵乘法单元之后,为了在芯片运行时能够有足够的数据来使用这些矩阵乘法单元,HBM内存的容量也需要进行翻倍,从现在的16GB增加到32GB。最合理的设计方式是在TPU芯片里直接增加TPU Core,将TPU Core的数量从2个增加到4个。这样一来,TPUv4芯片的每核心内存容量(memory/core)就保持了不变。得益于工艺制程的缩减,增加的这部分电路不会使得芯片功耗出现大幅增加,估计仍然保持在400W~500W之间。唯一不确定的是Google仍然在使用单一大Die的芯片设计方式,还是已经开始使用Chiplet技术了。如果使用了Chiplet技术,TPU Core间的互连技术将是值得关注的一个点。

总 结

基于已有信息并进行一番推理分析,TPUv4芯片、TPUv4板卡和TPUv4 Pod的基本情况已经跃然纸上。剩下唯一需要做的,就是等待Google在后续的Hot Chip、ISSCC等全球性的芯片大会上一点点地确认这些信息。

部署在Google数据中心的TPUv4 Pod

TPUv4 Pod目前已经部署在了Google的数据中心里,跟上一代产品一样,外部用户只能通过Google Cloud的方式来访问新一代的TPU芯片。基于TPUv4 Pod,AI模型开发者将会拥有更为强大的算力,从而更快速地进行AI模型调优和训练。

下一篇
« Prev Post
上一篇
Next Post »