利用 Google Cloud TPU 和 GPU 加快 AI 推理速度

在发展迅猛的人工智能领域中，人们对高性能、高性价比 AI 推理（服务）的需求从未如此之高。本周我们宣布推出两款全新开源软件产品：JetStream 和 MaxDiffusion。

JetStream 是针对 XLA 设备（首先是 Cloud TPU）的新推理引擎。它专为大语言模型 (LLM) 设计，可大幅提高性能和成本效益。在费用相同的情况下，其 LLM 推理效率是之前 Cloud TPU 推理引擎的 3 倍。JetStream 通过 PyTorch/XLA 支持 PyTorch 模型，并通过 MaxText（这是我们针对 LLM 提供的高度可伸缩的高性能参考实现，客户可以创建分支，从而加快开发速度）支持 JAX 模型。

MaxDiffusion 与 MaxText 类似，不同之处在于它适用于潜在的 diffusion 模型，可帮助您在 XLA 设备（首先是 Cloud TPU）上轻松训练和提供针对高性能进行了优化的 diffusion 模型。

此外，我们很开心在此分享 MLPerf™ Inference v4.0 的最新性能结果，这些结果充分展示了搭载 NVIDIA H100 GPU 的 Google Cloud A3 虚拟机 (VM) 的强大功能和广泛用途。

JetStream：经济高效的高性能 LLM 推理

LLM 处于 AI 革新的前沿，广泛应用于自然语言理解、文本生成和语言翻译等领域。为了降低客户的 LLM 推理费用，我们打造了推理引擎 JetStream：在费用相同的情况下，其 LLM 推理效率是之前 Cloud TPU 推理引擎的 3 倍。

图 1：JetStream 堆栈。

JetStream 包括先进的性能优化功能，例如连续批处理、滑动窗口注意力 (sliding window attention) 以及用于权重、激活值和键值对 (KV) 缓存的 int8 量化。无论您使用 JAX 还是 PyTorch，JetStream 都可以支持您喜欢的框架。为了进一步简化您的 LLM 推理工作流，我们提供常用开放模型（如 Gemma 和 Llama）的 MaxText 和 PyTorch/XLA 实现，它们经过优化，可实现超高的成本效益与性能。

在 Cloud TPU v5e-8 上，JetStream 可为开放模型（包括 MaxText 中的 Gemma 和 PyTorch/XLA 中的 Llama 2）提供高达每秒 4783 个词元的吞吐量。

图 2：JetStream 吞吐量（每秒输出词元数）。Google 内部数据。在 Cloud TPU v5e-8 上使用 Gemma 7B (MaxText)、Llama 2 7B (PyTorch/XLA) 和 Llama 2 13B (PyTorch/XLA) 进行测量。输入长度上限：1024；输出长度上限：1024。连续批量处理，用于权重、激活值、KV 缓存的 int8 量化。PyTorch/XLA 采用滑动窗口注意力。截至 2024 年 4 月的数据。

JetStream 的高性能和高效率意味着为 Google Cloud 客户降低了推理费用，使得 LLM 推理价格更亲民，能够让更多人使用。

图 3：生成 100 万个输出词元的 JetStream 费用。Google 内部数据。在 Cloud TPU v5e-8 上使用 Gemma 7B (MaxText)、Llama 2 7B (PyTorch/XLA) 和 Llama 2 13B (PyTorch/XLA) 进行测量。输入长度上限：1024；输出长度上限：1024。连续批量处理，用于权重、激活值、KV 缓存的 int8 量化。PyTorch/XLA 采用滑动窗口注意力。在运行 Gemma 7B 时，与之前的 Cloud TPU LLM 推理堆栈（每 100 万词元的费用为 1.10 美元）相比，JetStream（每 100 万词元的费用为 0.30 美元）LLM 推理的成本效益提高了 3 倍。费用基于美国境内 Cloud TPU v5e-8 的 3 年期承诺使用折扣价格。截至 2024 年 4 月的数据。

Osmos 等客户使用 JetStream 来加快其 LLM 推理工作负载的运行速度：

“在 Osmos，我们开发了依托 AI 技术的数据转换引擎，帮助很多公司通过自动化数据处理来扩展其业务关系。从客户和业务合作伙伴那里收到的数据通常是杂乱且不标准的，因此需要对每一行数据进行智能处理，对其进行映射、验证并转换为良好的可用数据。这需要高性能、可伸缩且经济高效的 AI 基础设施来进行训练、微调和推理。为此，我们为自己的端到端 AI 工作流选择搭配了 MaxText、JAX 和 JetStream 的 Cloud TPU v5e 方案。借助 Google Cloud，我们能够使用 MaxText 在数十亿个词元的基础上快速轻松地微调 Google 最新的 Gemma 开放式模型，并使用 JetStream 部署该模型以进行推理，而这一切都是在 Cloud TPU v5e 上完成的。得益于 Google 经优化的 AI 硬件和软件堆栈，我们能够在几小时而不是几天内取得成果。”

– Kirat Pandya，Osmos 首席执行官

Google 为研究人员和开发者提供强大且经济高效的 LLM 推理开源基础，为下一代 AI 应用提供可靠助力。无论您是经验丰富的 AI 从业者，还是刚刚开始接触 LLM，JetStream 都能加快您的学习历程并在自然语言处理领域发掘新的可能。

快来使用 JetStream 体验未来的 LLM 推理吧！访问我们的 GitHub 代码库，详细了解 JetStream 并开始创建新的 LLM 项目。Google 将在 GitHub 上及通过 Google Cloud Customer Care 长期开发和支持 JetStream，还会邀请社区成员一起，共同进行构建和优化，进一步推动这项先进技术的发展。

MaxDiffusion：高性能 diffusion 模型推理

正如 LLM 彻底改变了自然语言处理一样，diffusion 模型也在改变着计算机视觉领域。为了降低客户部署这些模型的成本，我们开发了 MaxDiffusion，这是一系列开源 diffusion 模型参考实现的集合。这些实现采用 JAX 编写，具有高性能、高可伸缩性和可定制性。您可以将 MaxDiffusion 视为 MaxText 的计算机视觉版本。 

MaxDiffusion 提供 diffusion 模型核心组件的高性能实现，例如交叉注意力、卷积和高吞吐量图片数据加载。MaxDiffusion 在设计上保证了高适应性和可定制性，无论您是希望突破图片生成界限的研究人员，还是希望将尖端的生成式 AI 功能集成到应用中的开发者，MaxDiffusion 都能为您提供成功所需的基础。

新 SDXL-Lightning 模型的 MaxDiffusion 实现在 Cloud TPU v5e-4 上每秒可生成 6 张图片，在 Cloud TPU v5e-8 上的吞吐量可线性扩展至每秒 12 张图片，充分利用了 Cloud TPU 的高性能和可伸缩性。

图 4：MaxDiffusion 吞吐量（图片数/秒）。Google 内部数据。使用 Cloud TPU v5e-4 和 Cloud TPU v5e-8 上的 SDXL-Lightning 模型进行测量。分辨率：1024x1024；每个设备的批次大小：2；解码步数：4。截至 2024 年 4 月的数据。

与 MaxText 和 JetStream 一样，MaxDiffusion 也极具成本效益：在 Cloud TPU v5e-4 或 Cloud TPU v5e-8 上生成 1000 张图片仅需 0.10 美元。

图 5：MaxDiffusion 生成 1000 张图片的费用。Google 内部数据。使用 Cloud TPU v5e-4 和 Cloud TPU v5e-8 上的 SDXL-Lightning 模型进行测量。分辨率：1024x1024；每个设备的批次大小：2；解码步数：4。费用基于美国境内 Cloud TPU v5e-4 和 Cloud TPU v5e-8 的 3 年期承诺使用折扣价格。截至 2024 年 4 月的数据。

Codeway 等许多客户采用 Google Cloud 来最大限度地提高大规模 diffusion 模型推理的成本效益。

“Codeway 开发了多款受人追捧的热门应用和游戏，在全球 160 个国家/地区拥有超过 1.15 亿用户。其中，《Wonder》是一款依托 AI 技术的应用，可以根据文字创建数字艺术作品；而《Facedance》可以通过各种有趣的动画让您的脸部“跳舞”。将 AI 交到数百万用户手中需要高度可伸缩且经济高效的推理基础设施。与其他推理解决方案相比，使用 Cloud TPU v5e 提供 diffusion 模型时，我们能够将处理时间减少 45%，将每小时处理的请求数增加 3.6 倍。就我们的规模而言，这大幅降低了基础设施成本，让我们能够经济高效地将依托 AI 技术的应用带给更多用户。”

– Uğur Arpacı，Codeway DevOps 负责人

MaxDiffusion 为图片生成提供高性能、高度可伸缩且灵活的基础。无论您是经验丰富的计算机视觉专家，还是刚刚涉足图片生成领域的新手，MaxDiffusion 都可以全程为您提供支持。

立即访问我们的 GitHub 代码库，详细了解 MaxDiffusion 并开始构建新的创意项目。

A3 虚拟机：MLPerf™ Inference 4.0 结果出众

我们在 2023 年 8 月推出了 A3 虚拟机系列正式版。A3 系列在单个虚拟机中使用 8 个 NVIDIA H100 Tensor Core GPU，专为训练和提供要求苛刻的生成式 AI 工作负载及 LLM 而构建。搭载 NVIDIA H100 GPU 的 A3 Mega 于下个月全面推出，将提供双倍于 A3 的 GPU 到 GPU 网络带宽。

在 MLPerf™ Inference v4.0 基准测试中，Google 使用 A3 虚拟机提交了 7 个不同模型的 20 个结果，包括新的 Stable Diffusion XL 和 Llama 2 (70B) 基准。

• RetinaNet（服务器和离线）

• 3D U-Net：99% 和 99.9% 准确率（离线）

• BERT：99% 和 99% 准确率（服务器和离线）

• DLRM v2：99.9% 准确率（服务器和离线）

• GPT-J：99% 和 99% 准确率（服务器和离线）

• Stable Diffusion XL（服务器和离线）

• Llama 2：99% 和 99% 准确率（服务器和离线）

所有结果均与 NVIDIA 提交的结果中展示的峰值性能相差 0-5% 以内。这些结果表明，Google Cloud 与 NVIDIA 密切合作，打造出了专用于 LLM 和生成式 AI 且针对工作负载进行了优化的端到端解决方案。

借助 Google Cloud TPU 和 NVIDIA GPU 助力未来 AI 持续发展

依托 Google Cloud TPU 和 NVIDIA GPU 的硬件进步以及 JetStream、MaxText 和 MaxDiffusion 等软件创新，Google 在 AI 推理方面的创新助力客户构建和扩展 AI 应用。利用 JetStream，开发者能够进一步提升 LLM 推理的性能和成本效益，在自然语言处理应用领域发掘出新的机遇。MaxDiffusion 为研究人员和开发者提供坚实基础，让他们可以深入探索 diffusion 模型的潜力并加快图片生成速度。配备 NVIDIA H100 Tensor Core GPU 的 A3 虚拟机在 MLPerf™ Inference 4.0 中表现出众，证明了 Cloud GPU 的强大功能和广泛用途。 

立即访问我们的网站了解详情，并开始使用 Google Cloud TPU 和 GPU 推理。