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, optimisé pour Entraînement, affinage et inférence basés sur les modèles de texte à image et CNN basés sur les transformations Transformer (inférence). 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é de production, où il peut être utilisé pour l'inférence. Les SLO de latence pour la diffusion.

Ce document traite de la diffusion d'un modèle sur un TPU à hôte unique. les tranches de TPU avec Huit puces ou moins possèdent une VM ou un hôte TPU et sont appelées TPU à hôte unique.

Commencer

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

Pour utiliser Cloud TPU, vous devez disposer d'un compte et d'un projet Google Cloud. Pour plus plus d'informations, consultez la page Configurer un environnement Cloud TPU.

Remarque :Pour utiliser TPU v5e, vous devez installer les composants gcloud alpha.

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

Inférence et inférence de modèles Cloud TPU

La manière dont vous inférez un modèle dépend du framework de ML utilisé dans lequel vous avez écrit. Le TPU v5e permet d'inférer des 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 dans le SavedModel TensorFlow format
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 qui inclut paramètres et des calculs. Il n'a pas besoin du code de création du modèle d'origine s'exécuter.

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

Convertisseur d'inférence

Le convertisseur d'inférence Cloud TPU prépare et optimise un modèle exporté dans Format SavedModel pour TPU l'inférence. 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écessaire à l'exécution du convertisseur. Pour en savoir plus sur l'objet Inference Consultez le Guide de l'utilisateur de l'outil Inference Converter.

Exigences liées au convertisseur d'inférence

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

2. Vous devez définir un alias pour la fonction TPU. Pour plus d'informations, 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 eXtensions vectorielles avancées (AVX) les instructions, comme dans la bibliothèque TensorFlow (la dépendance du Cloud TPU Inference Converter) est compilé pour utiliser des instructions AVX. La plupart des processeurs prendre en charge AVX.

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

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 Inférence.

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. Configurer vos identifiants Docker et extraire le convertisseur d'inférence et le Cloud TPU Diffuser une image Docker:

   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:

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

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é de 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, vérifier le 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êtes à être traité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. 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 le modèle de diffusion stable pré-entraîné de Hugging Face.

1. Télécharger le modèle Stable Diffusion dans un modèle enregistré TF2 compatible TPU format:

   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, vérifiez dans le journal du conteneur du serveur de modèles les serveurs gRPC et HTTP s'exécutent:

   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 sont prêts à diffuser 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. Vous 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 requête. 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 vous expliquent comment diffuser 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. Configurer les identifiants Docker et extraire le convertisseur d'inférence 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:

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

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

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

5. Inférer un modèle TensorFlow à l'aide de TensorFlow Serving une image Docker 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 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èles Python pour le chargement et l'inférence à l'aide de TorchDynamo et 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 langage Python Compilateur JIT conçu pour accélérer les programmes PyTorch. Il fournit une API propre pour que les backends de compilation puissent s'associer. Il modifie dynamiquement le bytecode Python 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 avec Dynamo.

Dynamo fournit un graphique Torch FX (FX) lorsqu'il reconnaît un modèle de modèle et que PyTorch/XLA utilise une approche de Tensor différé pour compiler le graphe FX et renvoyer la fonction compilée. Pour en savoir plus, sur Dynamo, consultez les ressources 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 vos archives pytorch 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 indique au Le nœud de calcul du serveur de modèles TorchServe charge votre modèle, traite les données d'entrée et pour exécuter l'inférence. Vous pouvez utiliser les gestionnaires d'inférence par défaut TorchServe (source) ou Développez votre propre gestionnaire de modèles personnalisés à la suite du 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 l'API par défaut Gestionnaire de classificateur d'images 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 de 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