谷歌超强AI超算碾压英伟达A100！TPU v4性能提升10倍，细节首次公开

虽然谷歌早在2020年，就在自家的数据中心上部署了当时最强的AI芯片——TPU v4。

但直到今年的4月4日，谷歌才首次公布了这台AI超算的技术细节。

论文地址：https://arxiv.org/abs/2304.01433

相比于TPU v3，TPU v4的性能要高出2.1倍，而在整合4096个芯片之后，超算的性能更是提升了10倍。

另外，谷歌还声称，自家芯片要比英伟达A100更快、更节能。

与A100对打，速度快1.7倍

论文中，谷歌表示，对于规模相当的系统，TPU v4可以提供比英伟达A100强1.7倍的性能，同时在能效上也能提高1.9倍。

另外，谷歌超算速度还要比Graphcore IPU Bow快约4.3倍至4.5倍。

谷歌展示了TPU v4的封装，以及4个安装在电路板上的封装。

与TPU v3一样，每个TPU v4包含两个TensorCore（TC）。每个TC包含四个128×128矩阵乘法单元（MXU），一个具有128个通道（每个通道16个ALU），以及16 MiB向量存储器（VMEM）的向量处理单元（VPU）。

两个TC共享一个128 MiB的公共存储器（CMEM）。

值得注意的是，A100芯片与谷歌第四代TPU同时上市，那么其具体性能对比如何？

谷歌分别展示了在5个MLPerf基准测试中每个DSA的最快性能。其中包括BERT、ResNET、DLRM、RetinaNet、MaskRCNN。

其中，Graphcore IPU在BERT和ResNET提交了结果。

如下展示了两个系统在ResNet和BERT的结果，点之间的虚线是基于芯片数量的插值。

TPU v4和A100的MLPerf结果都扩展到比IPU更大的系统（4096个芯片对比256个芯片）。

对于相似规模的系统，TPU v4在BERT上比A100快1.15倍，比IPU快大约4.3倍。对于ResNet，TPU v4分别快1.67倍和大约4.5倍。

对于在MLPerf基准测试上的功耗使用情况，A100平均上使用了1.3倍至1.9倍的功率。

峰值每秒浮点运算次数是否能预测实际性能？许多机器学习领域的人认为峰值每秒浮点运算次数是一个很好的性能代理指标，但实际上并非如此。

例如，尽管在峰值每秒浮点运算次数上仅具有1.10倍的优势，TPU v4在两个MLPerf基准测试上比IPU Bow在相同规模的系统上快4.3倍至4.5倍。

另一个例子是，A100的峰值每秒浮点运算次数是TPU v4的1.13倍，但对于相同数量的芯片，TPU v4却快1.15倍至1.67倍。

如下如图使用Roofline模型展示了峰值FLOPS/秒与内存带宽之间的关系。

那么，问题来了，谷歌为什么不和英伟达最新的H100比较？

谷歌表示，由于H100是在谷歌芯片推出后使用更新技术制造的，所以没有将其第四代产品与英伟达当前的旗舰H100芯片进行比较。

不过谷歌暗示，它正在研发一款与Nvidia H100竞争的新TPU，但没有提供详细信息。谷歌研究员Jouppi在接受路透社采访时表示，谷歌拥有「未来芯片的生产线」。

TPU vs GPU

在ChatGPT和Bard「决一死战」的同时，两个庞然大物也在幕后努力运行，以保持它们的运行——英伟达CUDA支持的GPU（图形处理单元）和谷歌定制的TPU（张量处理单元）。

换句话说，这已经不再是关于ChatGPT与Bard的对抗，而是TPU与GPU之间的对决，以及它们如何有效地进行矩阵乘法。

由于在硬件架构方面的出色设计，英伟达的GPU非常适合矩阵乘法任务——能有效地在多个CUDA核心之间实现并行处理。

因此从2012年开始，在GPU上训练模型便成为了深度学习领域的共识，至今都未曾改变。

而随着NVIDIA DGX的推出，英伟达能够为几乎所有的AI任务提供一站式硬件和软件解决方案，这是竞争对手由于缺乏知识产权而无法提供的。

相比之下，谷歌则在2016年推出了第一代张量处理单元（TPU），其中不仅包含了专门为张量计算优化的定制ASIC（专用集成电路），并且还针对自家的TensorFlow框架进行了优化。而这也让TPU在矩阵乘法之外的其他AI计算任务中具有优势，甚至还可以加速微调和推理任务。

此外，谷歌DeepMind的研究人员还找到了一种能够创造出更好矩阵乘法算法的方法——AlphaTensor。

然而，即便谷歌通过自研的技术和新兴的AI计算优化方法取得了良好的成果，但微软与英伟达长久以来的深度合作，则通过利用各自在行业上的积累，同时扩大了双方的竞争优势。

第四代TPU

时间回到21年的谷歌I/O大会上，劈柴首次公布了谷歌最新一代AI芯片TPU v4。

「这是我们在谷歌上部署的最快的系统，对我们来说是一个具有历史意义的里程碑。」

这次的改进已经成为构建AI超算的公司之间竞争的关键点，因为像谷歌的Bard、或OpenAI的ChatGPT类似的大型语言模型已经在参数规模上实现爆炸式增长。

这意味着它们远远大于单个芯片所能存储的容量，对算力需求是一个巨大的「黑洞」。

因此这些大模型必须分布在数千个芯片上，然后这些芯片必须协同工作数周，甚至更长时间来训练模型。

目前，谷歌迄今为止公开披露的最大的语言模型PaLM，有5400亿参数，便是在50天内将其分割到两台4000芯片的超级计算机上进行训练的。

谷歌表示，自家的超级计算机能够轻松地重新配置芯片之间的连接，能够避免问题，并进行性能调优。

谷歌研究员Norm Jouppi和谷歌杰出工程师David Patterson在关于该系统的博客文章中写道，

「电路交换使得绕过失效组件变得容易。这种灵活性甚至允许我们改变超算互连的拓扑结构，以加速机器学习模型的性能。」

尽管谷歌现在才发布有关其超级计算机的详细信息，但自2020年以来，该超级计算机已在位于俄克拉荷马州梅斯县的数据中心内上线。

谷歌表示，Midjourney使用该系统训练了其模型，最新版的V5让所有人见识到图像生成的惊艳。

最近，劈柴在接受纽约时报采访称，Bard将从LaMDA转到PaLM上。

现在有了TPU v4超算的加持，Bard只会变得更强。