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

Remarque :Si vous ne connaissez pas les Cloud TPU, nous vous recommandons de consulter la page Présentation de Cloud TPU.

Présentation et avantages

Cloud TPU v5e est un accélérateur d'IA développé par Google et optimisé pour l'entraînement, l'affinage et l'inférence (inférence) basées sur un transformateur, du texte à l'image et des réseaux de neurones convolutifs. Les tranches de TPU v5e peuvent contenir jusqu'à 256 puces.

L'inférence 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 constituent une priorité pour la diffusion.

Ce document traite de la diffusion d'un modèle sur un TPU à hôte unique. Les tranches de TPU comportant huit puces ou moins disposent d'une VM ou d'un hôte TPU et sont appelées TPU à hôte unique.

Démarrer

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.

Pour utiliser Cloud TPU, vous devez disposer d'un compte et d'un projet Google Cloud. Pour en savoir plus, consultez la page 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 page Types de Cloud TPU v5e pour la diffusion.

Inférence et inférence 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. Le TPU v5e est compatible avec l'inférence de modèles écrits en JAX, TensorFlow et PyTorch.

Inférence et inférence 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. Préparer le modèle enregistré en vue de l'inférence à l'aide du convertisseur d'inférence
3. Utiliser TensorFlow Serving pour diffuser le modèle

Format SavedModel

Un SavedModel contient un programme TensorFlow complet incluant les paramètres entraînés et les calculs. L'exécution du code de création du modèle d'origine n'est pas nécessaire.

Si votre modèle a été écrit en JAX, vous devez 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 une interface système locale ou votre VM TPU. Nous vous recommandons d'utiliser l'interface système de votre VM TPU, car elle dispose de tous les outils de ligne de commande nécessaires à l'exécution du convertisseur. Pour en savoir plus sur le convertisseur d'inférence, consultez le guide de l'utilisateur du convertisseur d'inférence.

Exigences liées au 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 pour la fonction TPU. Pour en savoir plus, consultez le guide de l'utilisateur du convertisseur d'inférence. 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 (la 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 inférence 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 le modèle enregistré en vue de son inférence
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ération entre JAX et Cloud TPU.

Préparer le modèle enregistré en vue de l'inférence à l'aide du convertisseur d'inférence

Les instructions d'utilisation du convertisseur d'inférence sont décrites dans le guide du convertisseur d'inférence.

Utiliser TensorFlow Serving

Les instructions d'utilisation de TensorFlow Serving sont décrites dans la section TensorFlow Serving.

Exemples d'inférence de modèle JAX

Prérequis

1. Configurez vos identifiants Docker et extrayez le convertisseur d'inférence et l'image Docker de diffusion 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 avec SSH et installez le code de démonstration d'inférence:

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

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

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

Inférer le modèle JAX BERT 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. Exportez un modèle enregistré TensorFlow compatible avec les TPU à partir d'un modèle Flax BERT:

   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, consultez le journal du conteneur du serveur de modèles, 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 à 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. Envoyer 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 du modèle:

   sudo rm -rf /tmp/jax/
   

Diffuser la diffusion stable JAX pour l'inférence

Vous pouvez télécharger un modèle de diffusion stable pré-entraîné depuis Hugging Face.

1. Téléchargez le modèle Stable Diffusion dans un format de modèle enregistré TF2 compatible 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èles 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, consultez le journal du conteneur du serveur de modèles 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 avec les informations suivantes, cela signifie que les serveurs sont prêts à diffuser des 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. Vous envoyez une requête au serveur de modèles.

   python3 stable_diffusion_request.py
   

   Ce script envoie l'invite "Peinture d'un écureuil en train de faire du patin à New York". L'image de sortie sera 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 du modèle

   sudo rm -rf /tmp/jax/
   

TensorFlow Serving

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

Workflow d'inférence TensorFlow

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 le convertisseur d'inférence et l'image Docker 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:

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

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

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

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

   • API cliente REST
   • API cliente gRPC

Lancer la démonstration de TensorFlow ResNet-50 Serving

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

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

2. Lancer 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 gRPC et le serveur HTTP sont opérationnels:

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

   Si le journal se termine avec 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. Vous envoyez la requête au serveur de modèles.

   L'image de demande 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 du modèle:

   sudo rm -rf /tmp/tf/
   

Inférence et inférence 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. Utiliser TorchModelArchiver pour créer une archive de modèle
3. Utiliser TorchServe pour diffuser le modèle

TorchDynamo et PyTorch/XLA

TorchDynamo (Dynamo) est un compilateur JIT de niveau Python conçu pour accélérer les programmes PyTorch. Il fournit une API propre à laquelle les backends de compilation peuvent s'associer. Il modifie de façon dynamique le bytecode Python juste avant son 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 graphe Torch FX (FX) lorsqu'il reconnaît un modèle de modèle, et PyTorch/XLA utilise une approche de Tensor différé pour compiler le graphe d'effets et renvoyer la fonction compilée. Pour en savoir plus sur Dynamo, consultez les pages suivantes:

• Post Discussions pour les développeurs Pytorch
• Documentation TorchDynamo
• PyTorch 2.0 et XLA pour en savoir plus

Voici un petit exemple de code d'exécution de 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 sur 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 demande au nœud de calcul 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 TorchServe (source) ou développer votre propre gestionnaire de modèles personnalisé en suivant le fichier 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 TorchServe sont compatibles avec PyTorch/XLA

Dans l'exemple Densenet 161 suivant, nous utilisons des artefacts de modèle et le gestionnaire de classificateur 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 des artefacts de modèle à partir de l'exemple de classificateur 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 Torch Model Archiver.

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 doit être généré 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èle TorchServe et diffuser des requêtes d'inférence.

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

   curl http://localhost:8080/ping
   

   Si le serveur de modèles est opérationnel, le message suivant s'affiche:

   {
     "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 à des fins d'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èles

   Utilisez les commandes suivantes pour accéder aux journaux:

   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 du TPU. Pour en savoir plus sur le profilage, consultez les pages suivantes:

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