Introdução à inferência da Cloud TPU v5e
Informações gerais e benefícios
A Cloud TPU v5e é um acelerador de IA desenvolvido pelo Google, otimizado para treinamento, ajuste fino e veiculação (inferência) de texto para imagem e baseado em transformador e CNN. As fatias da TPU v5e podem conter até 256 chips.
A veiculação se refere ao processo de implantação de um modelo de aprendizado de máquina treinado em um ambiente de produção, onde ele pode ser usado para inferência. Os SLOs de latência são uma prioridade para veiculação.
Este documento discute a disponibilização de um modelo em uma TPU de host único. As frações de TPU com 8 chips ou menos têm uma VM ou um host de TPU e são chamadas de TPUs de host único.
Primeiros passos
Você vai precisar de cota para TPUs v5e. Os TPUs sob demanda exigem a cota tpu-v5s-litepod-serving
. TPUs reservados exigem cota de tpu-v5s-litepod-serving-reserved
. Para mais informações, entre em contato com a equipe de vendas do Cloud.
Você vai precisar de uma conta e um projeto do Google Cloud para usar o Cloud TPU. Para mais informações, consulte Configurar um ambiente do Cloud TPU.
Você provisiona TPUs v5e usando recursos na fila. Para mais informações sobre as configurações disponíveis da v5e para veiculação, consulte Tipos de Cloud TPU v5e para veiculação.
Inferência e exibição de modelos do Cloud TPU
A maneira de veicular um modelo para inferência depende do framework de ML usado para criar o modelo. A TPU v5e oferece suporte à veiculação de modelos escritos em JAX, TensorFlow e PyTorch.
Inferência e veiculação de modelos JAX
Para disponibilizar um modelo em uma VM da TPU, você precisa:
- Serialize seu modelo no formato SavedModel do TensorFlow.
- Use o Converter de inferência para preparar o modelo salvo para veiculação
- Usar o TensorFlow Serving para mostrar o modelo
Formato SavedModel
Um SavedModel contém um programa completo do TensorFlow, incluindo parâmetros treinados e computação. Ele não exige que o código de criação do modelo original seja executado.
Se o modelo foi escrito no JAX, use jax2tf
para serializar
o modelo no formato SavedModel.
Conversor de inferência
O conversor de inferência da 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 de TPU porque ele tem todas as ferramentas de linha de comando necessárias para executar o conversor. Para mais informações sobre o Converter de inferência, consulte o Guia do usuário do Converter de inferência.
Requisitos do Inference Converter
O modelo precisa ser exportado do TensorFlow ou do JAX no formato SavedModel.
É necessário definir um alias de função para a função TPU. Para mais informações, consulte o Guia do usuário do Inference Converter. Os exemplos neste guia usam
tpu_func
como o alias da função TPU.Verifique se a CPU da máquina oferece suporte a instruções de Extensões de vetor avançadas (AVX, na sigla em inglês). A biblioteca do TensorFlow (a dependência do Conversor de inferência do Cloud TPU) é compilada para usar instruções AVX. A maioria das CPUs tem suporte a AVX.
Inferência e veiculação de modelos JAX
Esta seção descreve como disponibilizar modelos do JAX usando jax2tf
e o TensorFlow
Serving.
- Use
jax2tf
para serializar seu modelo no formato SavedModel. - Use o Converter de inferência para preparar o modelo salvo para exibição
- Usar o TensorFlow Serving para mostrar o modelo
Use jax2tf
para serializar um modelo JAX no formato SavedModel
A função 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 do Cloud TPU.
Use o Converter de inferência para preparar o modelo salvo para veiculação
As instruções para usar o Inference Converter estão descritas no guia do Inference Converter.
Usar o TensorFlow Serving
As instruções para usar o TensorFlow Serving estão descritas em Serviço do TensorFlow.
Exemplos de disponibilização de modelos JAX
Pré-requisitos
Configure suas credenciais do Docker e extraia a imagem do Docker do Converter de inferência e do serviço do 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
Conecte-se à VM do TPU com SSH e instale o código de demonstração de inferência:
gcloud storage cp \ "gs://cloud-tpu-inference-public/demo" \ . \ --recursive
Instale as dependências da demonstração do JAX:
pip install -r ./demo/jax/requirements.txt
Exibir o modelo BERT JAX para inferência
Você pode fazer o download do modelo BERT pré-treinado no Hugging Face.
Exporte um modelo salvo do TensorFlow compatível com TPU de um modelo BERT do Flax:
cd demo/jax/bert python3 export_bert_model.py
Inicie o contêiner do servidor de modelos do 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 o início do contêiner, verifique o registro do contêiner do servidor de modelo e verifique 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ê encontrar uma entrada de registro que termina com as informações a seguir, o servidor estará 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 ...
Envie uma solicitação de inferência para o 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?."
Fazer a limpeza.
Limpe 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/
Servir a difusão estável do JAX para inferência
Faça o download do modelo pré-treinado da Stable Diffusion no Hugging Face.
Faça o download do modelo de difusão estável em um formato de modelo salvo do TF2 compatível com TPU:
cd demo/jax/stable_diffusion python3 export_stable_diffusion_model.py
Inicie o contêiner do servidor de 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 de 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, 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 ...
Envie uma solicitação para o servidor do modelo.
python3 stable_diffusion_request.py
Esse script envia "Pintura de um esquilo patinando em Nova York" como instrução. A imagem de saída será salva como
stable_diffusion_images.jpg
no seu diretório atual.Fazer a limpeza.
Limpe 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
As instruções a seguir demonstram como você pode oferecer seu modelo do TensorFlow em VMs TPU.
Fluxo de trabalho do TensorFlow Serving
Faça o download da imagem do Docker do TensorFlow Serving para sua VM da TPU.
Definir variáveis de ambiente de exemplo
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 do Docker
docker pull ${IMAGE_NAME}
Configure as credenciais do Docker e extraia o Converter de inferência e a imagem do Docker do 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
Faça o download do código de demonstração:
gcloud storage cp \ "gs://cloud-tpu-inference-public/demo" \ . \ --recursive
Instale as dependências da demonstração do TensorFlow:
pip install -r ./demo/tf/requirements.txt
Exiba seu modelo do TensorFlow usando a imagem do Docker do TensorFlow Serving na VM da 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}
Use a API de cliente de veiculação para consultar seu modelo.
Executar a demonstração do TensorFlow ResNet-50 Serving
Exporte um modelo salvo do TF2 compatível com TPU do modelo Keras ResNet-50.
cd demo/tf/resnet-50 python3 export_resnet_model.py
Inicie 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 verifique 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 a seguir, significa que o servidor está pronto para atender solicitações. Isso 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 ...
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)]]
Fazer a limpeza.
Limpe 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 disponibilização de modelos do PyTorch
Para modelos criados com PyTorch, o fluxo de trabalho é o seguinte:
- Programar um gerenciador de modelo Python para carregar e fazer inferência usando
TorchDynamo
e PyTorch/XLA - Use
TorchModelArchiver
para criar um arquivo de modelo - Use
TorchServe
para disponibilizar o modelo
TorchDynamo e PyTorch/XLA
O TorchDynamo (Dynamo) é um compilador JIT no nível do Python projetado para tornar os programas do PyTorch mais rápidos. Ele fornece uma API limpa para que os back-ends do compilador se conectem. 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 Torch FX (FX) quando reconhece um padrão de modelo, e o PyTorch/XLA usa uma abordagem de tensor preguiçoso para compilar o gráfico FX e retornar a função compilada. Para mais informações sobre o Dynamo, consulte:
- Postagem do Pytorch Dev Discussions (em inglês)
- Documentação do TorchDynamo
- PyTorch 2.0 e XLA para mais detalhes.
Confira um exemplo de código para executar a inferência de 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 do Docker torchserve-tpu
fornecida para exibir o modelo
pytorch arquivado em uma VM TPU.
Configure a autenticação para o Docker:
sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
us-docker.pkg.dev
Extraia a imagem do Docker do TorchServe do Cloud TPU 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, é necessário fornecer um gerenciador de modelo, que instrui o trabalhador do servidor de modelo 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 TorchServe (fonte) 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.
No exemplo de Densenet 161 abaixo, usamos artefatos de modelo e o gerenciador de classificador de imagem padrão fornecido pelo TorchServe:
Configure algumas variáveis de ambiente:
CWD="$(pwd)" WORKDIR="${CWD}/densenet_161" mkdir -p ${WORKDIR}/model-store mkdir -p ${WORKDIR}/logs
Faça o download e copie os artefatos do modelo do exemplo de classificador de imagens do 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}
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}
Crie um arquivo de configuração do modelo TorchServe para usar o back-end do Dynamo:
echo 'pt2: "torchxla_trace_once"' >> ${WORKDIR}/model_config.yaml
Você vai encontrar 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 exibir seu modelo PyTorch com o TorchServe do Cloud TPU, é necessário empacotar
o gerenciador de modelo e todos os artefatos em um arquivo de modelo (*.mar)
usando o arquivador de modelo do Torch.
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
O arquivo de modelo gerado vai aparecer no diretório model-store:
>> ls ${WORKDIR}/model-store
Densenet161.mar
Atender solicitações de inferência
Agora que você tem o arquivo de modelo, pode iniciar o servidor de modelo do TorchServe e atender às solicitações de inferência.
Inicie o servidor de modelo 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
Consultar a integridade do servidor de modelo:
curl http://localhost:8080/ping
Se o servidor de modelo estiver em execução, 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 uso de inferência:
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 de 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 }
Registros do servidor do modelo
Use os comandos a seguir para acessar os registros:
ls ${WORKDIR}/logs/ cat ${WORKDIR}/logs/model_log.log
A seguinte mensagem vai aparecer no seu 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
Criação de perfil
Depois de configurar a inferência, use os profilers para analisar o desempenho e a utilização da TPU. Para mais informações sobre a criação de perfis, consulte:
- Criação de perfil no Cloud TPU
- Criação de perfil do TensorFlow
- Criação de perfil do PyTorch
- Criação de perfil do JAX