Cloud TPU 推理

服务是指将经过训练的机器学习模型部署到生产环境的过程,以便在生产环境中使用该模型进行推理。TPU v5e 及更高版本支持推断。延迟时间 SLO 是服务优先级。

本文档介绍如何在单主机 TPU 上提供模型。包含 8 个或更少芯片的 TPU 切片具有一个 TPU 虚拟机或主机,称为单主机 TPU。

开始使用

您需要有 Google Cloud 账号和项目才能使用 Cloud TPU。如需了解详情,请参阅设置 Cloud TPU 环境

您需要申请以下配额才能在 TPU 上传送数据:

  • 按需 v5e 资源:TPUv5 lite pod cores for serving per project per zone
  • 抢占式 v5e 资源:Preemptible TPU v5 lite pod cores for serving per project per zone
  • 按需 v6e 资源:TPUv6 cores per project per zone
  • 抢占式 v6e 资源:Preemptible TPUv6 cores per project per zone

如需详细了解 TPU 配额,请参阅 TPU 配额

使用 JetStream 为 LLM 提供服务

JetStream 是一个吞吐量和内存优化引擎,用于在 XLA 设备 (TPU) 上进行大语言模型 (LLM) 推断。您可以将 JetStream 与 JAX 和 PyTorch/XLA 模型结合使用。如需查看使用 JetStream 处理 JAX LLM 的示例,请参阅在 v6e TPU 上利用 JetStream MaxText 推断

使用 vLLM 提供 LLM 模型

vLLM 是一个开源库,专为快速推理和提供大型语言模型 (LLM) 而设计。您可以将 vLLM 与 PyTorch/XLA 搭配使用。如需查看使用 vLLM 处理 PyTorch LLM 的示例,请参阅通过 vLLM 在 GKE 上使用 TPU Trillium 提供 LLM

分析

设置推理后,您可以使用分析器来分析性能和 TPU 利用率。如需详细了解性能剖析,请参阅: