Cloud TPU v5e 推断简介

概览和优势

Cloud TPU v5e 是 Google 开发的 AI 加速器，针对基于转换器、文本到图像和基于 CNN 的训练、微调和服务（推断）进行了优化。TPU v5e 切片最多可包含 256 个芯片。

服务是指将经过训练的机器学习模型部署到生产环境的过程，该模型可用于推断。延迟时间 SLO 是服务优先级。

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

开始使用

您将需要 v5e TPU 的配额。按需 TPU 需要 tpu-v5s-litepod-serving 配额。预留 TPU 需要 tpu-v5s-litepod-serving-reserved 配额。如需了解详情，请联系 Cloud 销售人员。

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

您可以使用已加入队列的资源预配 v5e TPU。如需详细了解可用于提供服务的 v5e 配置，请参阅提供服务的 Cloud TPU v5e 类型。

Cloud TPU 模型推断和服务

如何提供模型进行推理取决于编写模型时使用的机器学习框架。TPU v5e 支持使用 JAX、TensorFlow 和 PyTorch 编写的模型。

JAX 模型推断和服务

如需在 TPU 虚拟机上提供模型，您需要执行以下操作：

以 TensorFlow SavedModel 格式序列化模型
使用推断转换器准备保存的模型以提供服务
使用 TensorFlow Serving 应用模型

SavedModel 格式

SavedModel 包含完整的 TensorFlow 程序，其中包括经过训练的参数和计算。它不需要运行原始模型构建代码。

如果模型是用 JAX 编写的，您需要使用 jax2tf 以 SavedModel 格式序列化模型。

推理转换器

Cloud TPU Inference Converter 可准备和优化以 SavedModel 格式导出的模型，以便进行 TPU 推断。您可以在本地 shell 或 TPU 虚拟机中运行推理转换器。我们建议您使用 TPU 虚拟机 shell，因为它包含运行转换器所需的所有命令行工具。如需详细了解推断转换器，请参阅推断转换器用户指南。

推断转换器要求

您的模型必须以 SavedModel 格式从 TensorFlow 或 JAX 中导出。
您必须为 TPU 函数定义函数别名。如需了解详情，请参阅推断转换器用户指南。本指南中的示例使用 tpu_func 作为 TPU 函数别名。
确保您的机器 CPU 支持高级矢量扩展 (AVX) 指令，因为已编译 TensorFlow 库（Cloud TPU 推断转换器的依赖项）以使用 AVX 指令。大多数 CPU 都支持 AVX。

注意：您可以运行 lscpu | grep avx 来检查设备是否支持 AVX 指令集。

JAX 模型推断和服务

本部分介绍如何使用 jax2tf 和 TensorFlow Serving 提供 JAX 模型。

使用 jax2tf 将模型序列化为 SavedModel 格式
使用推断转换器准备已保存的模型，以便提供服务
使用 TensorFlow Serving 应用模型

使用 `jax2tf` 将 JAX 模型序列化为 SavedModel 格式

以下 Python 函数展示了如何在模型代码中使用 jax2tf：

# Inference function
def model_jax(params, inputs):
  return params[0] + params[1] * inputs

# Wrap the parameter constants as tf.Variables; this will signal to the model
# saving code to save those constants as variables, separate from the
# computation graph.
params_vars = tf.nest.map_structure(tf.Variable, params)

# Build the prediction function by closing over the `params_vars`. If you
# instead were to close over `params` your SavedModel would have no variables
# and the parameters will be included in the function graph.
prediction_tf = lambda inputs: jax2tf.convert(model_jax)(params_vars, inputs)

my_model = tf.Module()
# Tell the model saver what the variables are.
my_model._variables = tf.nest.flatten(params_vars)
my_model.f = tf.function(prediction_tf, jit_compile=True, autograph=False)
tf.saved_model.save(my_model)

如需详细了解 jax2tf，请参阅 JAX 和 Cloud TPU 互操作。

使用推断转换器准备保存的模型以提供服务

推断转换器指南中介绍了推断转换器的使用说明。

使用 TensorFlow Serving

TensorFlow Serving 中介绍了 TensorFlow Serving 的使用说明。

JAX 模型服务示例

前提条件

设置 Docker 凭据，并拉取 Inference Converter 和 Cloud TPU Serving Docker 映像：

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
   us-docker.pkg.dev
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

使用 SSH 连接到您的 TPU 虚拟机并安装推理演示代码：
```
gsutil -m cp -r \
"gs://cloud-tpu-inference-public/demo" \
.
```

安装 JAX 演示依赖项：

pip install -r ./demo/jax/requirements.txt

提供 JAX BERT 模型进行推理

您可以从 Hugging Face 下载预训练 BERT 模型。

从 Flax BERT 模型中导出与 TPU 兼容的 TensorFlow 保存模型：
```
cd demo/jax/bert
python3 export_bert_model.py
```

启动 Cloud TPU 模型服务器容器：

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/jax/bert_tpu,target=/models/bert \
  -e MODEL_NAME=bert \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

容器启动大约 30 秒后，检查模型服务器容器日志，确保 gRPC 和 HTTP 服务器已启动：

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

如果您看到以以下信息结尾的日志条目，则表示服务器已准备好处理请求。

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

向模型服务器发送推断请求。

python3 bert_request.py

输出将如下所示：

For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."

清理。

在运行其他演示之前，请务必清理 Docker 容器。

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}

清理模型工件：

sudo rm -rf /tmp/jax/

提供 JAX Stable Diffusion 以进行推理

您可以从 Hugging Face 下载预训练的 Stable Diffusion 模型。

以与 TPU 兼容的 TF2 已保存的模型格式下载 Stable Diffusion 模型：

注意：下载大约需要 3 分钟。
```
cd demo/jax/stable_diffusion
python3 export_stable_diffusion_model.py
```

启动模型的 Cloud TPU 模型服务器容器：

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/jax/stable_diffusion_tpu,target=/models/stable_diffusion \
  -e MODEL_NAME=stable_diffusion \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

大约两分钟后，检查模型服务器容器日志，确保 gRPC 和 HTTP 服务器正在运行：

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

如果您看到以以下信息结尾的日志，则表示服务器已准备好处理请求。

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

向模型服务器发送请求。
```
python3 stable_diffusion_request.py
```
此脚本发送“Painting of a squirrel skating in New York”（纽约在纽约滑冰的松鼠）作为提示。输出图片将作为 stable_diffusion_images.jpg 保存到当前目录中。

清理。

在运行其他演示之前，请务必清理 Docker 容器。

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}

清理模型工件

sudo rm -rf /tmp/jax/

TensorFlow Serving

以下说明演示了如何在 TPU 虚拟机上提供 TensorFlow 模型。

TensorFlow 服务工作流

下载适用于您的 TPU 虚拟机的 TensorFlow Serving Docker 映像。

设置示例环境变量

export YOUR_LOCAL_MODEL_PATH=model-path
export MODEL_NAME=model-name
# Note: this image name may change later.
export IMAGE_NAME=us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

下载 Docker 映像

docker pull ${IMAGE_NAME}

设置 Docker 凭据，并拉取 Inference Converter 和 TensorFlow Serving Docker 映像。

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
   us-docker.pkg.dev
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

下载演示代码：

gsutil -m cp -r \
"gs://cloud-tpu-inference-public/demo" \
.

安装 TensorFlow 演示版依赖项：

pip install -r ./demo/tf/requirements.txt

在 TPU 虚拟机上使用 TensorFlow Serving Docker 映像提供 TensorFlow 模型。

# PORT 8500 is for gRPC model server and 8501 is for HTTP/REST model server.
docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=${YOUR_LOCAL_MODEL_PATH},target=/models/${MODEL_NAME} \
  -e MODEL_NAME=${MODEL_NAME} \
  ${IMAGE_NAME}

使用 Serving Client API 来查询模型。
- REST 客户端 API
- gRPC 客户端 API

运行 TensorFlow ResNet-50 Serving 演示

从 Keras ResNet-50 模型中导出与 TPU 兼容的 TF2 保存模型。
```
cd demo/tf/resnet-50
python3 export_resnet_model.py
```

启动模型的 TensorFlow 模型服务器容器。

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/tf/resnet_tpu,target=/models/resnet \
  -e MODEL_NAME=resnet \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

检查模型服务器容器日志，确保 gRPC 和 HTTP Server 已启动：

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

如果您看到结尾为以下信息的日志，则表示服务器已准备好处理请求。此过程大约需要 30 秒。

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

向模型服务器发送请求。

请求图片是 https://i.imgur.com/j9xCCzn.jpeg 中的香蕉。

python3 resnet_request.py

输出将如下所示：

Predict result: [[('n07753592', 'banana', 0.94921875), ('n03532672', 'hook', 0.022338867), ('n07749582', 'lemon', 0.005126953)]]

清理。

在运行其他演示之前，请务必清理 Docker 容器。

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}

清理模型工件：

sudo rm -rf /tmp/tf/

PyTorch 模型推断和服务

对于使用 PyTorch 编写的模型，工作流为：

使用 TorchDynamo 和 PyTorch/XLA 编写用于加载和推断的 Python 模型处理程序
使用 TorchModelArchiver 创建模型归档文件
使用 TorchServe 提供模型服务

TorchDynamo 和 PyTorch/XLA

TorchDynamo (Dynamo) 是一种 Python 级 JIT 编译器，旨在加快 PyTorch 程序的运行速度。它提供了一个干净的 API，供编译器后端接入。它会在执行之前动态修改 Python 字节码。在 PyTorch/XLA 2.0 版本中，有一个实验性后端，用于使用 Dynamo 进行推断和训练。

Dynamo 在识别出模型模式，并且 PyTorch/XLA 使用延迟张量方法编译 FX 图并返回编译后的函数时，会提供 Torch FX (FX) 图。如需详细了解 Dynamo，请参阅：

下面是一个使用 torch.compile 运行 densenet161 推断的小型代码示例。

import torch
import torchvision
import torch_xla.core.xla_model as xm

def eval_model(loader):
  device = xm.xla_device()
  xla_densenet161 = torchvision.models.densenet161().to(device)
  xla_densenet161.eval()
  dynamo_densenet161 = torch.compile(
      xla_densenet161, backend='torchxla_trace_once')
  for data, _ in loader:
    output = dynamo_densenet161(data)

TorchServe

您可以使用提供的 torchserve-tpu Docker 映像在 TPU 虚拟机上提供已归档的 pytorch 模型。

为 Docker 设置身份验证：

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
    us-docker.pkg.dev

将 Cloud TPU TorchServe Docker 映像拉取到您的 TPU 虚拟机：

CLOUD_TPU_TORCHSERVE_IMAGE_URL=us-docker.pkg.dev/cloud-tpu-images/inference/torchserve-tpu:v0.9.0-2.1
docker pull ${CLOUD_TPU_TORCHSERVE_IMAGE_URL}

收集模型工件

首先，您需要提供一个模型处理程序，用于指示 ToorchServe 模型服务器工作器加载模型、处理输入数据并运行推断。您可以使用 TorchServe 默认推断处理程序（来源），也可以按照 base_handler.py 开发自己的自定义模型处理程序。您可能还需要提供经过训练的模型和模型定义文件。

在下面的 Densenet 161 示例中，我们使用了模型工件和 TorchServe 提供的默认图片分类器处理程序：

配置一些环境变量：

CWD="$(pwd)"

WORKDIR="${CWD}/densenet_161"

mkdir -p ${WORKDIR}/model-store
mkdir -p ${WORKDIR}/logs

从 TorchServe 图片分类器示例中下载和复制模型工件：

git clone https://github.com/pytorch/serve.git

cp ${CWD}/serve/examples/image_classifier/densenet_161/model.py ${WORKDIR}
cp ${CWD}/serve/examples/image_classifier/index_to_name.json ${WORKDIR}

下载模型权重：

wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth

mv densenet161-8d451a50.pth ${WORKDIR}

创建 TorchServe 模型配置文件以使用 Dynamo 后端：

echo 'pt2: "torchxla_trace_once"' >> ${WORKDIR}/model_config.yaml

您应该会看到以下文件和目录：

>> ls ${WORKDIR}
model_config.yaml
index_to_name.json
logs
model.py
densenet161-8d451a50.pth
model-store

生成模型归档文件

如需使用 Cloud TPU TorchServe 提供 PyTorch 模型，您需要使用 Torch Model Archiver 将模型处理程序和所有模型工件打包到模型归档文件 (*.mar) 中。

使用 torch-model-archiver 生成模型归档文件：

MODEL_NAME=Densenet161

docker run \
    --privileged  \
    --shm-size 16G \
    --name torch-model-archiver \
    -it \
    -d \
    --rm \
    --mount type=bind,source=${WORKDIR},target=/home/model-server/ \
    ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
    torch-model-archiver \
        --model-name ${MODEL_NAME} \
        --version 1.0 \
        --model-file model.py \
        --serialized-file densenet161-8d451a50.pth \
        --handler image_classifier \
        --export-path model-store \
        --extra-files index_to_name.json \
        --config-file model_config.yaml

您应该会看到在 model-store 目录中生成的模型归档文件：

>> ls ${WORKDIR}/model-store
Densenet161.mar

处理推理请求

现在，您已经有了模型归档文件，接下来可以启动 TorchServe 模型服务器并处理推理请求。

启动 TorchServe 模型服务器：

docker run \
    --privileged  \
    --shm-size 16G \
    --name torchserve-tpu \
    -it \
    -d \
    --rm \
    -p 7070:7070 \
    -p 7071:7071 \
    -p 8080:8080 \
    -p 8081:8081 \
    -p 8082:8082 \
    -p 9001:9001 \
    -p 9012:9012 \
    --mount type=bind,source=${WORKDIR}/model-store,target=/home/model-server/model-store \
    --mount type=bind,source=${WORKDIR}/logs,target=/home/model-server/logs \
    ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
    torchserve \
        --start \
        --ncs \
        --models ${MODEL_NAME}.mar \
        --ts-config /home/model-server/config.properties

查询模型服务器运行状况：

curl http://localhost:8080/ping

如果模型服务器已启动并运行，您会看到以下内容：

{
  "status": "Healthy"
}

如需查询当前已注册模型的默认版本，请使用：

curl http://localhost:8081/models

您应该会看到已注册的模型：

{
  "models": [
    {
      "modelName": "Densenet161",
      "modelUrl": "Densenet161.mar"
    }
  ]
}

如需下载图片进行推理，请使用：

curl -O https://raw.githubusercontent.com/pytorch/serve/master/docs/images/kitten_small.jpg

mv kitten_small.jpg ${WORKDIR}

如需向模型服务器发送推理请求，请使用以下代码：

curl http://localhost:8080/predictions/${MODEL_NAME} -T ${WORKDIR}/kitten_small.jpg

您应该会看到如下所示的响应：

{
  "tabby": 0.47878125309944153,
  "lynx": 0.20393909513950348,
  "tiger_cat": 0.16572578251361847,
  "tiger": 0.061157409101724625,
  "Egyptian_cat": 0.04997897148132324
}

模型服务器日志

请使用以下命令访问日志：

ls ${WORKDIR}/logs/
cat ${WORKDIR}/logs/model_log.log

您应该会在日志中看到以下消息：

"Compiled model with backend torchxla\_trace\_once"

清理

停止 Docker 容器：

rm -rf serve
rm -rf ${WORKDIR}

docker stop torch-model-archiver
docker stop torchserve-tpu

分析

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

Cloud TPU v5e 推断简介

概览和优势

开始使用

Cloud TPU 模型推断和服务

JAX 模型推断和服务

SavedModel 格式

推理转换器

推断转换器要求

JAX 模型推断和服务

使用 jax2tf 将 JAX 模型序列化为 SavedModel 格式

使用推断转换器准备保存的模型以提供服务

使用 TensorFlow Serving

JAX 模型服务示例

前提条件

提供 JAX BERT 模型进行推理

提供 JAX Stable Diffusion 以进行推理

TensorFlow Serving

TensorFlow 服务工作流

运行 TensorFlow ResNet-50 Serving 演示

PyTorch 模型推断和服务

TorchDynamo 和 PyTorch/XLA

TorchServe

收集模型工件

生成模型归档文件

处理推理请求

清理

分析

使用 `jax2tf` 将 JAX 模型序列化为 SavedModel 格式