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

Remarque :Si vous débutez avec Cloud TPU, nous vous recommandons de lire 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, le réglage et la diffusion basés sur des transformateurs, du texte à l'image et des réseaux de neurones convolutifs. Les tranches TPU v5e peuvent contenir jusqu'à 256 puces.

La diffusion fait référence au 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 traite de la diffusion d'un modèle sur un TPU à hôte unique. Les tranches TPU comportant huit puces ou moins disposent d'une VM ou d'un hôte TPU et sont appelées TPU à hôte unique.

Premiers pas

Vous avez besoin d'un quota pour les TPU v5e. Les TPU à la demande nécessitent un quota de tpu-v5s-litepod-serving. Les TPU réservés nécessitent un quota de 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 section 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 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. TPU v5e permet de diffuser 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 au format SavedModel TensorFlow
2. Utiliser le convertisseur d'inférence pour préparer le modèle enregistré en vue de l'inférence
3. Utiliser TensorFlow Serving pour diffuser le modèle

Format SavedModel

Un SavedModel contient un programme TensorFlow complet, y compris des paramètres entraînés et des opérations de calcul. Le code de création du modèle d'origine n'est pas nécessaire pour s'exécuter.

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 un shell local ou dans votre VM TPU. Nous vous recommandons d'utiliser l'interface système de votre VM TPU, car elle contient tous les outils de ligne de commande nécessaires à l'exécution du convertisseur. Pour plus d'informations sur le convertisseur d'inférence, consultez le Guide de l'utilisateur du convertisseur d'inférence.

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é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 l'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 page Interopération entre JAX et Cloud TPU.

Utiliser le convertisseur d'inférence pour préparer le modèle enregistré en vue de l'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, puis 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 à l'aide de 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
   

Déployer 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 TPU à partir d'un modèle Flax BERT:

   cd demo/jax/bert
   python3 export_bert_model.py
   

2. Démarrez le conteneur de 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 une entrée de journal se termine par les informations suivantes, le serveur est prêt à 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. 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/
   

Livrer la diffusion stable JAX pour l'inférence

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

1. Téléchargez le modèle Stable Diffusion dans un 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 de 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 de 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 par 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. Envoyer une requête au serveur de modèles.

   python3 stable_diffusion_request.py
   

   Ce script envoie l'invite "Peinture d'un écureuil qui patine à New York". L'image de sortie sera enregistrée sous 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 montrent comment diffuser votre modèle TensorFlow sur des VM TPU.

Workflow de diffusion 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, puis 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 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 du client de diffusion pour interroger votre modèle.

   • API du client REST
   • API client gRPC

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

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

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

2. Lancer le conteneur de 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
   

   Consultez le journal du conteneur de 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 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 tirée 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 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 permettant aux backends de compilation de se connecter. Il modifie de manière dynamique le bytecode Python juste avant l'exécution. La version PyTorch/XLA 2.0 propose 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. PyTorch/XLA utilise une approche de Tensor différée pour compiler le graphe de FX et renvoyer la fonction compilée. Pour en savoir plus sur Dynamo, consultez les pages suivantes:

• Post PyTorch Dev Discussions
• Documentation TorchDynamo
• PyTorch 2.0 et XLA pour en savoir plus

Voici un petit exemple de code permettant d'exécuter une 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 Cloud TPU TorchServe 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èles, 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és 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 certaines 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 du 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 à l'aide de 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 diffuser des 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. État du serveur du modèle de requête:

   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, exécutez la commande suivante:

   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 du modèle

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

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