Google谷歌第四代定制AI芯片TPU v4 Pods整合算力及性能详解

shuszhao · 发表于 2021-5-22 12:57:19

Google今天正式发布了其第四代人工智能TPU v4 AI芯片，其速度达到了TPU v3的2.7被。Google实际上已经于2020年就开始在自己的数据中心中使用了新的TPU v4。通过整合4096个TPU v4芯片成一个TPU v4 Pod，一个Pod性能就达到世界第一超算“富岳”的两倍。这些算力可能在今年晚些时候向Google Cloud用户开放此功能。且谷歌希望未来可能应用于量子计算。

Google TPU系列研发

Google于2016年宣布了其首款内部定制的定制AI芯片。

与传统的CPU和GPU组合（用于训练和部署AI模型的最常见架构）相比，第一款TPU ASIC提供了“数量级”更高的性能。
TPU V2于2018年上市，TPU V3于2019年上市，TPU v4于2020年推出。
借助TPU v3，这些服务器的机架需要大量的电源，以至于Google必须对其数据中心进行改造以支持液体冷却，这比传统的风冷系统具有更高的功率密度。
而今，在比TPU v3性能提高2.7倍的基础上再整合4096个TPU v4为一个TPU v4 Pod，许多Pod再组成TPU v4 Pods，以提供超强的算力。
除了将这些系统用于自己的AI应用程序（例如搜索建议，语言翻译或语音助手）外，Google还将TPU基础设施（包括整个TPU吊舱）以云的方式（付费）开放给Google Cloud客户。

1 TPU v4 Pod 整合算力 = 2 富岳

在其I/O开发者大会上，Google今天（美国时间5月18日）宣布了其下一代定制的张量处理单元（TPU）人工智能芯片。这是TPU芯片的第四代产品，Google称其速度是上一版本的两倍。正如Google首席执行官桑达尔·皮查伊所指出的，这些芯片整合了4096个TPU v4，一个pod就可以提供超过一个exaflop（每秒百亿亿次浮点运算）的AI计算能力，达到目前全球最快的超级计算机“富岳”的两倍性能。
Google使用定制芯片为其自己的许多机器学习服务提供动力，但与此同时它也将把这最新一代产品作为其Google云平台的一部分提供给开发者。
"这是我们在Google部署过的最快的系统，对我们来说是历史性的里程碑，"Google CEO桑达尔·皮查伊说。"以前要想获得1个exaflop的算力，通常需要建立一个定制的超级计算机，但我们今天已经部署了许多这样的计算机，很快就会在我们的数据中心有几十个TPUv4 pods，其中许多将以90%或接近90%的无碳能源运行。而我们的TPUv4 pods将在今年晚些时候提供给我们的云客户"。

TPU是Google的第一批定制芯片之一，当包括微软在内的其他公司决定为其机器学习服务采用更灵活的FPGA时，Google很早就在这些定制芯片上下了赌注。虽然它们的开发时间稍长，而且随着技术的变化很快就会过时，但用于特定场景的运算时可以提供明显更好的性能。

发布会上，Google还展示了其量子数据中心，尽管目前相关研究仍停留于早期阶段，但谷歌显然希望在未来有朝一日梦想成真。基于目前的材料和实验场所限制，量子计算必须在极低的超导温度下运行。

或用于量子计算

发布会上，Google还展示了其量子数据中心，尽管目前相关研究仍停留于早期阶段，但谷歌显然希望在未来有朝一日梦想成真。基于目前的材料和实验场所限制，量子计算必须在极低的超导温度下运行。

Google TPU v4 性能详解：打破MLPerf中的AI性能记录

机器学习（ML）模型的快速训练对于提供以前无法实现的新产品，服务和研究突破的研究和工程团队至关重要。在Google，最近启用ML的功能包括更有用的搜索结果和一个可以翻译100种不同语言的ML模型。
行业标准MLPerf基准测试对比的最新结果表明，谷歌已打造出世界上最快的ML训练超级计算机。Google使用此超级计算机以及最新的Tensor处理单元（TPU）芯片，在八项MLPerf基准测试中的六项中创造了性能记录。

图1：与任何可用性类别中最快的非Google提交相比，Google最好的MLPerf Training v0.7研究提交速度有所提高。不论系统大小如何，总训练时间都将比较标准化，而系统大小介于8到4096个芯片之间。
使用TensorFlow，JAX和Lingvo中的ML模型实现实现了这些结果。在不到30秒的时间里，八个模型中的四个从零开始接受了训练。为了正确理解这一点，考虑到在2015年，花了超过三周的时间在可用的最先进的硬件加速器上训练其中一种模型。谷歌最新的TPU超级计算机可以在五年后将相同模型的速度提高近五个数量级。
在此文中，我们将研究对比的一些细节，提交如何实现如此高的性能，以及它们对模型训练速度的所有意义。

MLPerf模型

选择MLPerf模型来代表整个行业和学术界常见的尖端机器学习工作负载。上图中，这是每个MLPerf模型的更多详细信息：
DLRM代表了排名和推荐模型，这些模型是从媒体到旅行到电子商务等在线业务的核心
Transformer是自然语言处理（包括BERT）的最新发展浪潮的基础
BERT使Google搜索获得了“过去五年中最大的飞跃”
ResNet-50是广泛用于图像分类的模型
SSD是一种对象检测模型，其重量很轻，可以在移动设备上运行
Mask R-CNN是广泛使用的图像分割模型，可用于自主导航，医学成像和其他领域（您可以在Colab中进行实验）
除了上述最大规模的行业领先结果之外，Google还使用Google Cloud Platform上的TensorFlow提供了MLPerf提交，供今天的企业使用。

世界上最快的ML训练超级计算机

Google在本次MLPerf训练回合中使用的超级计算机是Cloud TPU v3 Pod的四倍，后者在之前的比赛中创下了三项记录。该系统包括4096个TPU v3芯片和数百个CPU主机，所有这些都通过超快速，超大规模的定制互连进行连接。总体而言，该系统可提供430多个PFLOP峰值性能。

表1：所有这些MLPerf提交的内容都是在Google新型ML超级计算机上以33秒或更短的时间从零开始进行训练的。2个

使用TensorFlow，JAX，Lingvo和XLA进行大规模训练

使用数千个TPU芯片训练复杂的ML模型需要在TensorFlow，JAX，Lingvo和XLA中结合算法技术和优化。
XLA是支持所有Google MLPerf提交的基础编译器技术，TensorFlow是Google的端到端开源机器学习框架，Lingvo是使用TensorFlow构建的序列模型的高级框架，而JAX是一个基于可组合功能转换的以研究为中心的新框架。
上面的记录设置规模依赖于模型并行性，按比例缩放的批次归一化，有效的计算图启动以及基于树的权重初始化。
上表中的所有TensorFlow，JAX和Lingvo提交（ResNet-50，BERT，SSD和Transformer的实现）都在2048或4096 TPU芯片上进行了训练，每个不到33秒。

TPU v4 vs TPU v3，性能提高2.7倍

Google的第四代TPU ASIC提供了TPU v3的矩阵乘法TFLOP的两倍以上，显着提高了内存带宽，并且在互连技术方面取得了进步。Google的TPU v4 MLPerf提交利用了这些新的硬件功能以及互补的编译器和建模优势。结果表明，在上一届MLPerf培训比赛中，TPU v4比TPU v3的性能平均提高了2.7倍，达到了类似的规模。

图2：Google的MLPerf培训v0.7中的TPU v4结果研究提交的结果，与Google的MLPerf培训v0.6可用提交中的TPU v3结果相比，平均提高了2.7倍，在相同规模的64核心芯片条件下。此性能改进归功于TPU v4中的硬件创新以及软件的改进。
没有有关节点、体系结构或设计的信息，但可以假定它们位于7nm类节点上。同样，FLOP中只有2倍，而性能则是2.7倍，这意味着它们在提取原始Flop方面更加有效。

Google已经在使用这项新技术，并将很快扩大其使用范围

Google实际上已经在使用新的TPU v4。该公司表示已经在自己的数据中心中使用了该技术。至于何时可以访问Google Cloud客户，这家搜索巨头尚未提供确切的日期。但它确实说，它们将在“今年晚些时候”上市。
该公司表示，目前，它将很快在其数据中心运行数十个最新的Pod。并指出，其中许多将“以或接近” 90％的无碳能源运行。至少就Google自身的运营而言，使新系统也成为最高效的系统之一。