Présentation de l'inférence Cloud TPU v5e

Remarque :Si vous ne connaissez pas les Cloud TPU, nous vous recommandons de lire la présentation des Cloud TPU.

Présentation et avantages

Cloud TPU v5e est un accélérateur d'IA développé par Google, optimisé pour l'entraînement, le réglage et le traitement (inférence) basés sur des transformateurs, sur la conversion de texte en image et sur des CNN. Les tranches TPU v5e peuvent contenir jusqu'à 256 puces.

Le traitement désigne le processus de déploiement d'un modèle de machine learning entraîné dans un environnement de production, où il peut être utilisé pour l'inférence. Les SLO de latence sont une priorité pour la diffusion.

Ce document explique comment diffuser un modèle sur un TPU à hôte unique. Les tranches de TPU comportant huit puces ou moins ont une VM ou un hôte TPU et sont appelées TPU à hôte unique.

Premiers pas

Vous aurez besoin d'un quota pour les TPU v5e. Les TPU à la demande nécessitent un quota tpu-v5s-litepod-serving. Les TPU réservés nécessitent un quota tpu-v5s-litepod-serving-reserved. Pour en savoir plus, contactez le service commercial Cloud.

Vous avez besoin d'un Google Cloud compte et d'un projet pour utiliser Cloud TPU. Pour en savoir plus, consultez Configurer un environnement Cloud TPU.

Vous provisionnez des TPU v5e à l'aide de ressources en file d'attente. Pour en savoir plus sur les configurations v5e disponibles pour la diffusion, consultez la section Types de Cloud TPU v5e pour la diffusion.

Inférence et exécution de modèles Cloud TPU

La manière dont vous diffusez un modèle pour l'inférence dépend du framework de ML avec lequel votre modèle a été écrit. TPU v5e est compatible avec la diffusion de modèles écrits en JAX, TensorFlow et PyTorch.

Inférence et diffusion de modèles JAX

Pour diffuser un modèle sur une VM TPU, vous devez:

1. Sérialiser votre modèle au format SavedModel TensorFlow
2. Utiliser le convertisseur d'inférence pour préparer le modèle enregistré à la diffusion
3. Utiliser TensorFlow Serving pour diffuser le modèle

Format SavedModel

Un SavedModel contient un programme TensorFlow complet, y compris des paramètres et des calculs entraînés. Il ne nécessite pas le code de création du modèle d'origine pour s'exécuter.

Si votre modèle a été écrit en JAX, vous devrez utiliser jax2tf pour le sérialiser au format SavedModel.

Convertisseur d'inférence

Le convertisseur d'inférence Cloud TPU prépare et optimise un modèle exporté au format SavedModel pour l'inférence TPU. Vous pouvez exécuter le convertisseur d'inférence dans un shell local ou dans votre VM TPU. Nous vous recommandons d'utiliser le shell de votre VM TPU, car il contient tous les outils de ligne de commande nécessaires pour exécuter le convertisseur. Pour en savoir plus sur le convertisseur d'inférences, consultez le guide de l'utilisateur du convertisseur d'inférences.

Exigences concernant le convertisseur d'inférence

1. Votre modèle doit être exporté depuis TensorFlow ou JAX au format SavedModel.

2. Vous devez définir un alias de fonction pour la fonction TPU. Pour en savoir plus, consultez le guide de l'utilisateur du convertisseur d'inférences. Les exemples de ce guide utilisent tpu_func comme alias de fonction TPU.

3. Assurez-vous que le processeur de votre machine est compatible avec les instructions AVX (Advanced Vector Extensions), car la bibliothèque TensorFlow (dépendance du convertisseur d'inférence Cloud TPU) est compilée pour utiliser les instructions AVX. La plupart des processeurs sont compatibles avec AVX.

   Remarque :Vous pouvez exécuter lscpu | grep avx pour vérifier si l'ensemble d'instructions AVX est compatible.

Inférence et diffusion de modèles JAX

Cette section explique comment diffuser des modèles JAX à l'aide de jax2tf et de TensorFlow Serving.

1. Utiliser jax2tf pour sérialiser votre modèle au format SavedModel
2. Utiliser le convertisseur d'inférence pour préparer votre modèle enregistré à la diffusion
3. Utiliser TensorFlow Serving pour diffuser le modèle

Utiliser jax2tf pour sérialiser un modèle JAX au format SavedModel

La fonction Python suivante montre comment utiliser jax2tf dans le code de votre modèle:

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

Pour en savoir plus sur jax2tf, consultez la section Interopérabilité entre JAX et Cloud TPU.

Utiliser le convertisseur d'inférence pour préparer le modèle enregistré à la diffusion

Pour savoir comment utiliser le convertisseur d'inférences, consultez le guide du convertisseur d'inférences.

Utiliser TensorFlow Serving

Pour savoir comment utiliser TensorFlow Serving, consultez la section Diffusion TensorFlow.

Exemples de mise en service de modèles JAX

Prérequis

1. Configurez vos identifiants Docker et extrayez l'image Docker du convertisseur d'inférence et du service 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. Connectez-vous à votre VM TPU à l'aide de SSH et installez le code de démonstration d'inférence:

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

3. Installez les dépendances de la démonstration JAX:

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

Diffuser le modèle BERT JAX pour l'inférence

Vous pouvez télécharger le modèle BERT pré-entraîné depuis Hugging Face.

Remarque :Les exportations de modèles enregistrés PyTorch ne sont pas acceptées.

1. Exporter un modèle SavedModel TensorFlow compatible avec les TPU à partir d'un modèle BERT Flax:

   cd demo/jax/bert
   python3 export_bert_model.py
   

2. Démarrez le conteneur du serveur de modèles 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
   

   Environ 30 secondes après le démarrage du conteneur, vérifiez le journal du conteneur du serveur de modèle et assurez-vous que les serveurs gRPC et HTTP sont opérationnels:

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

   Si vous voyez une entrée de journal se terminant par les informations suivantes, le serveur est prêt à répondre aux requêtes.

   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. Envoyez une requête d'inférence au serveur de modèles.

   python3 bert_request.py
   

   Le résultat doit ressembler à ce qui suit :

   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. effectuer un nettoyage.

   Veillez à nettoyer le conteneur Docker avant d'exécuter d'autres démonstrations.

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

   Nettoyez les artefacts de modèle:

   sudo rm -rf /tmp/jax/
   

Diffuser Stable Diffusion JAX pour l'inférence

Vous pouvez télécharger le modèle Stable Diffusion pré-entraîné sur Hugging Face.

1. Téléchargez le modèle Stable Diffusion au format de modèle enregistré TF2 compatible avec TPU:

   Remarque :Le téléchargement prend environ trois minutes.

   cd demo/jax/stable_diffusion
   python3 export_stable_diffusion_model.py
   

2. Démarrez le conteneur du serveur de modèle Cloud TPU pour le modèle:

   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
   

   Après environ deux minutes, vérifiez le journal du conteneur du serveur de modèle pour vous assurer que les serveurs gRPC et HTTP sont en cours d'exécution:

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

   Si le journal se termine par les informations suivantes, cela signifie que les serveurs sont prêts à traiter les requêtes.

   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. Envoyez une requête au serveur de modèles.

   python3 stable_diffusion_request.py
   

   Ce script envoie "Peinture d'un écureuil en train de patiner à New York" comme invite. L'image de sortie est enregistrée sous le nom stable_diffusion_images.jpg dans votre répertoire actuel.

4. effectuer un nettoyage.

   Veillez à nettoyer le conteneur Docker avant d'exécuter d'autres démonstrations.

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

   Nettoyer les artefacts de modèle

   sudo rm -rf /tmp/jax/
   

TensorFlow Serving

Les instructions suivantes montrent comment diffuser votre modèle TensorFlow sur des VM TPU.

Workflow de TensorFlow Serving

1. Téléchargez l'image Docker TensorFlow Serving pour votre VM TPU.

   Définir des exemples de variables d'environnement

   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

   Télécharger l'image Docker

   docker pull ${IMAGE_NAME}
   

2. Configurez les identifiants Docker et extrayez l'image Docker du convertisseur d'inférence et 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. Téléchargez le code de démonstration:

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

4. Installez les dépendances de la démonstration TensorFlow:

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

5. Diffuser votre modèle TensorFlow à l'aide de l'image Docker TensorFlow Serving sur votre VM 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. Utilisez l'API Client de diffusion pour interroger votre modèle.

   • API cliente REST
   • API cliente gRPC

Exécuter la démonstration de TensorFlow ResNet-50 Serving

1. Exportez un modèle SavedModel TF2 compatible avec TPU à partir du modèle Keras ResNet-50.

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

2. Lancez le conteneur du serveur de modèles TensorFlow pour le modèle.

   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
   

   Vérifiez le journal du conteneur du serveur de modèles et assurez-vous que le serveur gRPC et HTTP est opérationnel:

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

   Si le journal se termine par les informations suivantes, cela signifie que le serveur est prêt à traiter les requêtes. Cela prend environ 30 secondes.

   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. Envoyez la requête au serveur de modèles.

   L'image demandée est une banane provenant de https://i.imgur.com/j9xCCzn.jpeg .

   python3 resnet_request.py
   

   Le résultat doit ressembler à ce qui suit :

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

4. effectuer un nettoyage.

   Veillez à nettoyer le conteneur Docker avant d'exécuter d'autres démonstrations.

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

   Nettoyez les artefacts de modèle:

   sudo rm -rf /tmp/tf/
   

Inférence et exécution de modèles PyTorch

Pour les modèles écrits avec PyTorch, le workflow est le suivant:

1. Écrire un gestionnaire de modèle Python pour le chargement et l'inférence à l'aide de TorchDynamo et de PyTorch/XLA
2. Créer une archive de modèle à l'aide de TorchModelArchiver
3. Utiliser TorchServe pour diffuser le modèle

TorchDynamo et PyTorch/XLA

TorchDynamo (Dynamo) est un compilateur JIT au niveau Python conçu pour accélérer les programmes PyTorch. Il fournit une API propre pour les backends de compilateur. Il modifie dynamiquement le bytecode Python juste avant l'exécution. Dans la version 2.0 de PyTorch/XLA, il existe un backend expérimental pour l'inférence et l'entraînement à l'aide de Dynamo.

Dynamo fournit un graphique Torch FX (FX) lorsqu'il reconnaît un modèle, et PyTorch/XLA utilise une approche de tenseur paresseux pour compiler le graphique FX et renvoyer la fonction compilée. Pour en savoir plus sur Dynamo, consultez les ressources suivantes:

• Post sur les discussions de développement Pytorch
• Documentation TorchDynamo
• Pour en savoir plus, consultez PyTorch 2.0 et XLA.

Voici un petit exemple de code pour exécuter l'inférence densenet161 avec 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

Vous pouvez utiliser l'image Docker torchserve-tpu fournie pour diffuser votre modèle PyTorch archivé sur une VM TPU.

Configurez l'authentification pour Docker:

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

Extrayez l'image Docker TorchServe Cloud TPU vers votre VM 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}

Collecter des artefacts de modèle

Pour commencer, vous devez fournir un gestionnaire de modèle, qui indique au worker du serveur de modèles TorchServe de charger votre modèle, de traiter les données d'entrée et d'exécuter l'inférence. Vous pouvez utiliser les gestionnaires d'inférence par défaut de TorchServe (source) ou développer votre propre gestionnaire de modèle personnalisé en suivant base_handler.py. Vous devrez peut-être également fournir le modèle entraîné et le fichier de définition du modèle.

Remarque :Tous les gestionnaires de modèles par défaut de TorchServe sont compatibles avec PyTorch/XLA.

Dans l'exemple Densenet 161 suivant, nous utilisons des artefacts de modèle et le gestionnaire de classifieur d'images par défaut fourni par TorchServe:

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

1. Configurez des variables d'environnement:

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

2. Téléchargez et copiez les artefacts de modèle de l'exemple de classifieur d'images 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. Téléchargez les pondérations du modèle:

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

4. Créez un fichier de configuration de modèle TorchServe pour utiliser le backend Dynamo:

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

   Les fichiers et répertoires suivants doivent s'afficher:

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

Générer un fichier d'archive de modèle

Pour diffuser votre modèle PyTorch avec Cloud TPU TorchServe, vous devez empaqueter votre gestionnaire de modèle et tous vos artefacts de modèle dans un fichier d'archive de modèle (*.mar) à l'aide de l'archiveur de modèles Torch.

Générez un fichier d'archive de modèle avec 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

Le fichier d'archive du modèle généré doit s'afficher dans le répertoire model-store:

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

Diffuser des requêtes d'inférence

Maintenant que vous disposez du fichier d'archive du modèle, vous pouvez démarrer le serveur de modèles TorchServe et traiter les requêtes d'inférence.

1. Démarrez le serveur de modèles 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. Interrogez l'état du serveur de modèles:

   curl http://localhost:8080/ping
   

   Si le serveur de modèles est opérationnel, les informations suivantes s'affichent:

   {
     "status": "Healthy"
   }
   

   Pour interroger les versions par défaut du modèle enregistré actuel, utilisez:

   curl http://localhost:8081/models
   

   Le modèle enregistré doit s'afficher:

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

   Pour télécharger une image à utiliser pour l'inférence:

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

   Pour envoyer une requête d'inférence au serveur de modèles, utilisez:

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

   Un résultat semblable à celui-ci doit s'afficher :

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

3. Journaux du serveur de modèle

   Pour accéder aux journaux, utilisez les commandes suivantes:

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

   Le message suivant doit s'afficher dans votre journal:

   "Compiled model with backend torchxla\_trace\_once"
   

Effectuer un nettoyage

Arrêtez le conteneur Docker:

rm -rf serve
rm -rf ${WORKDIR}

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

Profilage

Après avoir configuré l'inférence, vous pouvez utiliser des profileurs pour analyser les performances et l'utilisation des TPU. Pour en savoir plus sur le profilage, consultez les pages suivantes:

• Profilage sur Cloud TPU
• Profilage TensorFlow
• Profilage PyTorch
• Profilage JAX