谷歌正式发布第六代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解决方案。