新智元报道
谷歌如今能保持人工智能领域霸主地位,很大程度上依赖于自主研发的TPU。
自去年开始,谷歌就开始不断放风TPUv4,在论文中描述新一代TPU的架构设计,从前几代TPU设计中不断吸取经验。
5月19日,谷歌正式推出新一代人工智能ASIC芯片TPUv4,运算效能是上一代产品的2倍以上,由4096个TPUv4单芯片组成的pod运算集群,可释放高达1exaflop(每秒10的18次方浮点运算)的算力,超过了目前全球运算速度最快的超级计算机。
前段时间谷歌发布的5400亿参数语言模型PaLM就是用两个TPUv4 pods训练了大约64天得到的。
10月11日,谷歌举办Google Cloud Next 2022大会,正式面向大众开放第四代TPU使用权限。
有外媒记者与TPU的幕后人员进行深入交流,并研究大量的论文和技术报告后,写成了一篇TPUv4最全面、深入的报道,对第四代TPU的计算引擎及其相关系统进行概述。
第四代TPU
早在2013年,当谷歌正在为Google搜索引擎开发语音激活搜索功能时,谷歌研究员Jeff Dean在一张废纸上做了一会数学计算,发现如果这种人工智能辅助搜索投入生产,那么谷歌将不得不将其数据中心的规模扩大一倍。
题外话,Jeff Dean参与了许多谷歌的关键技术开发,包括谷歌的网络爬虫、索引、查询系统、AdSense、谷歌翻译的设计与实现、TensorFlow等等。
扩容数据中心对于一个「免费」的产品来说,并不会让搜索广告客户突然增加两倍,因此,谷歌必须努力为人工智能训练和推理建立一个更好的数学引擎。
TPU的目标是在一个特定领域对芯片进行体系结构设计,在去年发布的一篇论文中,谷歌详细介绍了TPUv4i的各项参数及性能,专门为推理进行调整优化,其中i就代表了推理。
论文链接:https://www.gwern.net/docs/ai/2021-jouppi.pdf
事实证明,TPUv4i推理芯片的生产大大领先于TPUv4通用引擎,其生产方式与Nvidia相反,这也表明了一个事实:谷歌真的非常需要廉价的推理引擎来驱动数十个应用程序,有TPU加持的人工智能可以扩展到更丰富的应用场景。
Google并没有准确的时间表报告何时将完整TPUv4的系统投入生产,据猜测大约是在2021年第四季度。
谷歌并不耻于在TPU设计上多留下一两个节点,这样做反而可以保持芯片设计和生产的低成本。
前沿计算引擎的成本很大一部分来自于高精尖设计工艺的不良率,虽然可以通过在设计中添加备用组件来降低成本,但是这同样也是有成本的。
而具有多个向量的芯片更便宜,也更容易通过供应链进入谷歌手中。
TPUv4i 推理晶片采用台积电公司的七纳米工艺制造,一年半前投产,作者猜测通用的TPUv4也是采用同样的七纳米工艺。
并且作者预测TPUv5i和TPUv5计算引擎将会采用5纳米进程,几乎可以肯定的是,这些引擎即将开始测试了,如果Google生产了 TPUv4i和TPUv4,以及 Google Cloud上的客户能够普遍使用 TPUv4的话。
TPU 核心的基本结构是,它有一个标量和并行向量处理机的前端,带有一个自主开发的矩阵数学单元,可以进行 BF16(谷歌为其 TPU 发明的一种格式)乘法和 FP32累加运算(有时还可以根据模型进行 INT8推理)。
最初的 TPUv1处理器有一个256 × 256的矩阵数学单元,这个单元非常庞大,而且正如 Google设计过程中改进的,这并不是最佳配置。
事实证明,谷歌自己的人工智能应用程序也不是64 × 64的数组。
矩阵吞吐量和利用率之间的最佳点是128 × 128数组,至少对于 Google 自己的代码来说是这样的。
Google发现了这一点之后,公司通过增加TPU核心(其中嵌入了 MXU 单元)或者在每个核心中增加 MXU 单元来扩大 TPU 的规模。
为了实现这个配置,每一代TPU的服务器板的样式都有变化。
在 TPUv4服务器板的右侧,可以看到一些相邻的芯片。这是一个6端口6 Tb/秒的交换机作为网络接口卡和3D 环形曲面网络(torus network)的基础。
刚听到TPUv4芯片消息的时候,大家并不确定Google是否会增加内核或者扩展 MXUs,而事实证明,Google一直将内核数量保持在2,并且再次将每个内核内的MXUs数量翻了一番。
除非有什么变化,否则这也将是Google创建TPUv5、 TPUv6等所有后代的设计方式。
在某个时候点,谷歌采用四核心取决于其自主开发的XLA机器学习编译器能做什么,但目前,提供更少数量的核是扩展应用程序最简单的方法。
可以注意到,TPUv4i 推理芯片本身就相当强大,单个TPU 核心就有四个128 × 128 MXU,运行频率为1.05 GHz,但只占用了400平方毫米的面积。
根据一些平面图估计,在单个内核上将 MXU 计数增加一倍,缓存也增加一倍,并添加一些其他的非核心特性来创建单个内核的全功能 TPUv4,可能只会将芯片尺寸增加到563平方毫米左右。
但是不管出于什么原因,可能与 XLA 编译器中双核 TPU 的优化有关,Google 做了一个双核心TPUv4 AI训练引擎,它的面积大约在780平方毫米左右。
可以看到,这两个核心TPUv4芯片表现为一个逻辑核心,一个32GB HBM 内存空间。
所以有理由怀疑它是一个整体设计,但它可能是由两个 TPUv4芯片捆绑在一个插口(socket)上,可以降低成本,但可能影响性能,具体取决于插口内互连。
中杯、大杯、超大杯
在过去三代TPU核心中,考虑向量规模的话,TPUv3是对TPUv2的一次渐进式「中年升级」,在输入和速度上提高了30%左右,在同样的16纳米工艺下 ,晶体管数量和芯片尺寸只有名义上的11%增加,芯片性能提高了2.67倍,HBM 主存容量提高了2倍(可以输入更大的数据集)。
除此之外,最大的区别在于2D 环面互连从 TPUv2芯片中的256个芯片扩展到TPUv3代中的1024个芯片,这也导致pod处理能力增加了10.7倍,从12petaflops增加到126petaflops(在 BF16操作中测量)。
TPUv4是计算引擎的真正升级,进程从16纳米缩小到7纳米,并且具有其他所有优点:MXU 的数量再次增加了一倍,缓存内存增加了九倍达到244 MB,HBM2内存带宽增加了33%到1.2 TB/秒,但有趣的是,HBM2内存容量保持在32 GB。
谷歌之所以能够做到这一点,是因为它可以疯狂地扩展TPUv4 pods,比如Nvidia 吹嘘自己可以将多达256个Hopper GH100 GPU与其NVSwitch结构紧密耦合,但新的3D 环面互连首次与带宽更大、基数更高的 TPUv4连接,可以将4096个TPUv4引擎紧密耦合在一起,总计达到1.126 exaflops 的 BF16计算。
其中8个 TPUv4 pods位于谷歌俄克拉荷马州梅斯县的数据中心,能够产生9 exaflops的原生人工智能计算速度。
TPUv4 总览
谷歌云 TPU 产品的出境产品经理Vaibhav Singh表示,在大部分情况下,工程师使用64个芯片,甚至更少,TPUv4做的一件有趣的事情是多维数据集级别的环绕连接,其中一个多维数据集的单位是四乘四乘四,因为就最常见的工作负载大小而言,这也是最优的尺寸。
从而可以构造出由四乘四乘四组合而成更大的切片。
在上面的表格中,立方体切片的大小是计算核心,而不是芯片,切片大小是动态配置的。
Singh还补充道,当我们讨论的模型不适用于单个芯片,必须分布在数百个芯片上时,对物理拓扑结构的认识,以及你的软件在计算和通信重叠方面实际上尝试进行某些优化后,会对系统的整体性能产生巨大的影响。
高性能计算机用户非常清楚3D 和更高维环形互连的好处—— IBM 的 BlueGene 超级计算机有一个3D 环形,富士通的Tofu互连用于K和 Fukagu 的超级计算机是6D 环形/网状互连——即使如果加入机器,重新接线是件痛苦的事情,通常需要升级和重新连接的系统选择 Clos 或蜻蜓拓扑。
TPUv4i 和 TPUv4,特别是与它们的前辈相比,以及与同样针对人工智能训练和推理的竞争性图形处理器相比,展示了极简主义设计的价值,这种设计架构只为了做一件事:用于人工智能处理的矩阵数学计算。
尽管也有一些研究人员正试图利用 TPU 来运行加速 HPC 模拟和建模应用程序。
随着 Hopper GPU的普及,Google 使用 TPUv4引擎与其他加速器进行的比较必然要与Nvidia 的安培 A100图形处理器进行比较。
谷歌系统和服务基础设施副总裁 Amin Vahdat 提供了一张图,通过运行不同的神经网络模型,展示了 TPU与A100的对比结果,以及谷歌云上的 TPUv4 pod如何与微软 Azure 上的类似规模的 A100集群对抗。
在较低规模的模型中,性能差异并不像Meta平台的深度学习推荐模型(DLRM)那样大。
Vahdat 说,平均而言,TPUv4在 MLPerf 测试中的表现比 Nvidia A100s 高出40% ,有趣的是,Google 在这五个测试中提交的 MLPerf 2.0结果比使用 MLPerf 1.0套测试的 TPUv4高出50%
归根结底,比硬件供应和速度更重要的是人工智能训练的成本,使用4096个TPUv4 pod和一个同样堆叠规模的Nvidia A100 GPU对比。
Vahdat的团队根据提交的 MLPerf 结果以及他们自己对 BERT 和 ResNet 如何在 GPU 上扩展的了解建立了成本模型。
目前还不清楚与 H100GPU 加速器相比会是什么样子。Nvidia 可能提供3到6倍的性能(取决于工作负载) ,成本为2.5到3倍。
在价格上,TPUv4芯片上花费一个小时的价格比在TPUv3芯片上花费的价格高出61% ,但是它提供的性能是TPUv3芯片的2.23倍,也就是说性价比提高了28%
在TPUv2到 TPUv3的跃升时,性能提高了2.67% ,性价比提高了33.5%