谷歌新一代 AI 芯片发布，Jeff Dean：AI 硬件性能提升更难了

昨夜晚间，谷歌突然对外披露公司新一代 AI 加速器 Cloud TPU v5e。

谷歌在博客中指出，我们正处于计算领域千载难逢的拐点。设计和构建计算基础设施的传统方法不再足以满足生成式人工智能和 LLM 等工作负载呈指数级增长的需求。事实上，过去五年里 LLM 的参数数量每年增加 10 倍。因此，客户需要具有成本效益且可扩展的人工智能优化基础设施。

Google Cloud 则通过提供领先的人工智能基础设施技术、TPU 和 GPU，以满足开发者的需求，今天，谷歌宣布，公司在这两个产品组合均得到重大增强。首先，我们正在通过 Cloud TPU v5e 扩展我们的 AI 优化基础设施产品组合，这是迄今为止最具成本效益、多功能且可扩展的 Cloud TPU，现已提供预览版。TPU v5e 提供与 Google Kubernetes Engine (GKE)、Vertex AI 以及 Pytorch、JAX 和 TensorFlow 等领先框架的集成，因此您可以通过易于使用、熟悉的界面开始使用。

同时，谷歌还宣布，公司基于 NVIDIA H100 GPU 的 GPU 超级计算机 A3 VM 将于下个月全面上市，为您的大规模 AI 模型提供支持。、

TPU v5e，为大模型而生

谷歌表示，Cloud TPU v5e 是 Google Cloud 最新一代 AI 加速器，专为提供中大规模训练和推理所需的成本效益和性能而构建。虽然谷歌并没披露更多关于这个芯片的工艺，但据猜测，这可能是基于 5nm 打造的。

了解谷歌 TPU 芯片的读者都知道，在之前的产品中，他们都有一个或两个 TensorCore 来运行矩阵乘法。与 v2 和 v3 Pod 类似，v5e 每个芯片都有一个 TensorCore。每个 TensorCore 有 4 个矩阵乘法单元 (MXU)、一个向量单元和一个标量单元。下表显示了 v5e 的关键规格及其值。相比之下，v4 Pod 每个芯片有 2 个 TensorCore。

与 Cloud TPU v4 相比，TPU v5e 可为 LLM 和新一代 AI 模型提供高达 2 倍的训练性能和高达 2.5 倍的推理性能。TPU v5e 的成本不到 TPU v4 的一半，使更多组织能够训练和部署更大、更复杂的 AI 模型。

谷歌指出，v5e 每个 Pod 的芯片占用空间较小，为 256 个，经过优化，成为 transformer、文本到图像和卷积神经网络 (CNN) 训练、微调和服务的最高价值产品。

换而言之，通过 TPU v5e Pod 平衡性能、灵活性和效率，允许最多 256 个芯片互连，总带宽超过 400 Tb/s 和 100 petaOps INT8 性能。TPU v5e 的用途也非常广泛，支持八种不同的虚拟机 (VM) 配置，范围从一个芯片到单个片内的 250 多个芯片。这使得客户能够选择正确的配置来服务于各种 LLM 和 gen AI 模型大小。

与此同时，Cloud TPU v5e 还为领先的 AI 框架（如 JAX、PyTorch 和 TensorFlow）以及流行的开源工具（如 Hugging Face 的 Transformers 和 Accelerate、PyTorch Lightning 和 Ray）提供内置支持。谷歌透露，公司将通过即将发布的 PyTorch/XLA 2.1 版本进一步加强对 Pytorch 的支持，其中包括 Cloud TPU v5e 支持，以及用于大规模模型训练的模型和数据并行性等新功能。

据介绍，在公共预览版，推理客户可以使用 1 芯片 (v5litepod-1)、4 芯片 (v5litepod-4) 和 8 芯片 (v5litepod-8)。“为了为具有更多芯片的虚拟机腾出空间，调度程序可能会抢占具有较少芯片的虚拟机。因此，8 芯片虚拟机很可能会抢占 1 芯片和 4 芯片虚拟机。” 谷歌表示。

他们进一步指出，为了更轻松地扩展训练作业，谷歌还在预览版中引入了 Multislice 技术，该技术允许用户轻松扩展 AI 模型，使其超出物理 TPU pod 的范围，最多可容纳数万个 Cloud TPU v5e 或 TPU v4 芯片。

谷歌在博客中写道，到目前为止，使用 TPU 的训练作业仅限于单个 TPU 芯片 slice，TPU v4 的最大 slice 大小为 3,072 个芯片。借助 Multislice，开发人员可以在单个 Pod 内通过芯片间互连 (ICI)：inter-chip interconnect 或通过数据中心网络 (DCN：Data center network) 跨多个 Pod 将工作负载扩展到数万个芯片。Multislice 技术同时还为谷歌最先进的 PaLM 模型的创建提供了动力。

“我们真正致力于使其成为一个可扩展的解决方案，” 谷歌的 Lohmeyer 说。“我们跨软件和硬件进行设计。在这种情况下，该软件 [和] 硬件与 Multislice 等新软件技术协同工作的魔力，使我们的客户能够轻松扩展其 AI 模型，超越单个 TPU pod 或单个 GPU 集群的物理边界，” 他解释道。“换句话说，单个大型人工智能工作负载现在可以跨越多个物理 TPU 集群，扩展到数万个芯片，而且这样做非常经济高效。” Lohmeye 强调。

Jeff Dean：AI 硬件性能的下一个 100 倍将会更加困难

十年前，Jeff Dean 在一张纸上做了一些数学计算，并发现：如果谷歌将人工智能功能添加到其同名搜索引擎中，它就必须将其数据中心占地面积扩大一倍，这就趋势谷歌走上创建定制张量处理单元（TPU）矩阵数学引擎的道路。

十年过去了，人工智能变得更加复杂、计算密集，备受讨论的 TPUv4 铁虽然现在和未来很多年都很有用，但看起来有点过时了。全新推出的 TPUv5e 将替代 TPUv4 系统。

在日前开幕的 Hotchips 演讲中，Jeff Dean 表示，谷歌专注于驱动人工智能模型的三种不同方法——稀疏性、自适应计算和动态神经网络——并且还试图让人工智能蛇吃掉它的尾巴而不是蚕食它。

在 Jeff Dean 看来，真正让人工智能专家系统开始设计人工智能处理器可加快整个芯片开发周期，从而帮助不断改进的硬件进入该领域，以满足更快增长的模型。

Dean 解释说，到目前为止创建的人工智能模型，整个模型的层数不断增加，参数数量呈爆炸式增长，由数十亿、数百亿、数千亿的 token snippets 数据驱动，每次人工智能都会被激活。在新 token 上训练的模型或针对已完成的模型提出 token 以进行人工智能推理。

但是，有了像 Pathways 这样的框架（它是 Google PaLM 系列模型的基础），世界正在从专门针对不同任务的单独 AI 模型转向拥有单一基础模型。

对于稀疏模型，AI 模型的各个部分会在需要时激活（activate ），并且仅激活这些部分。模型如何知道要激活哪些部分尚不清楚，这就是 Pathways 框架中的秘密武器，该框架已通过 Gemini 模型得到完善，毫无疑问使用了 Dean 所说的技术。

值得注意的是，Pathways 框架不像 Google 创建的早期且可能更为初级的 TensorFlow 框架（该框架于 2015 年 11 月开源）那样是开源的。因此，我们只能知道 Google 告诉我们的有关 Pathways 和 Gemini 的信息。

“与这个巨大的模型相比，稀疏模型的效率要高得多，” Dean 解释道。“他们只是调用整个模型的正确部分——正确的部分也是在训练过程中学到的东西。然后模型的不同部分可以专门用于不同类型的输入。最终的结果是，您最终会得到一些非常大的模型的正确 1% 或正确 10% 的结果，这会提高您的响应能力和更高的准确性，因为您现在拥有比您大得多的模型容量可以进行其他训练，然后可以调用正确的部分。”

根据 Dean 的说法，稀疏性的另一个方面对于系统架构师来说很重要，它与加速器中通常谈论的细粒度稀疏性不同，在加速器中，单个向量或张量内的稀疏性（通常每四个中的两个）矩阵中的值设置为零，将其从密集型转换为稀疏型），这也不同于粗粒度稀疏性，其中模型中的大型模块要么被激活，要么不被激活。这种稀疏性看起来像这样，我们将几个 Dean 的图表合并到一页上，这样您就可以将其全部理解：

“当今大多数稀疏工作对每个专家都使用相同的大小和结构，” Dean 说。“所以这里有一些绿色专家为他们服务。这里有一些学习路由功能，可以了解哪个专家擅长哪种事情，然后将一些示例发送给适当的专家。计算平衡通常是通过每个专家的计算量相等以及每个专家的示例数量相等的流量来实现的。对于计算机架构师来说，这意味着跨加速器的全面混洗性能非常重要。对于所有稀疏模型来说都是如此——您希望能够以正确的方式将事物从模型的一个部分快速路由到另一部分。”

“不过，您可能想做的一件事是，不是具有固定的计算成本，而是改变模型不同部分的计算成本。在每个示例上花费相同的计算量是没有意义的，因为有些示例的难度是原来的 100 倍。我们在真正困难的事情上花费的计算量应该是非常简单的事情上的 100 倍。” Dean 强调。

事实证明，一些小型专家可能只需要少量计算，并且可以用于生产中使用的模型中大约 90% 的提示。专家们为了处理更复杂的事情而变得更大，具有不同的计算结构和可能更多的层，并且它们的计算量更大，因此运行成本更高。如果您正在运行人工智能服务，您将希望能够将成本归因于所提供的专家答案的价值，以便您可以适当收费。

提高效率和计算能力的途径

当然，这不是谷歌的独家理论——该公司之所以谈论它是因为 Pathways 框架这样做：

这就是稀疏性和自适应计算。

Dean 说，上图中提到的最后一件事是动态神经网络，这意味着可以在正在运行的系统中添加或删除容量——这是我们几十年来通用服务器所拥有的（虽然不是在 X86 平台上，但奇怪的是，这就是 Arm 和 RISC-V 可能能够赶上大型机和 RISC/Unix 系统的地方）。CPU 及其工作负载的情况（虚拟机管理程序级别肯定存在动态分配）对于 GPU、TPU 和其他 AI 计算引擎来说也是如此。您希望能够在运行推理或训练时动态地为任何给定模型添加或减少核心池的容量。

来自 Google 的具有 5000 亿个参数的 PaLM 模型在 Pathways 上进行了训练，并通过在具有 6,144 个 TPUv4 引擎的一对 pod 之间动态分配资源来实现这一目标，但 TPUv4 引擎实际上分布在总共 24,576 个引擎的 6 个 pod 中，所有这些都链接在一起通过高速数据中心网络。像这样：

以下是 Dean 希望给系统架构师留下深刻印象的关键要点：

• 加速器的连接性（带宽和延迟）很重要

• 规模对于训练和推理都很重要

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

• 机器学习软件必须能够轻松表达有趣的模型 - 就像上图中所示的函数稀疏性

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

Google Fellow 以及工程副总裁 Amin Vahdat 在随后的演讲中展示了人工智能行业面临的模型大小增长的指数曲线：

据观察，我们绝对没有理由相信模型的复杂性会因此而降低，计算能力的需求将会降低。但根据我们的估计，模型每年增长 10 倍，GPU 和 TPU 的性能最多每年增长 2 到 3 倍。公司必须通过扩展来弥补这一点，这很困难，并改进他们的模型，这也很困难。

这就是为什么谷歌已经在其机群中部署了 TPUv5e 引擎。在我们看来，谷歌的 TPUv6 可能正在开发中并准备就绪很快就会部署，以帮助支持 Gemini 模型的商业化。

为了使每 TCO 的性能提高 100 倍，Vahdat 进行了一个深入的讲座，介绍了如何衡量人工智能或通用计算平台的相对价值，在人工智能出现之前，我们一直同意这一点系统——谷歌必须做很多事情，当中包括：

• 创建专门的硬件 - TPU - 用于密集矩阵乘法。

• 使用 HBM 内存将这些矩阵数学引擎的内存带宽提高 10 倍。

• 创建专门的硬件加速器，用于稀疏矩阵中的分散/聚集操作 - 我们现在称之为 Sparsecore，它嵌入在 TPUv4i、TPUv4 和可能的 TPUv5e 引擎中。

• 采用液体冷却可最大限度地提高系统电源效率，从而提高经济效益。

• 使用混合精度和专门的数字表示来提高设备的实际吞吐量（Vahdat 称之为 “有效吞吐量”）。

• 并具有用于参数分配的同步、高带宽互连，事实证明，这是一种光路开关，可以在系统上的作业发生变化时几乎即时重新配置网络，并且还提高了机器的容错能力。对于一个拥有数万个计算引擎且工作负载需要数月才能运行的系统来说，这是一件大事，世界各地的 HPC 中心都非常清楚这一点。

“我们为应对这一挑战而必须构建的计算基础设施类型必须改变，” Vahdat 在主题演讲中说道。“我认为，值得注意的是，如果我们试图在通用计算上做到这一点，我们就不会达到今天的水平。换句话说，我们在过去 50 到 60 年间所形成的传统智慧实际上已经被抛到了九霄云外。我认为可以公平地说，在谷歌，但更重要的是，在整个社区，机器学习周期将占据我们想做的事情中越来越多的部分。”

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

对于受内存限制的部分，电压和电流强度可能会有很大差异，正如 Vahdat 所说，试图管理数千到数万个计算引擎集群的功耗是 “介于困难和不可能之间的事情”。通过不在集群中创建大量热点（这可能发生在 Google 训练 PaLM 模型时），可以延长设备的使用寿命并减少中断，这对于 AI 训练等同步工作非常具有破坏性，就像 HPC 模拟一样和建模。与其回滚到检查点并从那里开始，不如从一开始就避免中断。

以下是如何利用核心频率和电压来平衡一些事情。

当然，快速更新 TPU 芯片是谷歌工作最重要的一部分。

Dean 表示，目前芯片的问世大约需要三年时间。这需要六到十二个月的时间进行设计和探索，一年的时间来实施设计，六个月的时间与晶圆进行流片，十二个月的时间用于将其投入生产、测试和提升。

目前尚不清楚人工智能可以在多大程度上缩短芯片开发周期或可以减少多少人力，Dean 也没有提供任何估计。但显然，硬件设计越接近新兴人工智能模型就越好。