Introduzione all'inferenza di Cloud TPU v5e
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
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:
- Serializza il modello in TensorFlow SavedModel formato
- Utilizza il convertitore di inferenza per preparare il modello salvato per la pubblicazione
- 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
Il modello deve essere esportato da TensorFlow o JAX nel nel formato SavedModel.
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.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.
Inferenza e pubblicazione del modello JAX
Questa sezione descrive come gestire modelli JAX utilizzando jax2tf
e TensorFlow
Pubblicazione.
- Usa
jax2tf
per serializzare il tuo modello nel formato SavedModel - Utilizza il convertitore di inferenza per preparare il modello salvato per la pubblicazione
- 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
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
Connettiti alla VM TPU con SSH e installa il codice demo di inferenza:
gcloud storage cp \ "gs://cloud-tpu-inference-public/demo" \ . \ --recursive
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.
Esporta un modello salvato di TensorFlow compatibile con TPU da un modello BERT Flax:
cd demo/jax/bert python3 export_bert_model.py
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 ...
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?."
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.
Scarica il modello Stable Diffusion in un modello salvato TF2 compatibile con TPU formato:
cd demo/jax/stable_diffusion python3 export_stable_diffusion_model.py
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 ...
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.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
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}
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
Scarica il codice demo:
gcloud storage cp \ "gs://cloud-tpu-inference-public/demo" \ . \ --recursive
Installa le dipendenze demo di TensorFlow:
pip install -r ./demo/tf/requirements.txt
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}
Utilizza l'API Serving Client per eseguire query sul tuo modello.
Esegui la demo di TensorFlow ResNet-50 Serving
Esportare un modello salvato TF2 compatibile con TPU dal modello Keras ResNet-50.
cd demo/tf/resnet-50 python3 export_resnet_model.py
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 ...
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)]]
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 è:
- Scrivi un gestore di modelli Python per il caricamento e l'inferenza utilizzando
TorchDynamo
e PyTorch/XLA - Utilizza
TorchModelArchiver
per creare un archivio dei modelli - 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.
Nel seguente esempio di Densenet 161 vengono utilizzati gli artefatti del modello e gestore classificatore di immagini fornito da TorchServe:
Configura alcune variabili di ambiente:
CWD="$(pwd)" WORKDIR="${CWD}/densenet_161" mkdir -p ${WORKDIR}/model-store mkdir -p ${WORKDIR}/logs
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}
Scarica i pesi del modello:
wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth mv densenet161-8d451a50.pth ${WORKDIR}
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.
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
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 }
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: