Introdução à inferência do Cloud TPU v5e

Observação: se você ainda não conhece os Cloud TPUs, recomendamos que leia Introdução ao Cloud TPU.

Informações gerais e benefícios

O Cloud TPU v5e é um acelerador de IA desenvolvido pelo Google e otimizado para treinamento, ajuste e disponibilização (inferência) baseado em transformadores, texto em imagem e CNN. As fatias de TPU v5e podem conter até 256 chips.

Exibição refere-se ao processo de implantação de um modelo de machine learning treinado em um ambiente de produção, onde ele pode ser usado para inferência. SLOs de latência são uma prioridade para disponibilização.

Neste documento, falamos sobre a disponibilização de um modelo em uma TPU de host único. Frações de TPU com oito ou menos chips têm uma VM ou host de TPU e são chamadas de TPUs de host único.

Começar

Você precisará de cota para TPUs v5e. TPUs sob demanda exigem uma cota de tpu-v5s-litepod-serving. As TPUs reservadas exigem uma cota de tpu-v5s-litepod-serving-reserved. Para mais informações, entre em contato com a equipe de vendas do Cloud.

É preciso ter uma conta e um projeto do Google Cloud para usar o Cloud TPU. Para mais informações, consulte Configurar um ambiente do Cloud TPU.

Observação: para usar a TPU v5e, é preciso instalar os componentes Alfa da gcloud.

Você provisiona TPUs v5e usando recursos na fila. Para mais informações sobre as configurações de v5e disponíveis para exibição, consulte Tipos do Cloud TPU v5e para exibição.

Inferência e veiculação de modelos do Cloud TPU

A forma como você disponibiliza um modelo para inferência depende do framework de ML com que o modelo foi escrito. A TPU v5e é compatível com a veiculação de modelos escritos em JAX, TensorFlow e PyTorch.

Inferência e exibição de modelo JAX

Para disponibilizar um modelo em uma VM de TPU, é preciso:

1. serializar o modelo no formato SavedModel do TensorFlow
2. Usar o conversor de inferência para preparar o modelo salvo para exibição
3. Usar o TensorFlow Serving para disponibilizar o modelo

Formato SavedModel

Um SavedModel contém um programa completo do TensorFlow, incluindo computação e parâmetros treinados. Ele não requer o código de criação do modelo original para ser executado.

Se o modelo foi escrito em JAX, você precisará usar jax2tf para serializar o modelo no formato SavedModel.

Conversor de inferência

O conversor de inferência do Cloud TPU prepara e otimiza um modelo exportado no formato SavedModel para inferência de TPU. É possível executar o conversor de inferência em um shell local ou na VM da TPU. Recomendamos o uso do shell da VM da TPU porque ele tem todas as ferramentas de linha de comando necessárias para executar o conversor. Para saber mais, consulte o Guia do usuário do conversor de inferência (link em inglês).

Requisitos do conversor de inferência

1. O modelo precisa ser exportado do TensorFlow ou do JAX no formato SavedModel.

2. É preciso definir um alias de função para a função da TPU. Para mais informações, consulte o Guia do usuário do conversor de inferência. Os exemplos neste guia usam tpu_func como o alias da função da TPU.

3. Verifique se a CPU da máquina é compatível com as instruções Advanced Vector eXtensions (AVX), já que a biblioteca do TensorFlow (a dependência do Cloud TPU Inference Converter) é compilada para usar as instruções do AVX. A maioria das CPUs é compatível com AVX.

   Observação: você pode executar lscpu | grep avx para verificar se o conjunto de instruções do AVX tem suporte.

Inferência e exibição de modelo JAX

Esta seção descreve como disponibilizar modelos JAX usando jax2tf e o TensorFlow Serving.

1. Usar jax2tf para serializar o modelo no formato SavedModel
2. Usar o conversor de inferência para preparar o modelo salvo para exibição
3. Usar o TensorFlow Serving para disponibilizar o modelo

Usar jax2tf para serializar um modelo JAX para o formato SavedModel

A função do Python a seguir mostra como usar jax2tf no código do 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 mais informações sobre jax2tf, consulte Interoperação do JAX e Cloud TPU.

Usar o conversor de inferência para preparar o modelo salvo para exibição

As instruções para usar o conversor de inferência são descritas no guia dele.

Usar o TensorFlow Serving

As instruções para usar o TensorFlow Serving são descritas em TensorFlow Serving.

Exemplos de exibição do modelo JAX

Pré-requisitos

1. Configure suas credenciais do Docker e extraia a imagem do Docker de veiculação do Cloud TPU e do Inference Converter:

   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. Conecte-se à VM da TPU com SSH e instale o código de demonstração de inferência:

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

3. Instale as dependências de demonstração do JAX:

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

Disponibilizar o modelo JAX BERT para inferência

Você pode fazer o download do modelo BERT pré-treinado da Hugging Face.

Observação :as exportações de modelos salvos do PyTorch não são aceitas.

1. Exporte um modelo salvo do TensorFlow compatível com TPU de um modelo Flax BERT:

   cd demo/jax/bert
   python3 export_bert_model.py
   

2. Inicie o contêiner do servidor do modelo da 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
   

   Cerca de 30 segundos após a inicialização do contêiner, verifique o registro do contêiner do servidor do modelo e se os servidores gRPC e HTTP estão ativos:

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

   Se você vir uma entrada de registro terminando com as seguintes informações, o servidor está pronto para atender às solicitações.

   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. Envia uma solicitação de inferência ao servidor do modelo.

   python3 bert_request.py
   

   A saída será semelhante a esta:

   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. Fazer a limpeza.

   Não se esqueça de limpar o contêiner do Docker antes de executar outras demonstrações.

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

   Limpe os artefatos do modelo:

   sudo rm -rf /tmp/jax/
   

Disponibilizar o JAX Stable Diffusion para inferência

Você pode fazer o download do modelo de difusão estável pré-treinado da Hugging Face.

1. Faça o download do modelo de difusão estável em um formato de modelo salvo do TF2 compatível com TPU:

   Observação :o download leva cerca de 3 minutos.

   cd demo/jax/stable_diffusion
   python3 export_stable_diffusion_model.py
   

2. Inicie o contêiner do servidor do modelo do Cloud TPU para o 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
   

   Após cerca de dois minutos, verifique o registro do contêiner do servidor do modelo para garantir que os servidores gRPC e HTTP estejam em execução:

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

   Se o registro terminar com as informações a seguir, isso significa que os servidores estão prontos para atender às solicitações.

   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. Envie uma solicitação ao servidor do modelo.

   python3 stable_diffusion_request.py
   

   O script envia "Painting of a squirrel skating in New York" como comando. A imagem de saída será salva como stable_diffusion_images.jpg no diretório atual.

4. Fazer a limpeza.

   Não se esqueça de limpar o contêiner do Docker antes de executar outras demonstrações.

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

   Limpar os artefatos do modelo

   sudo rm -rf /tmp/jax/
   

TensorFlow Serving

Com as instruções a seguir, mostramos como disponibilizar seu modelo do TensorFlow em VMs de TPU.

Fluxo de trabalho do TensorFlow Serving

1. Faça o download da imagem do Docker do TensorFlow Serving para sua VM da TPU.

   Definir variáveis de ambiente de amostra

   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
   

   Fazer o download da imagem Docker

   docker pull ${IMAGE_NAME}
   

2. Configure as credenciais do Docker e extraia a imagem do Docker do TensorFlow Serving e do conversor de inferência.

   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. Faça o download do código de demonstração:

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

4. Instale as dependências de demonstração do TensorFlow:

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

5. Exiba seu modelo do TensorFlow usando a imagem do Docker do TensorFlow Serving na sua 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. Use a API Serving Client para consultar o modelo.

   • API REST Client
   • API gRPC Client

Executar a demonstração do TensorFlow ResNet-50 Serving

1. Exportar do modelo Keras ResNet-50 um modelo salvo do TF2 compatível com TPU.

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

2. Iniciar o contêiner do servidor de modelos do TensorFlow para o 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
   

   Verifique o registro do contêiner do servidor de modelo e confirme se o gRPC e o servidor HTTP estão ativos:

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

   Se o registro terminar com as informações abaixo, isso significa que o servidor está pronto para atender às solicitações. O processo leva cerca 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. Envie a solicitação ao servidor do modelo.

   A imagem da solicitação é uma banana de https://i.imgur.com/j9xCCzn.jpeg .

   python3 resnet_request.py
   

   A saída será semelhante a esta:

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

4. Fazer a limpeza.

   Não se esqueça de limpar o contêiner do Docker antes de executar outras demonstrações.

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

   Limpe os artefatos do modelo:

   sudo rm -rf /tmp/tf/
   

Inferência e exibição de modelos do PyTorch

Para modelos escritos com PyTorch, o fluxo de trabalho é:

1. Criar um gerenciador de modelo do Python para carregar e inferir usando TorchDynamo e PyTorch/XLA
2. Use TorchModelArchiver para criar um arquivo de modelo.
3. Usar TorchServe para exibir o modelo

TorchDynamo e PyTorch/XLA

O TorchDynamo (Dynamo) é um compilador JIT no nível do Python projetado para tornar os programas PyTorch mais rápidos. Ele fornece uma API limpa para os back-ends do compilador se conectarem. Ele modifica dinamicamente o bytecode do Python antes da execução. Na versão PyTorch/XLA 2.0, há um back-end experimental para inferência e treinamento usando o Dynamo.

O Dynamo fornece um gráfico de Torch FX (FX) quando reconhece um padrão de modelo e o PyTorch/XLA usa uma abordagem de tensor lento para compilar o gráfico de FX e retornar a função compilada. Para mais informações sobre o Dynamo, consulte:

• Postagem do Pytorch Dev Discussions
• Documentação do TochDynamo
• PyTorch 2.0 e XLA para mais detalhes.

Este é um pequeno exemplo de código executando a inferência do densenet161 com 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

É possível usar a imagem torchserve-tpu do Docker fornecida para exibir seu modelo Pytorch arquivado em uma VM de TPU.

Configure a autenticação do Docker:

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

Extraia a imagem do Docker do Cloud TPU TorchServe para a VM da 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}

Coletar artefatos de modelo

Para começar, forneça um gerenciador de modelo, que instrui o worker do servidor de modelos do TorchServe a carregar seu modelo, processar os dados de entrada e executar a inferência. É possível usar os gerenciadores de inferência padrão do TochServe (source) ou desenvolver seu próprio gerenciador de modelo personalizado seguindo o base_handler.py. Talvez também seja necessário fornecer o modelo treinado e o arquivo de definição do modelo.

Observação :todos os gerenciadores de modelo padrão do TorchServe são compatíveis com PyTorch/XLA

No exemplo do Densenet 161 a seguir, usamos artefatos de modelo e o gerenciador de classificador de imagem padrão fornecido pelo TorchServe:

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

1. Configure algumas variáveis de ambiente:

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

2. Faça o download e copie os artefatos de modelo do exemplo do classificador de imagem 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. Faça o download dos pesos do modelo:

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

4. Crie um arquivo de configuração de modelo do TorchServe para usar o back-end do Dynamo:

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

   Você verá os seguintes arquivos e diretórios:

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

Gerar um arquivo de modelo

Para disponibilizar seu modelo PyTorch com o Cloud TPU TorchServe, é necessário empacotar o gerenciador de modelos e todos os artefatos dele em um arquivo de modelo (*.mar) usando o Torch Model Archiver.

Gere um arquivo de modelo com 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

Você verá o arquivo do modelo gerado no diretório model-store:

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

Exibir solicitações de inferência

Agora que você tem o arquivo do modelo, é possível iniciar o servidor de modelo do TorchServe e exibir solicitações de inferência.

1. Inicie o servidor de modelo do 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. Integridade do servidor do modelo de consulta:

   curl http://localhost:8080/ping
   

   Se o servidor do modelo estiver funcionando, você verá:

   {
     "status": "Healthy"
   }
   

   Para consultar as versões padrão do modelo registrado atual, use:

   curl http://localhost:8081/models
   

   O modelo registrado vai aparecer:

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

   Para fazer o download de uma imagem para inferência, use:

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

   Para enviar uma solicitação de inferência ao servidor do modelo, use:

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

   Uma resposta semelhante a esta vai aparecer:

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

3. Registros do servidor do modelo

   Use os seguintes comandos para acessar os registros:

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

   Você verá a seguinte mensagem no registro:

   "Compiled model with backend torchxla\_trace\_once"
   

Limpar

Pare o contêiner do Docker:

rm -rf serve
rm -rf ${WORKDIR}

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

Caracterização de perfil

Depois de configurar a inferência, use os criadores de perfil para analisar o desempenho e o uso da TPU. Para mais informações sobre a criação de perfis, consulte:

• Como criar o perfil no Cloud TPU
• Criação de perfis do TensorFlow
• Criação de perfil do PyTorch
• Criação de perfil do JAX