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 利用率。如需详细了解性能剖析,请参阅: