Introduzione all'inferenza di Cloud TPU v5e

Nota: se non hai mai utilizzato le Cloud TPU, ti consigliamo di leggere Introduzione a Cloud TPU.

Panoramica e vantaggi

Cloud TPU v5e è un acceleratore di AI sviluppato da Google e ottimizzato per l'addestramento, l'ottimizzazione e la pubblicazione basati su trasformatori, da testo a immagine e CNN (inferenza). Le sezioni TPU v5e possono contenere fino a 256 chip.

La pubblicazione si riferisce al processo di deployment di un modello di machine learning addestrato in un ambiente di produzione, in cui può essere utilizzato per l'inferenza. Gli SLO di latenza sono una priorità di gestione.

Questo documento illustra la gestione di un modello su una TPU singolo host. Le sezioni TPU con 8 chip o meno hanno una VM o un host TPU e sono chiamate TPU a host singolo.

Inizia

Avrai bisogno di una quota per le TPU v5e. Le TPU on demand richiedono una quota di tpu-v5s-litepod-serving. Le TPU riservate richiedono una quota pari a tpu-v5s-litepod-serving-reserved. Per ulteriori informazioni, contatta il team di vendita Cloud.

Per utilizzare Cloud TPU devi avere un account e un progetto Google Cloud. Per ulteriori informazioni, consulta Configurare un ambiente Cloud TPU

Nota: per utilizzare TPU v5e, devi installare i componenti gcloud alpha.

Per eseguire il provisioning delle TPU v5e, utilizza le risorse in coda. Per ulteriori informazioni sulle configurazioni v5e disponibili per la pubblicazione, consulta Tipi di Cloud TPU v5e per la pubblicazione.

Inferenza e gestione del modello Cloud TPU

Il modo in cui gestisci un modello per l'inferenza dipende dal framework ML con cui è stato scritto il modello. TPU v5e supporta modelli di gestione scritti in JAX, TensorFlow e PyTorch.

Inferenza e gestione del modello JAX

Per gestire un modello su una VM TPU, devi:

1. Serializza il modello in formato SavedModel di TensorFlow
2. Utilizza il convertitore di inferenza per preparare il modello salvato per la pubblicazione
3. Utilizza TensorFlow Serving per gestire il modello

Formato SaveModel

Un SaveModel contiene un programma TensorFlow completo, inclusi i parametri addestrati e il calcolo. Non richiede il codice di creazione del modello originale per essere eseguito.

Se il modello è stato scritto in JAX, dovrai utilizzare jax2tf per serializzare il modello nel formato SaveModel.

Convertitore di inferenza

Il convertitore di inferenza Cloud TPU prepara e ottimizza un modello esportato in formato SavedModel per l'inferenza TPU. Puoi eseguire il convertitore di inferenza in una shell locale o nella VM TPU. Ti consigliamo di utilizzare la shell della VM TPU perché dispone di tutti gli strumenti a riga di comando necessari per l'esecuzione del convertitore. Per ulteriori informazioni sullo strumento di conversione di inferenza, consulta la Guida dell'utente del convertitore di inferenza.

Requisiti di Inference Converter

1. Il modello deve essere esportato da TensorFlow o JAX nel formato SavedModel.

2. Devi definire un alias di funzione per la funzione TPU. Per ulteriori informazioni, consulta la Guida dell'utente per la conversione di inferenza. Gli esempi in questa guida utilizzano tpu_func come alias della funzione TPU.

3. Assicurati che la CPU della macchina supporti le istruzioni Advanced Vector eXtensions (AVX), poiché la libreria TensorFlow (la dipendenza di Cloud TPU Inference Converter) è compilata in modo da utilizzare le istruzioni AVX. La maggior parte delle CPU supporta AVX.

   Nota: puoi eseguire lscpu | grep avx per verificare se il set di istruzioni AVX è supportato.

Inferenza e gestione del modello JAX

Questa sezione descrive come gestire modelli JAX utilizzando jax2tf e TensorFlow Serve.

1. Usa jax2tf per serializzare il modello nel formato SaveModel
2. Utilizza il convertitore di inferenza per preparare il modello salvato per la pubblicazione
3. Utilizza TensorFlow Serving per gestire il modello

Usa jax2tf per serializzare un modello JAX nel formato SalvatoModel

La seguente funzione Python mostra come utilizzare jax2tf nel codice del modello:

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

Per maggiori informazioni su jax2tf, consulta Interoperabilità JAX e Cloud TPU.

Utilizza il convertitore di inferenza per preparare il modello salvato per la pubblicazione

Le istruzioni per l'utilizzo del Convertitore di inferenza sono descritte nella Guida al convertitore di inferenza.

Utilizzo di TensorFlow Serving

Le istruzioni per l'utilizzo di TensorFlow Serving sono descritte in Gestione di TensorFlow.

Esempi di pubblicazione di modelli JAX

Prerequisiti

1. Configura le credenziali Docker ed esegui il pull dell'immagine Docker di Inference Converter e 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. Connettiti alla VM TPU con SSH e installa il codice demo di inferenza:

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

3. Installa le dipendenze demo JAX:

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

Serve il modello BERT JAX per l'inferenza

Puoi scaricare il modello BERT preaddestrato da Hugging Face.

Nota: le esportazioni dei modelli salvati di PyTorch non sono supportate.

1. Esporta un modello salvato di TensorFlow compatibile con TPU da un modello BERT Flax:

   cd demo/jax/bert
   python3 export_bert_model.py
   

2. Avvia il container del server del modello 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
   

   Circa 30 secondi dopo l'avvio del container, controlla il log del container del server del modello e assicurati che i server gRPC e HTTP siano attivi:

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

   Se vedi una voce di log che termina con le informazioni seguenti, il server è pronto a gestire le richieste.

   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. Inviare una richiesta di inferenza al server del modello.

   python3 bert_request.py
   

   L'output sarà simile al seguente:

   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. Eseguire la pulizia.

   Assicurati di eseguire la pulizia del container Docker prima di eseguire altre demo.

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

   Esegui la pulizia degli artefatti del modello:

   sudo rm -rf /tmp/jax/
   

Servi la diffusione stabile JAX per l'inferenza

Puoi scaricare il modello Stable Diffusion preaddestrato da Hugging Face.

1. Scarica il modello Stable Diffusion in un formato di modello salvato TF2 compatibile con TPU:

   Nota: il download richiede circa 3 minuti.

   cd demo/jax/stable_diffusion
   python3 export_stable_diffusion_model.py
   

2. Avvia il container del server del modello Cloud TPU per il modello:

   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
   

   Dopo circa due minuti, controlla il log del container del server del modello per assicurarti che i server gRPC e HTTP siano in esecuzione:

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

   Se il log termina con le seguenti informazioni, significa che i server sono pronti a gestire le richieste.

   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. Invia una richiesta al server del modello.

   python3 stable_diffusion_request.py
   

   Questo script invia "Pittura di uno scoiattolo che pattina a New York" come prompt. L'immagine di output verrà salvata come stable_diffusion_images.jpg nella directory attuale.

4. Eseguire la pulizia.

   Assicurati di eseguire la pulizia del container Docker prima di eseguire altre demo.

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

   Esegui la pulizia degli artefatti del modello

   sudo rm -rf /tmp/jax/
   

Pubblicazione su TensorFlow

Le seguenti istruzioni mostrano come gestire il tuo modello TensorFlow su VM TPU.

Flusso di lavoro di gestione di TensorFlow

1. Scarica l'immagine Docker TensorFlow che gestisce la VM TPU.

   Imposta variabili di ambiente di esempio

   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

   Scarica l'immagine Docker

   docker pull ${IMAGE_NAME}
   

2. Configura le credenziali Docker ed esegui il pull dell'immagine Docker di Inference Converter e TensorFlow che gestisce il 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
   

3. Scarica il codice demo:

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

4. Installa le dipendenze della demo di TensorFlow:

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

5. Gestisci il tuo modello TensorFlow utilizzando l'immagine Docker di TensorFlow sulla 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. Utilizza l'API Serving Client per eseguire query sul modello.

   • API client REST
   • API client gRPC

Esegui la demo di gestione di TensorFlow ResNet-50

1. Esporta un modello salvato TF2 compatibile con TPU dal modello Keras ResNet-50.

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

2. Avvia il container del server del modello TensorFlow per il modello.

   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
   

   Controlla il log del container del server del modello e assicurati che gRPC e il server HTTP siano attivi:

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

   Se il log termina con le informazioni riportate di seguito, significa che il server è pronto a gestire le richieste. Impiega circa 30 secondi.

   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. Invia la richiesta al server del modello.

   L'immagine richiesta è una banana da https://i.imgur.com/j9xCCzn.jpeg .

   python3 resnet_request.py
   

   L'output sarà simile al seguente:

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

4. Eseguire la pulizia.

   Assicurati di eseguire la pulizia del container Docker prima di eseguire altre demo.

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

   Esegui la pulizia degli artefatti del modello:

   sudo rm -rf /tmp/tf/
   

Inferenza e pubblicazione del modello PyTorch

Per i modelli scritti con PyTorch, il flusso di lavoro è:

1. Scrivi un gestore del modello Python per il caricamento e l'inferenza utilizzando TorchDynamo e PyTorch/XLA
2. Usa TorchModelArchiver per creare un archivio dei modelli
3. Utilizza TorchServe per pubblicare il modello

TorchDynamo e PyTorch/XLA

TorchDynamo (Dynamo) è un compilatore JIT a livello di Python progettato per rendere più veloci i programmi PyTorch. Fornisce un'API pulita a cui agganciarsi ai backend di compilazione. Modifica dinamicamente il bytecode Python prima dell'esecuzione. Nella versione PyTorch/XLA 2.0, è presente un backend sperimentale per l'inferenza e l'addestramento utilizzando Dynamo.

Dynamo fornisce un grafico Torch FX (FX) quando riconosce un pattern del modello e PyTorch/XLA utilizza un approccio con tensore lento per compilare il grafico FX e restituire la funzione compilata. Per ulteriori informazioni su Dynamo, consulta:

• Post sulle discussioni di Pytorch Dev
• Documentazione di TorchDynamo
• PyTorch 2.0 e XLA per maggiori dettagli.

Ecco un piccolo esempio di codice dell'esecuzione dell'inferenza densenet161 con 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

Puoi utilizzare l'immagine Docker torchserve-tpu fornita per gestire il tuo modello Python archiviato su una VM TPU.

Configura l'autenticazione per Docker:

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

Esegui il pull dell'immagine Docker TorchServe Cloud TPU sulla 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}

Raccogli artefatti del modello

Per iniziare, devi fornire un gestore del modello, che indica al worker del server del modello TorchServe di caricare il modello, elaborare i dati di input ed eseguire l'inferenza. Puoi utilizzare i gestori di inferenza predefiniti di TorchServe (origine) oppure sviluppare il tuo gestore di modelli personalizzati seguendo base_handler.py. Potresti anche dover fornire il modello addestrato e il file di definizione del modello.

Nota: tutti i gestori di modelli predefiniti di TorchServe supportano PyTorch/XLA

Nel seguente esempio di Densenet 161, utilizziamo gli artefatti del modello e il gestore di classificazione delle immagini predefinito fornito da TorchServe:

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

1. Configura alcune variabili di ambiente:

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

2. Scarica e copia gli artefatti del modello dall'esempio di classificatore di immagini 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. Scarica le ponderazioni del modello:

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

4. Crea un file di configurazione del modello TorchServe per utilizzare il backend Dynamo:

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

   Dovresti vedere i seguenti file e directory:

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

Genera un file di archivio dei modelli

Per gestire il tuo modello PyTorch con Cloud TPU TorchServe, devi pacchettizzare il gestore del modello e tutti gli artefatti del modello in un file di archivio dei modelli (*.mar) utilizzando Torch Model Archiver.

Genera un file di archivio del modello con 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

Dovresti vedere il file di archivio del modello generato nella directory model-store:

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

Gestisci richieste di inferenza

Ora che hai il file di archivio dei modelli, puoi avviare il server dei modelli TorchServe e gestire le richieste di inferenza.

1. Avvia il server del modello 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. Integrità del server del modello di query:

   curl http://localhost:8080/ping
   

   Se il server del modello è in esecuzione, vedrai:

   {
     "status": "Healthy"
   }
   

   Per eseguire una query sulle versioni predefinite del modello attualmente registrato, utilizza:

   curl http://localhost:8081/models
   

   Dovresti vedere il modello registrato:

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

   Per scaricare un'immagine per l'inferenza utilizza:

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

   Per inviare una richiesta di inferenza al server del modello utilizza:

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

   Dovresti visualizzare una risposta simile alla seguente:

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

3. Log del server del modello

   Utilizza i seguenti comandi per accedere ai log:

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

   Nel log dovrebbe essere visualizzato il seguente messaggio:

   "Compiled model with backend torchxla\_trace\_once"
   

Esegui la pulizia

Arresta il container Docker:

rm -rf serve
rm -rf ${WORKDIR}

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

Profilazione

Dopo aver configurato l'inferenza, puoi utilizzare i profiler per analizzare le prestazioni e l'utilizzo di TPU. Per ulteriori informazioni sulla profilazione, vedi:

• Profilazione su Cloud TPU
• Profilazione TensorFlow
• Profilazione di PyTorch
• Profilazione JAX