谷歌正式发布第六代TPU——Trillium，并将其开放给Google Cloud客户使用。Trillium是谷歌迄今为止性能最强大的TPU，用于训练其最强大的AI模型Gemini 2.0。其显著提升了训练性能、推理吞吐量和能源效率，并实现了更低的成本。 本文将深入探讨Trillium TPU的性能提升、关键功能以及在不同AI工作负载中的出色表现，并展示其在AI21Labs等客户中的实际应用案例。
今年早些时候，谷歌发布了第六代也是迄今为止性能最强大的 TPU——Trillium。今天，Trillium 正式面向 Google Cloud 客户开放使用。
谷歌使用 Trillium TPU 训练了最新的 Gemini2.0，这是谷歌迄今为止最强大的 AI 模型。现在，企业和初创公司都可以利用同样强大、高效和可持续的基础设施。
AI 超级计算机的核心:Trillium TPU
Trillium TPU 是 Google Cloud AI Hypercomputer 的关键组成部分。AI Hypercomputer 是一种突破性的超级计算机架构，它采用性能优化的硬件、开放软件、领先的 ML 框架和灵活的消费模型集成系统。随着 Trillium TPU 的正式推出，谷歌还对 AI Hypercomputer 的开放软件层进行了关键增强，包括优化 XLA 编译器和 JAX、PyTorch 和 TensorFlow 等流行框架，以在 AI 训练、调整和服务方面实现领先的性价比。
此外，使用大规模主机 DRAM（补充高带宽内存或 HBM）的主机卸载等功能可提供更高水平的效率。AI Hypercomputer 使您能够从每个 Jupiter 网络架构中前所未有的超过10万个 Trillium 芯片部署中提取最大价值，该架构具有13Petabits/秒 的双向带宽，能够将单个分布式训练作业扩展到数十万个加速器。
AI21Labs 等客户已经在使用 Trillium，以更快地向其客户交付有意义的 AI 解决方案:
AI21Labs 首席技术官 Barak Lenz表示:“在 AI21，我们不断努力提高 Mamba 和 Jamba 语言模型的性能和效率。作为 TPU v4的长期用户，我们对 Google Cloud 的 Trillium 的功能印象深刻。在规模、速度和成本效率方面的进步非常显著。我们相信 Trillium 将在加速我们下一代复杂语言模型的开发方面发挥至关重要的作用，使我们能够为客户提供更强大和更易于访问的 AI 解决方案。”
Trillium 性能大幅提升，多项指标刷新纪录
与上一代相比，Trillium 在以下方面进行了显著改进:
训练性能提高4倍以上
推理吞吐量提高3倍
能源效率提高67%
每个芯片的峰值计算性能提高4.7倍
高带宽内存 （HBM） 容量翻倍
芯片间互连 （ICI） 带宽翻倍
单个 Jupiter 网络架构中包含10万个 Trillium 芯片
每美元训练性能提高2.5倍，每美元推理性能提高1.4倍
这些增强功能使 Trillium 能够在各种 AI 工作负载中表现出色，包括:
扩展 AI 训练工作负载
训练 LLM，包括密集型模型和混合专家 （MoE） 模型
推理性能和集合调度
嵌入密集型模型
提供训练和推理性价比
Trillium 如何在不同工作负载中表现出色?
扩展 AI 训练工作负载
训练像 Gemini2.0这样的大型模型需要大量的数据和计算。Trillium 的近乎线性的扩展能力使这些模型可以通过在多个 Trillium 主机之间有效且高效地分配工作负载来显著加快训练速度，这些主机通过高速芯片间互连连接在256芯片 pod 和我们最先进的 Jupiter 数据中心网络中。这通过 TPU 多片、用于大规模训练的全栈技术实现，并通过 Titanium 进一步优化，Titanium 是一种动态数据中心级卸载系统，范围从主机适配器到网络架构。
Trillium 在由3072个芯片组成的12个 pod 的部署中实现了99% 的扩展效率，并在具有6144个芯片的24个 pod 中展示了94% 的扩展效率，以预训练 gpt3-175b，即使在跨数据中心网络运行时也是如此。
训练 LLM，包括密集型模型和混合专家 （MoE） 模型
像 Gemini 这样的 LLM 本身就非常强大和复杂，具有数十亿个参数。训练这种密集型 LLM 需要巨大的计算能力以及协同设计的软件优化。Trillium 在训练 Llama-2-70b 和 gpt3-175b 等密集型 LLM 时，速度比上一代 Cloud TPU v5e 快4倍。
除了密集型 LLM 之外，使用混合专家 （MoE） 架构训练 LLM 是一种越来越流行的方法，它结合了多个“专家”神经网络，每个神经网络都专门负责 AI 任务的不同方面。与训练单个整体模型相比，在训练期间管理和协调这些专家增加了复杂性。Trillium 在训练 MoE 模型时的速度比上一代 Cloud TPU v5e 快3.8倍。
此外，与 Cloud TPU v5e 相比，Trillium TPU 提供了3倍的主机动态随机存取存储器 （DRAM）。这会将一些计算卸载到主机，有助于最大限度地提高大规模性能和良好吞吐量。Trillium 的主机卸载功能在训练 Llama-3.1-405B 模型时，在模型 FLOP 利用率 (MFU) 方面提供了超过50% 的性能提升。
推理性能和集合调度
在推理时，多步骤推理的重要性日益提高，这需要加速器能够有效地处理增加的计算需求。Trillium 为推理工作负载提供了重大进步，可以更快、更高效地部署 AI 模型。事实上，Trillium 为图像扩散和密集型 LLM 提供了我们最佳的 TPU 推理性能。我们的测试表明，与 Cloud TPU v5e 相比，Stable Diffusion XL 的相对推理吞吐量（每秒图像数）高出3倍以上，而 Llama2-70B 的相对推理吞吐量(每秒标记数)高出近2倍。
Trillium 是我们用于离线和服务器推理用例的性能最高的 TPU。下图显示，与 Cloud TPU v5e 相比，Stable Diffusion XL 的离线推理相对吞吐量（每秒图像数）高出3.1倍，服务器推理相对吞吐量高出2.9倍。
除了更好的性能之外，Trillium 还引入了新的集合调度功能。此功能允许 Google 的调度系统做出智能的作业调度决策，以提高集合中存在多个副本时推理工作负载的整体可用性和效率。它提供了一种管理运行单主机或多主机推理工作负载的多个 TPU 片段的方法，包括通过 Google Kubernetes Engine （GKE）。将这些片段分组到一个集合中可以轻松地调整副本的数量以匹配需求。
嵌入密集型模型
通过添加第三代 SparseCore，Trillium 使嵌入密集型模型的性能提高了2倍，DLRM DCNv2的性能提高了5倍。
SparseCore 是数据流处理器，它为嵌入密集型工作负载提供了更具适应性的架构基础。Trillium 的第三代 SparseCore 擅长加速动态和数据相关的操作，例如分散收集、稀疏段求和和分区。
提供训练和推理性价比
除了训练世界上一些最大规模 AI 工作负载所需的绝对性能和规模外，Trillium 还旨在优化每美元的性能。迄今为止，Trillium 在训练 Llama2-70b 和 Llama3.1-405b 等密集型 LLM 时，每美元的性能比 Cloud TPU v5e 高出2.1倍，比 Cloud TPU v5p 高出2.5倍。
Trillium 擅长以经济高效的方式并行处理大型模型。它的设计目的是使研究人员和开发人员能够以比以前低得多的成本提供强大而高效的图像模型。在 Trillium 上生成一千张图像的成本，对于离线推理，比 Cloud TPU v5e 低27%，对于 SDXL 上的服务器推理，则比 Cloud TPU v5e 低22%。
将 AI 创新提升到新的水平
Trillium 代表了 Google Cloud AI 基础设施的重大飞跃，为各种 AI 工作负载提供了令人难以置信的性能、可扩展性和效率。借助其使用世界一流的协同设计软件扩展到数十万个芯片的能力，Trillium 使您能够实现更快的突破并交付卓越的 AI 解决方案。此外，Trillium 卓越的性价比使其成为希望最大限度地提高 AI 投资价值的组织的经济高效的选择。随着 AI 格局的不断发展，Trillium 证明了 Google Cloud 致力于提供尖端基础设施，以帮助企业释放 AI 的全部潜力。
官方介绍：https://cloud.google.com/blog/products/compute/trillium-tpu-is-ga
总而言之，Trillium TPU 的出现标志着云端AI计算能力的显著提升，其强大的性能、可扩展性和经济效益将推动AI领域更快的发展，为企业和研究机构提供更强大的AI解决方案。