谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

芯东西(公众号:aichip001)
作者 |  ZeR0
编辑 |  漠影

芯东西8月30日报道,昨晚,在时隔四年首次回归线下的谷歌云年度大会Next23上,谷歌云发布最新云端AI芯片TPU v5e,并罕见放出一段2分39秒的视频,披露了其云计算TPU数据中心内部实景,包括TPU、数据中心网络、光电路交换、水冷系统、生物识别安全验证等。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

谷歌TPU芯片为加速AI计算而生,是AI芯片中专用集成电路(ASIC)的代表之作,近年一直被视作全球AI芯片霸主英伟达GPU的劲敌。因此谷歌AI芯片的一举一动都备受业界关注。

据半导体研究和咨询公司SemiAnalysis的分析师Dylan Patel和Daniel Nishball爆料,谷歌拥有的算力资源比OpenAI、Meta、亚马逊、甲骨文和CoreWeave加起来还要多,其下一代大模型Gemini已经开始在新的TPUv5 Pod上进行训练,算力达到GPT-4的5倍,基于其目前的基础设施建设情况,到明年年底可能达到20倍

值得一提的是,谷歌相信AI工作负载没有唯一的标准,因此不只用自研AI芯片,还宣布跟英伟达扩大合作伙伴关系,为全球谷歌云客户提供基于英伟达GPU和生成式AI技术的云服务。英伟达创始人兼CEO黄仁勋也特意来到现场为谷歌云站台,与谷歌云CEO托马斯·库里安(Thomas Kurian)进行炉边对话。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

▲谷歌云CEO托马斯·库里安(左)与英伟达创始人兼CEO黄仁勋(右)同台(图源:英伟达)

谷歌云将于下个月全面上市的A3虚拟机,配备8张英伟达H100 GPU,能将训练AI模型的效率提高3倍,将网络带宽分配效率提高10倍。此外,谷歌云不仅是全球首家提供英伟达L4 GPU的云供应商,还是全球首批获得英伟达DGX GH200 AI超级计算机访问权限的公司之一。

一、TPU v5e:单位成本训练性能翻倍,跑大模型和生成式AI能效高

上次谷歌发布旗舰云端AI芯片TPU v4,还是在2021年5月的谷歌IO开发者大会上,开发者直到去年才真正用上TPU v4。

业界对第五代旗舰AI芯片TPU v5的出场翘首已久。没想到谷歌不按套路出牌,这次先亮出了为大规模、中等规模AI训练和推理而打造的优化版芯片TPU v5e。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光▲TPU v5e的系统架构

这是一款训推一体产品,每颗TPU v5e芯片包含1个张量核心,每个张量核心有4个矩阵乘法单元(MXU)、1个向量单元和1个标量单元

谷歌云将它称作是迄今最具成本效益、多功能和可扩展的TPU,能确保用户“将其TPU集群扩展到以前无法达到的水平”。

与Cloud TPU v4相比,TPU v5e每美元可提供高达2倍的训练性能,对于大型语言模型和生成式AI模型,每美元可提供2.5倍的推理性能,而成本不到TPU v4的一半

谷歌通过TPU v5e Pod平衡性能、灵活性和效率。TPU v5e支持多达256个芯片互连,总带宽超过400Tb/s,INT8性能达到100petaOps,从而解决更复杂的计算任务。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光▲TPU v5e的关键规格

TPU v5e现已推出预览版,将支持8种不同的虚拟机配置,从一个TPU芯片到单个切片中超过250个TPU芯片。谷歌云客户可按需选择配置,以服务于不同规模的大模型和生成式AI模型。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光▲TPU v5e支持的2D切片情况

TPU切片中,每个TPU虚拟机包含1、4或8个芯片,4芯片或者更少芯片将具有相同的非均匀内存访问(NUMA)节点。对于8芯片的TPU虚拟机,CPU-TPU通信在NUMA分区内将更加高效。例如在下图中,CPU0-Chip0的通信比CPU0-Chip4的通信快。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

AssemblyAI技术副总裁Domenic Donato说,在他们的生产ASR模型上运行推理时,Cloud TPU v5e每美元的性能始终比市场上同类解决方案高出4倍

Gridspace机器学习主管Wonkyum Lee亦谈道,其速度基准显示,当在谷歌Cloud TPU v5e上训练和运行时,AI模型的速度提高了5倍。推理指标规模也有显著改进,他们现在可以在1秒内实时处理1000秒的内部语音到文本和情感预测模型,处理速度提高了6倍

二、引入Multislice优化横向扩展,提高AI模型开发效率

从历史来看,通过横向扩展基础设施协调大规模AI工作负载需要手动处理故障、日志记录、监控和其他基本操作。

谷歌正着力简化TPU的操作,TPU v5e提供了与Google Kubernetes Engine(GKE)、Vertex AI、PyTorch、JAX、TensorFlow等主流框架的集成,并对各种主流开源工具提供内置支持,方便开发者使用熟悉的界面。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

以知名生成式AI独角兽企业Anthropic为例,据Anthropic联合创始人Tom Brown分享,GKE使其能大规模运行和优化GPU和TPU基础设施,Vertex AI使其能通过Vertex AI模型商店向客户分发模型。此外,他提到由A3和带有Multislice的TPU v5e驱动的谷歌下一代AI基础设施,将为Anthropic的工作负载带来性价比优势。

为了更容易地扩展训练工作,谷歌在预览版中引入了Multislice技术,这是一种将AI模型分配给数万个TPU芯片的方法,已经为谷歌最先进的PaLM模型的构建提供了动力。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

该技术能够超越物理TPU Pod的边界,通过单个Pod内的芯片间互连(ICI)或数据中心网络(DCN)上的多个Pod,将工作负载扩展到多达数万个TPU v5e或TPU v4芯片。

截至目前,使用TPU的训练任务仅限于单个TPU芯片切片,TPU v4的最大切片大小限制为3072个芯片。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

三、为什么谷歌要在机群中部署TPU v5e?

恰逢国际芯片顶会Hot Chips 2023正在举行,两位谷歌Fellow——谷歌首席科学家Jeff Dean和谷歌网络技术主管Amin Vahdat——共同发表了主题为《机器学习模型的激动人心的方向和对计算硬件的影响》的演讲。

Vahdat分享了为什么谷歌在其机群中部署了TPU v5e引擎,以及为什么可能带有i或e字母版本的TPU v6正在开发中、很快就会部署。

过去五年,AI模型变得越来越复杂,大型语言模型的参数数量以每年10倍的速度增长。而GPU和TPU的性能每年最多增长2~3倍,单芯片性能增长根本撑不住,所以必须通过集群扩展和稀疏化模型来应对飙涨的算力需求。但Vahdat认为这两种方法很快就会遇到瓶颈。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

正如谷歌在博客文章中所写,我们正处于计算机行业千载难逢的拐点。传统的设计和构建计算基础设施的方式已无法满足生成式AI和大型语言模型等工作负载的指数级增长需求。客户需要既具有成本效益又具有可扩展性的AI优化基础设施,同时还不能牺牲掉性能或灵活性。

为了将每TCO性能提高100倍,谷歌必须做很多事情:

1、创建专门的硬件TPU,用于密集矩阵乘法。矩阵乘法单元等功能可优化大型矩阵运算。

2、采用HBM内存,将这些矩阵数学引擎的内存带宽提高10倍。

3、创建专门的硬件加速器,用于稀疏矩阵中的scatter/gather操作,这被称作Sparsecore,它嵌入在TPU v4i、TPU v4和可能的TPU v5e引擎中。

4、采用液冷可最大限度地提高系统电源效率,从而提高经济效益。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光▲谷歌输送给超算的垂直电源,由其最新一代水冷系统冷却,该系统设计考虑了可持续性

5、使用混合精度和专门的数字表示,来提高设备的实际有效吞吐量。

6、具有用于参数分配的同步、高带宽互连,这是一种光电路交换(OCS),支持高带宽内存(HBM),当系统上的任务发生变化时几乎能即时重新配置网络,还提高了机器的容错能力。对于一个拥有成千上万个计算引擎且工作负载需运行数月的系统来说,这是一件大事。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光▲谷歌光电路交换技术使超算能进行重新配置,并为工作负载创建高效的机器形态和规格,按需动态扩展或缩小,无需人工参与

为了应对这些挑战,构建的计算基础设施类型必须改变。Vahdat说:“我们在过去五六十年间所形成的传统智慧实际上已被抛到了九霄云外。”

谷歌专注于优化硬件和软件以跨系统集群动态管理工作负载和功耗:

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

对于内存受限的任务,电压和电流强度可能会有很大差异,正如Vahdat所说,试图管理数千到数万个计算引擎集群的功耗是“介于困难和不可能之间的事情”。

通过不在集群中创建大量热点,可延长设备的使用寿命并减少中断,这对于AI训练等同步工作非常具有破坏性。与其回滚到检查点并从那里开始,不如从一开始就避免中断。

以下是如何利用核心频率和电压来平衡一些事情。在集群上查看任务放置的控制平面与这些任务的电源参数之间存在持续的干扰,并且这些任务在运行时的放置和移动是一个持续的过程。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

四、用AI辅助设计AI芯片,三个方法加速AI计算

Dean在演讲中分享的内容,跟TPU v5e相关度较少,但有助于为加速AI模型提供启发。

他重点分享了谷歌专注于稀疏性、自适应计算、动态神经网络这3种方法来驱动AI模型,并试图真正让AI专家系统开始设计AI芯片,以加快整个芯片开发周期、帮助优化AI硬件。

为了更快推出新的TPU芯片,谷歌已经使用其AI增强型EDA工具来帮助设计TPU v4i和TPU v4芯片的逻辑块,很可能还包括TPU v5e。

据Dean分享,目前要花大约三年才能推出芯片,6~12个月用于设计和探索,1年用于实施设计,6个月用于将其放到晶圆厂进行流片上,12个月用于投入生产、测试和提升。显然,硬件设计越接近新兴的AI模型越好。

Dean希望让系统架构师记住这些关键要点

1、加速器的连接性(带宽和延迟)很重要。

2、扩展对于训练和推理都很重要。

3、稀疏模型给内存容量和高效路由带来压力。

4、机器学习软件必须能够轻松表达有趣的模型,例如函数稀疏性。

5、功率、可持续性和可靠性确实很重要。

他谈道,AI模型的层数不断增加,参数数量呈爆炸式增长,由数十亿、数百亿、数千亿个tokens数据片段驱动。每次AI模型在一个新token上训练或在一个已完成的模型呈现一个新token来进行AI推理时,整个模型都会被激活。

有了像Pathways这样的框架,世界正在从针对不同任务的单独AI模型转向拥有能解决数百万任务的单一基础模型

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

拥有5000亿个参数的谷歌PaLM模型,便是在Pathways上进行训练的,通过在有6144个TPU v4引擎的2个Pod之间动态分配资源实现,TPU v4实际上分布在总共24576个引擎的6个Pod中。PaLM花了56天进行训练,下图是第5.71天的快照。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

谷歌Pathways框架采用稀疏性和自适应计算,在运行推理或训练时动态地为给定模型添加或减少核心池的容量。该框架的一个秘密武器,是知道应该激活稀疏模型的哪些部分,它已经通过谷歌正在秘密筹备的“杀手级”大模型项目Gemini得到完善。不过谷歌没将Pathways框架开源。

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

以稀疏方式激活模型,能够大大降低进一步训练和生产推理所需的计算成本。

这里谈到的稀疏性,跟通常说的单个向量或张量内稀疏性的加速器细粒度稀疏性不是一回事。它也不同于粗粒度稀疏性,模型中的大型模块要么被激活,要么不被激活。为了直观反映Dean说的稀疏性,外媒The Next Platform将Dean的几张图拼到一张图中:

谷歌最新AI芯片发布!与英伟达黄仁勋同台,TPU系统实景曝光

据Dean解释,当前大多数稀疏化工作对每个专家都用相同的尺寸和结构,所以有一些绿色的专家,通过学习一些路由函数,可以了解哪个专家擅长做哪种事,然后将一些例子发送给合适的专家。计算平衡通常是通过每个专家的计算量相等以及每个专家的示例数量相等来实现的。

对于计算机架构师来说,这意味着跨加速器的全对全shuffle性能非常重要,对于所有稀疏模型都是成立的。不过你可能想能够做的,不是用固定的计算成本,而是改变模型不同部分的计算成本。“我们应该在真正困难的事情上花费的计算量,应该是非常简单的事情上花费计算量的100倍。”Dean说。

结语:拉大AI基础设施性价比优势,在云计算竞赛中越来越关键

二十年来,从构建为如今大模型盛世奠定基础的谷歌Transformer架构,到自研TPU芯片、构建为YouTube、Gmail、谷歌地图、谷歌Play和Android等数十亿用户提供服务的AI优化基础设施,谷歌已建立一系列行业领先的AI能力。

谷歌云人工智能业务和解决方案副总裁Phil Moyer称,Anthropic、Character AI、Typeface等生成式AI独角兽创企中,70%以上都是谷歌云的客户;获得总投资的AI初创公司中,有1/4正在采用谷歌云。

但这显然远远不够,如今生成式AI和大模型已经成为吸引云客户的一大关键因素,要应对亚马逊云科技、微软等强劲云计算对手的猛烈攻势,谷歌不仅要继续完善支持模型训练、调优和全球范围内服务的整个生命周期的端到端软件和服务,而且需不断提高其AI基础设施的性价比优势。