Introducción a la inferencia de Cloud TPU v5e

Nota: Si es la primera vez que usas Cloud TPU, te recomendamos que leas la Introducción a Cloud TPU.

Descripción general y beneficios

Cloud TPU v5e es un acelerador de IA desarrollado por Google que está optimizado para el entrenamiento, el ajuste y la entrega (inferencia) basados en transformadores, texto a imagen y CNN. Las porciones de TPU v5e pueden contener hasta 256 chips.

La publicación se refiere al proceso de implementar un modelo de aprendizaje automático entrenado en un entorno de producción, donde se puede usar para la inferencia. Los SLO de latencia son una prioridad para la publicación.

En este documento, se analiza la publicación de un modelo en una TPU de host único. Las porciones de TPU con 8 chips o menos tienen una VM o un host de TPU y se denominan TPU de host único.

Comenzar

Necesitarás una cuota para las TPU v5e. Las TPUs on demand requieren una cuota de tpu-v5s-litepod-serving. Las TPU reservadas requieren una cuota de tpu-v5s-litepod-serving-reserved. Para obtener más información, comunícate con Ventas de Cloud.

Necesitarás una cuenta y un proyecto de Google Cloud para usar Cloud TPU. Para obtener más información, consulta Configura un entorno de Cloud TPU.

Las TPU v5e se aprovisionan con recursos en cola. Para obtener más información sobre las configuraciones v5e disponibles para la publicación, consulta Tipos de Cloud TPU v5e para la publicación.

Inferencia y entrega de modelos de Cloud TPU

La forma en que se entrega un modelo para la inferencia depende del framework de AA con el que se escribió el modelo. La TPU v5e admite la publicación de modelos escritos en JAX, TensorFlow y PyTorch.

Inferencia y entrega de modelos de JAX

Para entregar un modelo en una VM de TPU, debes hacer lo siguiente:

1. Serializa tu modelo en el formato SavedModel de TensorFlow
2. Usa el convertidor de inferencia para preparar el modelo guardado para la entrega
3. Usa TensorFlow Serving para entregar el modelo

Formato de SavedModel

Un modelo guardado contiene un programa de TensorFlow completo, incluidos los parámetros y el procesamiento entrenados. No requiere el código de compilación del modelo original para ejecutarse.

Si tu modelo se escribió en JAX, deberás usar jax2tf para serializarlo en el formato de SavedModel.

Conversor de inferencia

El convertidor de inferencia de Cloud TPU prepara y optimiza un modelo exportado en formato SavedModel para la inferencia de TPU. Puedes ejecutar el convertidor de inferencia en un shell local o en tu VM de TPU. Te recomendamos que uses la shell de la VM de TPU, ya que tiene todas las herramientas de línea de comandos necesarias para ejecutar el convertidor. Para obtener más información sobre el convertidor de inferencia, consulta la Guía del usuario del convertidor de inferencia.

Requisitos del convertidor de inferencia

1. Tu modelo se debe exportar desde TensorFlow o JAX en el formato SavedModel.

2. Debes definir un alias de función para la función de TPU. Para obtener más información, consulta la Guía del usuario de Inference Converter. En los ejemplos de esta guía, se usa tpu_func como el alias de la función de TPU.

3. Asegúrate de que la CPU de tu máquina admita instrucciones de extensiones vectoriales avanzadas (AVX), ya que la biblioteca de TensorFlow (la dependencia del convertidor de inferencia de Cloud TPU) se compila para usar instrucciones AVX. La mayoría de las CPUs admiten AVX.

   Nota: Puedes ejecutar lscpu | grep avx para verificar si el conjunto de instrucciones AVX es compatible.

Inferencia y entrega de modelos de JAX

En esta sección, se describe cómo entregar modelos de JAX con jax2tf y TensorFlow Serving.

1. Usa jax2tf para serializar tu modelo en el formato de modelo guardado
2. Usa el convertidor de inferencia para preparar tu modelo guardado para la entrega
3. Usa TensorFlow Serving para entregar el modelo

Usa jax2tf para serializar un modelo de JAX al formato SavedModel

En la siguiente función de Python, se muestra cómo usar jax2tf en el código de tu modelo:

# 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)

Para obtener más información sobre jax2tf, consulta Interoperabilidad entre JAX y Cloud TPU.

Usa el convertidor de inferencia para preparar el modelo guardado para la entrega

Las instrucciones para usar el convertidor de inferencia se describen en la guía del convertidor de inferencia.

Usa TensorFlow Serving

Las instrucciones para usar TensorFlow Serving se describen en Publicación de TensorFlow.

Ejemplos de entrega de modelos de JAX

Requisitos previos

1. Configura tus credenciales de Docker y extrae la imagen de Docker del convertidor de inferencia y de la entrega de Cloud TPU:

   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
   

2. Conéctate a tu VM de TPU con SSH y, luego, instala el código de demostración de inferencia:

   gcloud storage cp \
   "gs://cloud-tpu-inference-public/demo" \
   . \
   --recursive
   

3. Instala las dependencias de la demostración de JAX:

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

Entrega el modelo BERT de JAX para la inferencia

Puedes descargar el modelo BERT previamente entrenado de Hugging Face.

Nota: No se admiten las exportaciones de modelos guardados de PyTorch.

1. Exporta un modelo guardado de TensorFlow compatible con TPU desde un modelo de BERT de Flax:

   cd demo/jax/bert
   python3 export_bert_model.py
   

2. Inicia el contenedor del servidor de modelos de 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
   

   Aproximadamente 30 segundos después de que se inicie el contenedor, verifica el registro del contenedor del servidor de modelos y asegúrate de que los servidores gRPC y HTTP estén activos:

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

   Si ves una entrada de registro que finaliza con la siguiente información, el servidor está listo para entregar solicitudes.

   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 ...
   

3. Envía una solicitud de inferencia al servidor de modelos.

   python3 bert_request.py
   

   El resultado será similar al siguiente ejemplo:

   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?."
   

4. Realizar una limpieza

   Asegúrate de limpiar el contenedor de Docker antes de ejecutar otras demostraciones.

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

   Limpia los artefactos del modelo:

   sudo rm -rf /tmp/jax/
   

Entrega la difusión estable de JAX para la inferencia

Puedes descargar el modelo de dispersión estable previamente entrenado de Hugging Face.

1. Descarga el modelo de difusión estable en un formato de modelo guardado de TF2 compatible con TPU:

   Nota: La descarga tarda unos 3 minutos.

   cd demo/jax/stable_diffusion
   python3 export_stable_diffusion_model.py
   

2. Inicia el contenedor del servidor de modelos de Cloud TPU para el modelo:

   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
   

   Después de unos dos minutos, verifica el registro del contenedor del servidor de modelos para asegurarte de que los servidores gRPC y HTTP se estén ejecutando:

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

   Si ves que el registro finaliza con la siguiente información, significa que los servidores están listos para entregar solicitudes.

   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 ...
   

3. Envía una solicitud al servidor de modelos.

   python3 stable_diffusion_request.py
   

   Esta secuencia de comandos envía "Pintura de una ardilla patinando en Nueva York" como instrucción. La imagen de salida se guardará como stable_diffusion_images.jpg en tu directorio actual.

4. Realizar una limpieza

   Asegúrate de limpiar el contenedor de Docker antes de ejecutar otras demostraciones.

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

   Limpia los artefactos del modelo

   sudo rm -rf /tmp/jax/
   

TensorFlow Serving

En las siguientes instrucciones, se muestra cómo puedes entregar tu modelo de TensorFlow en VMs de TPU.

Flujo de trabajo de TensorFlow Serving

1. Descarga la imagen de Docker de TensorFlow Serving para tu VM de TPU.

   Configura variables de entorno de muestra

   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

   Descarga la imagen de Docker

   docker pull ${IMAGE_NAME}
   

2. Configura las credenciales de Docker y extrae la imagen de Docker del convertidor de inferencia y de TensorFlow Serving.

   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
   

3. Descarga el código de demostración:

   gcloud storage cp \
   "gs://cloud-tpu-inference-public/demo" \
   . \
   --recursive
   

4. Instala las dependencias de la demostración de TensorFlow:

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

5. Publica tu modelo de TensorFlow con la imagen de Docker de TensorFlow Serving en tu VM de TPU.

   # 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}
   

6. Usa la API de cliente de publicación para consultar tu modelo.

   • API de REST Client
   • API cliente de gRPC

Ejecuta la demostración de TensorFlow ResNet-50 Serving

1. Exporta un modelo guardado de TF2 compatible con TPU desde el modelo ResNet-50 de Keras.

   cd demo/tf/resnet-50
   python3 export_resnet_model.py
   

2. Inicia el contenedor del servidor de modelos de TensorFlow para el modelo.

   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
   

   Verifica el registro del contenedor del servidor de modelos y asegúrate de que el servidor gRPC y HTTP estén activos:

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

   Si ves que el registro finaliza con la siguiente información, significa que el servidor está listo para entregar solicitudes. El proceso demora alrededor de 30 segundos.

   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 ...
   

3. Envía la solicitud al servidor de modelos.

   La imagen de la solicitud es un plátano de https://i.imgur.com/j9xCCzn.jpeg .

   python3 resnet_request.py
   

   El resultado será similar al siguiente ejemplo:

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

4. Realizar una limpieza

   Asegúrate de limpiar el contenedor de Docker antes de ejecutar otras demostraciones.

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

   Limpia los artefactos del modelo:

   sudo rm -rf /tmp/tf/
   

Inferencia y entrega de modelos de PyTorch

Para los modelos escritos con PyTorch, el flujo de trabajo es el siguiente:

1. Escribe un controlador de modelos de Python para la carga y la inferencia con TorchDynamo y PyTorch/XLA
2. Usa TorchModelArchiver para crear un archivo de modelo
3. Usa TorchServe para entregar el modelo

TorchDynamo y PyTorch/XLA

TorchDynamo (Dynamo) es un compilador JIT a nivel de Python diseñado para que los programas de PyTorch sean más rápidos. Proporciona una API clara para que se conecten los backends del compilador. Modifica de forma dinámica el código de bytes de Python justo antes de la ejecución. En la versión 2.0 de PyTorch/XLA, hay un backend experimental para la inferencia y el entrenamiento con Dynamo.

Dynamo proporciona un gráfico Torch FX (FX) cuando reconoce un patrón de modelo, y PyTorch/XLA usa un enfoque de tensor diferido para compilar el gráfico FX y mostrar la función compilada. Para obtener más información sobre Dynamo, consulta los siguientes recursos:

• Publicación de debates para desarrolladores de Pytorch
• Documentación de TorchDynamo
• PyTorch 2.0 y XLA para obtener más detalles

Este es un pequeño ejemplo de código para ejecutar la inferencia de densenet161 con torch.compile.

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

Puedes usar la imagen de Docker torchserve-tpu proporcionada para entregar tu modelo de PyTorch archivado en una VM de TPU.

Configura la autenticación para Docker:

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

Extrae la imagen de Docker de TorchServe de Cloud TPU a tu VM de 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}

Recopila artefactos de modelo

Para comenzar, debes proporcionar un controlador de modelos, que le indica al trabajador del servidor de modelos de TorchServe que cargue tu modelo, procese los datos de entrada y ejecute la inferencia. Puedes usar los controladores de inferencia predeterminados de TorchServe (fuente) o desarrollar tu propio controlador de modelos personalizados siguiendo base_handler.py. Es posible que también debas proporcionar el modelo entrenado y el archivo de definición del modelo.

Nota: Todos los controladores de modelos predeterminados de TorchServe admiten PyTorch/XLA.

En el siguiente ejemplo de Densenet 161, usamos artefactos de modelos y el controlador de clasificador de imágenes predeterminado que proporciona TorchServe:

• model.py
• index_to_name.json
• image_classifier.py

1. Configura algunas variables de entorno:

   CWD="$(pwd)"
   
   WORKDIR="${CWD}/densenet_161"
   
   mkdir -p ${WORKDIR}/model-store
   mkdir -p ${WORKDIR}/logs
   

2. Descarga y copia los artefactos del modelo del ejemplo de clasificador de imágenes de 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}
   

3. Descarga los pesos del modelo:

   wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth
   
   mv densenet161-8d451a50.pth ${WORKDIR}
   

4. Crea un archivo de configuración del modelo de TorchServe para usar el backend de Dynamo:

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

   Deberías ver los siguientes archivos y directorios:

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

Genera un archivo de almacenamiento de modelos

Para entregar tu modelo de PyTorch con Cloud TPU TorchServe, debes empaquetar el controlador de tu modelo y todos sus artefactos en un archivo de almacenamiento de modelos (*.mar) con el archivador de modelos de PyTorch.

Genera un archivo de almacenamiento de modelos con 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

Deberías ver el archivo de almacenamiento del modelo generado en el directorio model-store:

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

Publica solicitudes de inferencia

Ahora que tienes el archivo de almacenamiento en caché del modelo, puedes iniciar el servidor de modelos de TorchServe y entregar solicitudes de inferencia.

1. Inicia el servidor de modelos 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
   

2. Consulta el estado del servidor de modelos:

   curl http://localhost:8080/ping
   

   Si el servidor de modelos está en funcionamiento, verás lo siguiente:

   {
     "status": "Healthy"
   }
   

   Para consultar las versiones predeterminadas del modelo registrado actual, usa lo siguiente:

   curl http://localhost:8081/models
   

   Deberías ver el modelo registrado:

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

   Para descargar una imagen para la inferencia, usa lo siguiente:

   curl -O https://raw.githubusercontent.com/pytorch/serve/master/docs/images/kitten_small.jpg
   
   mv kitten_small.jpg ${WORKDIR}
   

   Para enviar una solicitud de inferencia al servidor de modelos, usa lo siguiente:

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

   Deberías ver una respuesta similar a la siguiente:

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

3. Registros del servidor de modelos

   Usa los siguientes comandos para acceder a los registros:

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

   Deberías ver el siguiente mensaje en tu registro:

   "Compiled model with backend torchxla\_trace\_once"
   

Limpia

Detén el contenedor de Docker:

rm -rf serve
rm -rf ${WORKDIR}

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

Generación de perfiles

Después de configurar la inferencia, puedes usar generadores de perfiles para analizar el rendimiento y la utilización de TPU. Para obtener más información sobre la generación de perfiles, consulta los siguientes recursos:

• Creación de perfiles en Cloud TPU
• Perfilación de TensorFlow
• Creación de perfiles de PyTorch
• Creación de perfiles de JAX