谷歌在博客中表示，我们正处于计算领域历史上一个重要的转折点。过去设计和构建计算基础设施的方式已经无法满足生成式人工智能和大型语言模型等工作负载快速增长的需求。实际上，在过去五年中，大型语言模型的参数数量每年增长了十倍。因此，客户需要具备经济实惠且可扩展的人工智能优化基础设施。
谷歌云通过提供领先的人工智能基础设施技术、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模型提供支持。
谷歌的Cloud TPU v5e是Google Cloud最新一代的AI加速器，专为中大规模训练和推理提供成本效益和性能而设计。尽管谷歌没有透露更多关于该芯片工艺的信息，但据推测，它可能是基于5nm制造的。
对于熟悉谷歌TPU芯片的读者来说，在之前的产品中，每个芯片都有一个或两个TensorCore来执行矩阵乘法。与v2和v3 Pod类似，v5e每个芯片都包含一个TensorCore。每个TensorCore有4个矩阵乘法单元(MXU)、一个向量单元和一个标量单元。下表列出了v5e的关键规格及其值。相比之下，v4 Pod每个芯片有两个TensorCore。
与Cloud TPU v4相比，TPU v5e可以提供高达2倍的训练性能和高达2.5倍的推理性能给LLM和新一代AI模型。此外，TPU v5e的成本不到TPU v4的一半，使更多的组织能够训练和部署更大、更复杂的AI模型。
谷歌指出，v5e每个Pod的芯片占用空间较小，仅为256个。经过优化，它成为了转换器、文本到图像和卷积神经网络(CNN)训练、微调和服务的最有价值的产品。
换句话说，TPU v5e Pod通过平衡性能、灵活性和效率，最多允许256个芯片互连，总带宽超过400 Tb/s和100 petaOps INT8性能。它的用途也非常广泛，支持八种不同的虚拟机(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)或数据中心网络(DCN)跨多个Pod将工作负载扩展到数万个芯片。Multislice技术还为谷歌最先进的PaLM模型的创建提供了动力。
"我们真正致力于使其成为一个可扩展的解决方案，"谷歌的Lohmeyer说。 "我们跨软件和硬件进行设计。在这种情况下，该软件[和]硬件与Multislice等新软件技术协同工作的魔力，使我们的客户能够轻松扩展其AI模型，超越单个TPU pod或单个GPU集群的物理边界，"他解释道。 "换句话说，单个大型人工智能工作负载现在可以跨越多个物理TPU集群，扩展到数万个芯片，而且这样做非常经济高效。”Lohmeye强调。
在十年前，Jeff Dean 在一张纸上进行了一些数学计算，并得出结论：如果谷歌要将人工智能功能添加到其搜索引擎中，就必须将数据中心的面积扩大一倍。这促使谷歌开始开发定制张量处理单元（TPU）矩阵数学引擎。
如今十年过去了，人工智能变得更加复杂和计算密集。尽管备受讨论的 TPUv4 仍然有用，但它看起来有些过时了。因此，谷歌推出了全新的 TPUv5e，以取代 TPUv4。
在最近的 Hotchips 演讲中，Jeff Dean 表示谷歌专注于三种不同的方法来驱动人工智能模型——稀疏性、自适应计算和动态神经网络，并试图让人工智能能够自我学习而不是依赖于大量的数据。在 Dean 看来，人工智能专家系统的出现可以加快整个芯片的开发周期，从而帮助硬件不断改进以满足快速增长的模型需求。
目前为止，创建的人工智能模型层数不断增加，参数数量呈爆炸式增长，由数十亿、数百亿甚至数千亿的 token snippets 数据驱动，每次人工智能都会被激活。在新 token 上训练模型或为已完成的模型提供 token 以进行推理。
然而，有了像 Pathways 这样的框架作为基础，世界正在从为每个任务单独设计 AI 模型转向拥有单一基础模型的阶段。对于稀疏模型，AI 模型的各个部分会在需要时被激活，并且仅激活这些部分。Pathways 框架通过 Gemini 模型得到了完善，其中使用了 Dean 提到的技术。值得注意的是，Pathways 框架并不像早期且可能更初级的 TensorFlow 框架那样是开源的。因此，我们只能根据 Google 提供的信息了解有关 Pathways 和 Gemini 的情况。
相对于庞大的模型，稀疏模型的效率要高得多。”Dean解释道。“他们只调用整个模型中正确的部分——这些部分也是在训练过程中学到的。然后，模型的不同部分可以专门用于不同类型的输入。最终结果是，您将获得一些非常大的模型的正确 1% 或正确 10% 的结果，这提高了响应能力和准确性，因为您现在拥有比原始模型大得多的容量来进行其他训练，并调用正确的部分。”
根据 Dean 的说法，稀疏性的另一方面对于系统架构师很重要，这与加速器中通常讨论的细粒度稀疏性不同。在加速器中，单个向量或张量内的稀疏性（通常每四个中的两个）矩阵中的值设置为零，将其从密集型转换为稀疏型），这也不同于粗粒度稀疏性，其中模型中的大型模块要么被激活，要么不被激活。这种稀疏性看起来像这样：我们将几个 Dean 的图表合并到一页上，以便您可以全面理解：
“如今，大多数稀疏工作对每个专家都使用相同的大小和结构。”Dean说。“这里有一些绿色专家为他们服务。还有一些学习路由功能，可以了解哪个专家擅长哪种任务，然后将示例发送给适当的专家。计算平衡通常是通过每个专家的相等计算量以及每个专家的示例数量的流量来实现的。对于计算机架构师来说，这意味着跨加速器的全面混洗性能非常重要。所有稀疏模型都需要考虑这一点——您希望能够以正确的方式快速路由事物从模型的一个部分到另一个部分。”
“然而，您可能想做的一件事是改变模型不同部分的计算成本，而不是具有固定的计算成本。在每个示例上花费相同的计算量是没有意义的，因为有些示例的难度是原来的 100 倍。我们在真正困难的事情上花费的计算量应该是在简单事情上的 100 倍。”Dean强调道。
事实证明，一些小型专家可能只需要少量计算，并且可以用于生产中使用的模型中大约 90% 的提示。专家们为了处理更复杂的事情而变得更大，具有不同的计算结构和可能更多的层，并且它们的计算量更大，因此运行成本更高。如果您正在运行人工智能服务，您希望能够将成本归因于所提供的专家答案的价值，以便您可以适当收费。
提高效率和计算能力的途径是谷歌提出的 Pathways 框架的核心理念。该框架通过稀疏性和自适应计算来实现这一目标。Dean 指出，动态神经网络是实现此目标的关键，因为它允许在系统运行过程中增加或减少容量。无论是通用服务器、X86 平台还是 Arm 和 RISC-V，CPU 及其工作负载的情况都适用于 GPU、TPU 和其他 AI 计算引擎。通过动态调整核心池容量，您可以在推理或训练期间满足不同模型的需求。
谷歌的 PaLM 模型是一个拥有 5000 亿个参数的例子，它使用 Pathways 框架进行训练并通过在具有 6,144 个 TPUv4 引擎的一对 pod 之间动态分配资源来实现这一目标。然而，实际上 TPUv4 引擎分布在总共 24,576 个引擎的 6 个 pod 中，并通过高速数据中心网络相互连接。
对于系统架构师而言，以下是一些关键的要点：
- 加速器的连接性（带宽和延迟）非常重要。
- 规模对于训练和推理都至关重要。
- 稀疏模型给内存容量和高效路由带来挑战。
- 机器学习软件必须能够轻松表达有趣的模型，例如上图所示的函数稀疏性。
- 功率、可持续性和可靠性确实是重要的考虑因素。
在随后的演讲中，Google Fellow 和工程副总裁 Amin Vahdat 展示了人工智能行业面临的模型规模增长的指数曲线。他指出，没有理由相信模型的复杂性会降低，而计算能力的需求将会增加。据估计，模型每年增长约10倍，而 GPU 和 TPU 的性能最多每年增长2到3倍。为了弥补这一点并改进模型，公司必须进行扩展，但这很困难。
这就是谷歌已经在其机群中部署了 TPUv5e 引擎的原因。谷歌的 TPUv6 可能正在开发中并准备就绪很快就会部署，以帮助支持Gemini模型的商业化。为了提高每 TCO 的性能100倍，Vahdat 进行了一个深入的讲座，介绍了如何衡量人工智能或通用计算平台的相对价值。他提到了创建专门的硬件，如TPU用于密集矩阵乘法，使用HBM内存提高这些矩阵数学引擎的内存带宽，创建专门的硬件加速器用于稀疏矩阵中的分散/聚集操作等。此外，他还强调了液体冷却、混合精度和专门的数字表示等技术的应用，以提高设备的实际吞吐量和系统的容错能力。
总的来说，Vahdat认为我们为应对这一挑战而必须构建的计算基础设施类型必须改变。如果我们试图在通用计算上做到这一点，我们将无法达到今天水平。换句话说，过去50到60年间所形成的传统智慧实际上已经被抛到了九霄云外。在谷歌以及整个社区中，机器学习周期将占据我们想做的事情中越来越多的部分。
谷歌致力于优化硬件和软件，以便能够跨系统集群动态管理工作负载和功耗。其中一个关键问题是如何管理受内存限制的部分的功耗。电压和电流强度可能会有很大差异，这在管理数千到数万个计算引擎集群的功耗时是一个“困难且不可能”的任务。通过避免在集群中创建大量热点（这可能发生在 Google 训练 PaLM 模型时），可以延长设备的使用寿命并减少中断。这对于 AI 训练等同步工作具有很大的破坏性，就像 HPC 模拟一样和建模。因此，与其回滚到检查点并从那里开始，不如从一开始就避免中断。
此外，快速更新 TPU 芯片是谷歌工作的重要部分。
Dean 表示，目前芯片的制造大约需要三年时间。这包括六到十二个月的时间进行设计和探索，一年的时间实施设计，六个月的时间与晶圆进行流片，以及十二个月的时间用于投入生产、测试和提升。
目前尚不清楚人工智能可以在多大程度上缩短芯片开发周期或可以减少多少人力，并且 Dean 也没有提供任何估计值。但显然，硬件设计越接近新兴人工智能模型越好。
