通过 JetStream 使用 GKE 中的 TPU 应用 Gemma

本教程介绍如何在 Google Kubernetes Engine (GKE) 上使用张量处理单元 (TPU) 来提供 Gemma 大语言模型 (LLM)。您会将包含 JetStreamMaxText 的预构建容器部署到 GKE。您还需要配置 GKE,以便在运行时从 Cloud Storage 加载 Gemma 7B 权重。

本教程适用于机器学习 (ML) 工程师、平台管理员和运维人员,以及对使用 Kubernetes 容器编排功能提供 LLM 感兴趣的数据和 AI 专家。如需详细了解我们在Google Cloud 内容中提及的常见角色和示例任务,请参阅常见的 GKE 用户角色和任务

在阅读本页面之前,请确保您熟悉以下内容:

背景

本部分介绍本教程中使用的关键技术。

Gemma

Gemma 是一组公开提供的轻量级生成式人工智能 (AI) 模型(根据开放许可发布)。这些 AI 模型可以在应用、硬件、移动设备或托管服务中运行。您可以使用 Gemma 模型生成文本,但也可以针对专门任务对这些模型进行调优。

如需了解详情,请参阅 Gemma 文档

TPU

TPU 是 Google 定制开发的应用专用集成电路 (ASIC),用于加速机器学习和使用 TensorFlowPyTorchJAX 等框架构建的 AI 模型。

本教程介绍如何应用 Gemma 7B 模型。GKE 在单主机 TPUv5e 节点上部署模型,并根据模型要求配置 TPU 拓扑,以低延迟提供提示。

JetStream

JetStream 是由 Google 开发的开源推理服务框架。JetStream 可以在 TPU 和 GPU 上实现高性能、高吞吐量和内存优化的推理。它提供高级性能优化(包括连续批处理和量化技术),以协助 LLM 部署。JetStream 支持 PyTorch/XLA 和 JAX TPU 服务,从而实现最佳性能。

如需详细了解这些优化,请参阅 JetStream PyTorchJetStream MaxText 项目仓库。

MaxText

MaxText是一个高性能、可扩缩且适应性强的 JAX LLM 实现,基于如下开源 JAX 仓库构建:FlaxOrbaxOptax。MaxText 的仅解码器 LLM 实现是使用 Python 编写的。它大量利用 XLA 编译器来实现高性能,而无需构建自定义内核。

如需详细了解 MaxText 支持的最新模型和参数大小,请参阅 MaxtText 项目仓库