Introduzione all'inferenza di Cloud TPU v5e

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

Panoramica e vantaggi

Cloud TPU v5e è un acceleratore AI sviluppato da Google e ottimizzato addestramento, ottimizzazione e distribuzione basati su trasformatore, da testo a immagine e CNN (inferenza). Le sezioni TPU v5e possono contenere fino a 256 chip.

Per pubblicazione si intende il processo di deployment di un modello di machine learning addestrato a un ambiente di produzione in cui può essere utilizzato per l'inferenza. Gli SLO di latenza sono un la priorità per la pubblicazione.

Questo documento illustra la distribuzione di un modello su una TPU a host singolo. Sezione TPU con Fino a 8 chip hanno una VM o un host TPU e sono chiamati TPU a host singolo.

Inizia

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

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

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

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

Inferenza e gestione del modello Cloud TPU

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

Inferenza e pubblicazione del modello JAX

Per gestire un modello su una VM TPU, devi:

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

Formato SavedModel

Un SavedModel contiene un programma TensorFlow completo, che include parametri e calcolo. Non richiede il codice edilizio originale eseguire.

Se il modello è stato scritto in JAX, devi usare jax2tf per serializzare del modello nel formato SavedModel.

Convertitore di inferenza

Il convertitore di inferenza di Cloud TPU prepara e ottimizza un modello esportato in Formato SavedModel per TPU l'inferenza. Puoi eseguire il convertitore di inferenza in una shell locale o nella tua VM TPU. Ti consigliamo di utilizzare la shell della VM TPU perché include tutti gli strumenti a riga di comando necessarie per eseguire il convertitore. Per ulteriori informazioni sull'Inferenza Convertitore, consulta la Guida dell'utente per la conversione di inferenza.

Requisiti del convertitore di inferenza

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

2. Devi definire un alias di funzione per la funzione TPU. Per ulteriori informazioni, consulta la Guida dell'utente del convertitore di inferenza. Negli esempi in questa guida viene utilizzato tpu_func come alias della funzione TPU.

3. Assicurati che la CPU della macchina supporti Advanced Vector eXtensions (AVX) di base, la libreria TensorFlow (la dipendenza il convertitore di inferenza di Cloud TPU) viene compilato in modo da utilizzare le istruzioni AVX. La maggior parte delle CPU il supporto AVX.

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

Inferenza e pubblicazione del modello JAX

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

1. Usa jax2tf per serializzare il tuo modello nel formato SavedModel
2. Utilizza il convertitore di inferenza per preparare il modello salvato per la pubblicazione
3. Utilizzo di TensorFlow Serving per fornire il modello

Usa jax2tf per serializzare un modello JAX nel formato SavedModel

La seguente funzione Python mostra come utilizzare jax2tf all'interno del codice del tuo 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 ulteriori informazioni su jax2tf, consulta la sezione sull'interazione tra 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 sezione Guida per la conversione di inferenza.

Utilizzo di TensorFlow Serving

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

Esempi di pubblicazione del modello JAX

Prerequisiti

1. Configura le tue credenziali Docker ed esegui il pull del convertitore di inferenza e di Cloud TPU Immagine Docker in uso:

   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:

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

3. Installa le dipendenze demo JAX:

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

Fornisci il modello JAX BERT per l'inferenza

Puoi scaricare il modello BERT preaddestrato di Hugging Face.

Nota: le esportazioni di 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 di 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 server del modello log del container 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 riportate di seguito, il server pronti per 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 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 ripulire il 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/
   

Fornisci lo Stable Diffusion della JAX per l'inferenza

Puoi scaricare il modello Stable Diffusion preaddestrato di Hugging Face.

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

   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 di 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 verificare sui server gRPC e HTTP sono 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
   

   Il copione invia "Dipinto di uno scoiattolo che pattina a New York" come prompt. L'immagine di output verrà salvata come stable_diffusion_images.jpg nel della directory corrente.

4. Eseguire la pulizia.

   Assicurati di ripulire il 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/
   

Distribuzione di TensorFlow

Le seguenti istruzioni mostrano come utilizzare TensorFlow sulle VM TPU.

Flusso di lavoro di TensorFlow Serving

1. Scarica l'immagine Docker di TensorFlow Serving per la tua 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, esegui il pull del convertitore di inferenza e Immagine Docker di 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. Scarica il codice demo:

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

4. Installa le dipendenze demo di TensorFlow:

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

5. Pubblica il tuo modello TensorFlow utilizzando TensorFlow Serving Docker 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 tuo modello.

   • API client REST
   • API client gRPC

Esegui la demo di TensorFlow ResNet-50 Serving

1. Esportare 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 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 seguenti informazioni, significa che il server è pronto a gestire le richieste. L'operazione richiede 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 della 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 ripulire il 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 di modelli Python per il caricamento e l'inferenza utilizzando TorchDynamo e PyTorch/XLA
2. Utilizza TorchModelArchiver per creare un archivio dei modelli
3. Utilizza TorchServe per pubblicare il modello

TorchDynamo e PyTorch/XLA

TorchDynamo (Dynamo) è una funzionalità di livello Python Compilatore JIT progettato per rendere più veloci i programmi PyTorch. Offre un'API pulita per collegarsi ai backend del compilatore. Modifica dinamicamente il bytecode Python prima dell'esecuzione. Nella release PyTorch/XLA 2.0 è presente un backend sperimentale per l'inferenza e l'addestramento usando Dynamo.

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

• Post delle discussioni sugli sviluppatori di Pytorch
• Documentazione di TorchDynamo
• PyTorch 2.0 e XLA per ulteriori 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 i dati archiviati un modello pytorch 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 di Cloud TPU TorchServe sulla tua 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 indichi il worker del server del modello TorchServe per caricare il modello, elaborare i dati di input e eseguire l'inferenza. Puoi utilizzare i gestori di inferenza predefiniti di TorchServe (fonte), oppure sviluppare il tuo gestore del modello personalizzato 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 TorchServe supportano PyTorch/XLA

Nel seguente esempio di Densenet 161 vengono utilizzati gli artefatti del modello e gestore classificatore di immagini 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 dal classificatore di immagini TorchServe (esempio):

   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 i pesi 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 file e le directory seguenti:

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

Genera un file di archivio del modello

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

Genera un file di archivio dei modelli 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 dei modelli generato nella directory model-store:

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

Gestisci le richieste di inferenza

Ora che hai il file di archivio dei modelli, puoi avviare il server dei modelli TorchServe e gestire 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 è attivo e in esecuzione, vedrai:

   {
     "status": "Healthy"
   }
   

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

   curl http://localhost:8081/models
   

   Dovresti vedere il modello registrato:

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

   Per scaricare un'immagine da usare per l'inferenza:

   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 vedere una risposta simile alla seguente:

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

3. Log dei server del modello

   Utilizza i comandi seguenti 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 impostato l'inferenza, puoi usare i profiler per analizzare le prestazioni e l'utilizzo della TPU. Per ulteriori informazioni sulla profilazione, consulta:

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