Inferenza di Cloud TPU v5e [anteprima pubblica]

Inizia

Panoramica e vantaggi

Cloud TPU v5e è l'acceleratore AI di ultima generazione di Google Cloud. Con un ingombro di 256 chip per pod, i pod v5e sono ottimizzati per l'addestramento, il perfezionamento e la gestione basati su trasformatore, da testo a immagine e basati su CNN.

Concetti

Se non hai mai utilizzato le Cloud TPU, consulta la home page della documentazione TPU.

Patatine

Ci sono 256 chip in una singola v5e con 8 chip per host. Per ulteriori dettagli, vedi Architettura di sistema.

Core

I chip TPU hanno uno o due TensorCore per eseguire la moltiplicazione delle matrici. Analogamente ai pod v2 e v3, la v5e ha un TensorCore per chip. Al contrario, i pod v4 hanno 2 TensorCore per chip. Consulta Architettura di sistema per ulteriori dettagli su TensorCore v5e. Ulteriori informazioni su TensorCores sono disponibili in questo articolo su ACM.

Host

Un host è un computer fisico (CPU) che esegue le VM. Un host può eseguire più VM contemporaneamente.

Inferenza batch

L'inferenza batch o offline si riferisce all'inferenza al di fuori delle pipeline di produzione, in genere su una maggior parte degli input. L'inferenza batch viene utilizzata per attività offline come l'etichettatura dei dati e anche per valutare il modello addestrato. Gli SLO di latenza non sono una priorità per l'inferenza batch.

attiva

La pubblicazione si riferisce al processo di deployment di un modello di machine learning addestrato in un ambiente di produzione, dove può essere utilizzato per fare previsioni o prendere decisioni. Gli SLO di latenza sono una priorità per la gestione.

Host singolo e host multiplo

Le sezioni con meno di 8 chip utilizzano al massimo 1 host. Le sezioni con più di 8 chip hanno accesso a più di un singolo host e possono eseguire l'addestramento distribuito utilizzando più host. L'anteprima non supporta la pubblicazione multi-host. Il numero massimo di chip consentiti in un singolo job di pubblicazione è 8 chip.

Risorsa in coda

Una rappresentazione delle risorse TPU, utilizzata per accodare e gestire una richiesta per un ambiente TPU single-slice o multi-slice. Per ulteriori informazioni, consulta la guida dell'utente sulle risorse in coda.

VM Cloud TPU

Una macchina virtuale che esegue Linux e che ha accesso alle TPU sottostanti. Per le TPU v5e, ogni VM TPU ha accesso diretto a 1, 4 o 8 chip a seconda del tipo di acceleratore specificato dall'utente. Una VM TPU è nota anche come worker.

Architettura di sistema

Ogni chip v5e contiene un TensorCore. Ogni TensorCore ha 4 Matrix Multiply Unit (MXU), un'unità vettoriale e un'unità scalare. La tabella seguente mostra le specifiche principali e i relativi valori per una versione 5e. Le specifiche dei pod sono incluse nella tabella che segue le specifiche del chip.

Specifiche del chip di chiave	Valori v5e
Picco di computing per chip (bf16)	197 TFLOP
Picco di computing per chip (Int8)	393 TFLOP
Capacità e larghezza di banda HBM2	16 GB, 819 GB/s
BW Interchip Interconnect	1600 Gbps
Riduce completamente la larghezza di banda per pod	51,2 TB/s
Larghezza di banda di bisezione per pod	1,6 TB/s
Larghezza di banda della rete del data center per pod	6,4 Tbps

Specifiche chiave dei pod	Valori v5e
Dimensioni pod di TPU	256 chip
Topologia di interconnessione	Torus 2D
Picco di computing per pod	100 PetaOps(Int8)
Riduce completamente la larghezza di banda per pod	51,2 TB/s
Larghezza di banda di bisezione per pod	1,6 TB/s
Larghezza di banda della rete del data center per pod	6,4 Tbps

Chip TPU v5e

Il seguente diagramma illustra un chip TPU v5e.

disegno

Tipi di acceleratori

Cloud TPU v5e è un prodotto combinato di addestramento e inferenza. Il flag AcceleratorType viene utilizzato per distinguere l'ambiente TPU di cui è stato eseguito il provisioning per l'addestramento da quello di cui è stato eseguito il provisioning per la pubblicazione. I job di addestramento sono ottimizzati per la velocità effettiva e la disponibilità, mentre i job di addestramento sono ottimizzati per la latenza. Quindi, un job di addestramento sulle TPU di cui è stato eseguito il provisioning per la gestione potrebbe avere una disponibilità inferiore e, allo stesso modo, un job di addestramento eseguito sulle TPU di cui è stato eseguito il provisioning per l'addestramento potrebbe avere una latenza maggiore.

AcceleratorType utilizza il numero di TensorCore nel nodo per descrivere le dimensioni della sezione. AcceleratorType è una stringa formattata "v$VERSION_NUMBER-$CORES_COUNT".

Per la versione v5e sono supportate le seguenti forme di sezione 2D:

Topologia	Numero di chip TPU	Numero di host
1x1	1	8/1
2x2	4	1/2
2x4	8	1
4x4	16	2
4x8	32	4
8x8	64	8
8x16	128	16
16x16	256	32

Tipi di Cloud TPU v5e per la gestione

Esistono tipi di TPU v5e che offrono una migliore disponibilità per la gestione. Questi tipi supportano solo fino a 8 chip v5e (host singolo). Sono supportate le seguenti configurazioni: sezioni 1 x 1, 2 x 2 e 2 x 4. Ogni fetta ha rispettivamente 1, 4 e 8 fiches.

Configurazioni TPU v5e che supportano la pubblicazione: 1x1, 2x2 e 2x4.

La pubblicazione su più di 8 chip v5e, ad esempio la gestione di più host, è prevista per la roadmap di Cloud TPU e sarà disponibile in un secondo momento.

Per eseguire il provisioning delle TPU per un job di gestione, utilizza i seguenti tipi di acceleratore nella richiesta di creazione dell'interfaccia a riga di comando o delle TPU API. Tutta la capacità di gestione di v5e si trova attualmente nella zona us-west4-a:

V5E-1
V5E-4
V5E-8

Nota: la creazione della versione 5e-8 non è al momento disponibile.

Tipi di VM

Ogni VM in una sezione TPU può contenere 1, 4 o 8 chip. Le sezioni più piccole e a 4 chip avranno lo stesso nodo NUMA (per ulteriori informazioni sui nodi NUMA, vedi Architettura del sistema). Per le VM TPU a 8 chip, la comunicazione CPU-TPU sarà più efficiente all'interno delle partizioni NUMA. Ad esempio, nella figura mostrata di seguito, la comunicazione CPU0-Chip0 sarà più veloce della comunicazione CPU0-Chip4.

VM TPU a 8 chip contenente due partizioni NUMA

Confronto dei tipi di VM:

Per l'Anteprima pubblica, i clienti di inferenza hanno accesso a un chip (v5e-1), 4 chip (v5e-4) e 8 chip (v5e-8).

Tipo di VM	n2d-48-24-v5lite-tpu	n2d-192-112-v5lite-tpu	n2d-384-224-v5lite-tpu
N. di chip v5e	1	4	8
N. di vCPU	24	112	224
RAM (GB)	48	192	384
N. di NUMA nodi	1	1	2
Si applica a	V5E-1	V5E-4	V5E-8
Interruzione	Alti	Medie	Basso

Per fare spazio alle VM con più chip, gli scheduler potrebbero prerilasciare le VM con meno chip. Per questo motivo, è probabile che le VM a 8 chip prerilasciano le VM a 1 e 4 chip.

Inizia

Capacità di protezione

Cloud TPU v5e è ora in Anteprima pubblica. Contatta il team di vendita Cloud per iniziare a utilizzare Cloud TPU v5e per i tuoi carichi di lavoro di AI.

Prepara un progetto Google Cloud

Accedi al tuo Account Google. Se non l'hai ancora fatto, crea un nuovo account.
Nella console Cloud, seleziona o crea un progetto Cloud dalla pagina del selettore dei progetti.
La configurazione della fatturazione è obbligatoria per tutti gli utilizzi di Google Cloud, quindi assicurati che la fatturazione sia abilitata per il tuo progetto.

La fatturazione per l'utilizzo dell'Anteprima pubblica segue i prezzi regionali standard mostrati nella pagina dei prezzi di Cloud TPU.
Installa i componenti alpha di gcloud.
Abilita l'API TPU utilizzando il seguente comando gcloud in Cloud Shell. Puoi abilitarlo anche dalla console Google Cloud.
```
gcloud services enable tpu.googleapis.com
```
Abilita l'account di servizio TPU.

Gli account di servizio consentono al servizio Cloud TPU di accedere ad altri servizi Google Cloud. Un account di servizio gestito dall'utente è una pratica Google Cloud consigliata. Segui queste guide per creare e concedere i ruoli seguenti al tuo account di servizio. Sono necessari i seguenti ruoli:
- Amministratore TPU
- Amministratore Storage
- Writer log
- Writer metriche Monitoring

Configurare il progetto e la zona.

L'ID progetto è il nome del progetto visualizzato nella console Cloud. La zona predefinita per Cloud TPU v5e ia us-west4-a.

export PROJECT_ID=project-ID
export ZONE=us-west4-a

gcloud alpha compute tpus tpu-vm service-identity create --zone=${ZONE}

gcloud auth login
gcloud config set project ${PROJECT}
gcloud config set compute/zone ${ZONE}

Richiedi l'accesso agli artefatti di anteprima pubblica dell'inferenza di Cloud TPU.

Invia l'account di servizio che prevedi di utilizzare per la VM TPU al team TPU. Quello predefinito è l'account di servizio predefinito di Compute Engine (${PROJECT_ID}-compute@developer.gserviceaccount.com) che può essere trovato nella pagina IAM.
Esegui il provisioning dell'ambiente Cloud TPU v5e.

Una versione v5e viene gestita come risorsa in coda. È possibile eseguire il provisioning della capacità utilizzando il comando queued-resource create.

Crea variabili di ambiente per ID progetto, tipo di acceleratore, zona, versione di runtime e nome TPU.
```
export PROJECT_ID=project_ID
export ACCELERATOR_TYPE=v5e-1
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
```
Descrizioni delle variabili

PROJECT_ID
Nome progetto. Utilizza il nome del tuo progetto Google.
ACCELERATOR_TYPE
Consulta la sezione Tipi di acceleratore per i tipi di acceleratore supportati.
ZONE
L'intera capacità di Anteprima pubblica sarà nel formato us-west4-a.
RUNTIME_VERSION
v2-alpha-tpuv5-lite
SERVICE_ACCOUNT
Questo è l'indirizzo email dell'account di servizio che puoi trovare in Google Cloud Console -> IAM -> Account di servizio Ad esempio: tpu-service-account@myprojectID.iam.gserviceaccount.com.
TPU_NAME
L'ID della risorsa TPU assegnato dall'utente viene creato quando viene allocata la richiesta di risorsa TPU.
QUEUED_RESOURCE_ID
L'ID assegnato dall'utente della richiesta di risorsa in coda.

Creare una risorsa TPU.

gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
--node-id ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--accelerator-type ${ACCELERATOR_TYPE} \
--runtime-version ${RUNTIME_VERSION} \
--service-account ${SERVICE_ACCOUNT} \
--reserved

Se vuoi eliminare la risorsa che hai prenotato, devi prima eliminare la risorsa TPU_NAME, quindi anche la richiesta di risorsa in coda.

gcloud alpha compute tpus delete $TPU_NAME --zone ${ZONE} --project ${PROJECT_ID}
gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} \
--zone ${ZONE}

SSH nelle VM TPU

Per eseguire il codice sulle VM TPU, devi SSH in ogni VM TPU. In questo esempio, con una versione v5e-1, è presente una sola VM TPU.
```
gcloud compute config-ssh
gcloud compute tpus tpu-vm ssh $TPU_NAME --zone $ZONE --project $PROJECT_ID
```

Gestisci le VM TPU

Per tutte le opzioni di gestione delle TPU per le VM TPU, vedi Gestione delle TPU.

Sviluppa ed esegui

Questa sezione descrive il processo di configurazione generale per l'inferenza dei modelli personalizzati utilizzando JAX o PyTorch su Cloud TPU v5e. Il supporto di TensorFlow verrà abilitato a breve.

Per istruzioni sull'addestramento di v5e, consulta la guida per l'addestramento su v5e.

Esecuzione di Inference su v5e

Stack software di inferenza

I dettagli dello stack SW di inferenza sono trattati nelle sezioni seguenti. Questo documento è incentrato sulla gestione di un singolo host per i modelli addestrati con JAX, TensorFlow (TF) e PyTorch.

Questa sezione presuppone che tu abbia già configurato il tuo progetto Google Cloud secondo le istruzioni riportate in Preparare un progetto Google Cloud.

Inferenza e pubblicazione del modello JAX

La seguente sezione illustra il flusso di lavoro per l'inferenza del modello JAX. Come mostrato nel diagramma, esistono due percorsi per l'inferenza JAX. In questa sezione analizzeremo il percorso di produzione dei modelli JAX tramite jax2tf e TensorFlow Serving.

Utilizza jax2tf per convertire il modello in TensorFlow 2 e salvarlo
Utilizza il convertitore di inferenza per convertire il modello salvato
Utilizza TensorFlow Serving per gestire il modello

Utilizza `jax2tf` per convertire il modello e salvarlo

Fai riferimento all'interoperazione JAX e TensorFlow per convertire e salvare il tuo modello JAX in TensorFlow.

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

Utilizza il convertitore di inferenza per convertire il modello salvato

I passaggi per il convertitore di inferenza sono descritti nella Guida al convertitore di inferenza.

Utilizzo della pubblicazione su TensorFlow

I passaggi per la pubblicazione con TensorFlow sono descritti in Gestione di TensorFlow.

Esempio di pubblicazione di un modello JAX E2E

Prerequisito:

Devi configurare le credenziali Docker ed eseguire il pull dell'inference Converter e dell'immagine Docker che fornisce TensorFlow. Se non lo hai già fatto, esegui questi comandi:

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: SSH sulla VM TPU e installa il codice demo di inferenza.

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

Installa le dipendenze demo JAX Sulla VM TPU, installa requirements.txt.

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

Esegui la demo di JAX BERT E2E

Il modello BERT preaddestrato proviene da Hugging Face.

Esporta un modello salvato TF2 compatibile con TPU da un modello BERT Flax:
```
cd demo/jax/bert
```
```
python3 export_bert_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/jax/bert_tpu,target=/models/bert \
 -e MODEL_NAME=bert \
 us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Controlla il log del contenitore 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 seguenti informazioni, significa che il server è pronto per 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.

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 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/
```

Esegui la demo di pubblicazione di JAX Stable Diffusion E2E

Il modello di Diffusione stabile preaddestrato è di Hugging Face.

Esporta il modello salvato TF2 compatibile con TPU dal modello Flax Stable Diffusion:

Nota: l'esportazione richiede circa 3 minuti.
```
cd demo/jax/stable_diffusion
```
```
python3 export_stable_diffusion_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/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

Controlla il log del contenitore del server del modello e assicurati che gRPC e Http Server 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 per gestire le richieste. L'operazione richiede circa 2 minuti.

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.
```
python3 stable_diffusion_request.py
```
Il prompt è "Dipinto di uno scoiattolo che pattina a New York" e l'immagine di output verrà salvata come stable_diffusion_images.jpg nella directory attuale.
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/
```
Nota: questo è solo un esempio di riferimento e non è una versione del modello ottimizzata per le prestazioni che può essere utilizzata per casi d'uso di benchmarking o produzione.

Inferenza e pubblicazione del modello TensorFlow

Le sezioni seguenti illustrano il flusso di lavoro per l'inferenza del modello TensorFlow.

Utilizza il convertitore di inferenza per convertire il modello
Utilizza TensorFlow Serving per gestire il modello

Convertitore di inferenza

Cloud TPU Inference Converter prepara e ottimizza un modello esportato da TensorFlow o JAX per l'inferenza TPU. Il convertitore viene eseguito in una shell locale o nella shell della VM TPU. La shell della VM TPU è consigliata perché è preinstallata con gli strumenti a riga di comando necessari per il convertitore. Per ulteriori dettagli sul convertitore di inferenza, consulta la guida dell'utente sul convertitore di inferenza.

Prerequisiti

Il modello deve essere esportato da TensorFlow o JAX nel formato SaveModel.
Il modello deve avere un alias di funzione per la funzione TPU. Per istruzioni su come eseguire questa operazione, consulta gli esempi di codice nella Guida dell'utente del convertitore di inferenza. Gli esempi seguenti utilizzano tpu_func come alias della funzione TPU.
Assicurati che la CPU della tua macchina supporti le istruzioni AVX (Advanced Vector eXtensions), poiché la libreria Tensorflow (la dipendenza di Cloud TPU Inference Converter) è compilata per utilizzare le istruzioni AVX. La maggior parte delle CPU supporta AVX.
- Puoi eseguire lscpu | grep avx per verificare se il set di istruzioni AVX è supportato.

Per iniziare

Configura l'ambiente VM TPU Configura l'ambiente seguendo questi passaggi a seconda della shell che stai utilizzando:
Shell VM TPU
- Nella shell della VM TPU, esegui i comandi seguenti per consentire l'utilizzo di Docker non root:
```
sudo usermod -a -G docker ${USER}
newgrp docker
```
- Inizializza gli helper delle credenziali Docker:
```
gcloud auth configure-docker \
  us-docker.pkg.dev
```
Shell locale
Nella shell locale, configura l'ambiente seguendo questi passaggi:
- Installare Cloud SDK, che include lo strumento a riga di comando gcloud.
- Installa Docker:
- Consenti l'utilizzo di Docker non root:
```
sudo usermod -a -G docker ${USER}
newgrp docker
```
- Accedi al tuo ambiente:
```
gcloud auth login
```
- Inizializza gli helper delle credenziali Docker:
```
gcloud auth configure-docker \
us-docker.pkg.dev
```

Esegui il pull dell'immagine Docker del convertitore di inferenza:

  CONVERTER_IMAGE=us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
  docker pull ${CONVERTER_IMAGE}

Immagine convertitore

L'immagine è utilizzata per le conversioni del modello una tantum. Imposta i percorsi del modello e modifica le opzioni del convertitore in base alle tue esigenze. La sezione Esempi di utilizzo della Guida dell'utente per il convertitore di inferenza illustra diversi casi d'uso comuni.

docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      function_alias: "tpu_func"
    }
    batch_options {
      num_batch_threads: 2
      max_batch_size: 8
      batch_timeout_micros: 5000
      allowed_batch_sizes: 2
      allowed_batch_sizes: 4
      allowed_batch_sizes: 8
      max_enqueued_batches: 10
    }
'

La sezione seguente mostra come eseguire questo modello con un server modelli TensorFlow.

Pubblicazione su TensorFlow

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

Prerequisiti:

Configura le tue credenziali di Docker, se non l'hai ancora fatto:

Scarica l'immagine Docker che fornisce TensorFlow per 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}

Gestisci il tuo modello TensorFlow utilizzando l'immagine Docker che fornisce 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}

Segui l'API Serving Client per eseguire query sul tuo modello.
- API client REST
- API client gRPC

Esempio di pubblicazione del modello TensorFlow di E2E:

Prerequisito: assicurati di aver già configurato le credenziali Docker ed eseguito il pull dell'inference Converter e dell'immagine Docker che gestisce TensorFlow. In caso contrario, esegui questi comandi:

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:

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

Installa le dipendenze demo di TensorFlow:

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

Esegui la demo di pubblicazione di TensorFlow ResNet-50 E2E

Esporta 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 contenitore del server del modello e assicurati che gRPC e Http Server 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 per 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 richiesta è una banana tratta 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 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

Le seguenti sezioni descrivono il flusso di lavoro per l'inferenza del modello PyTorch:

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 gestire il modello

TorchDynamo e PyTorch/XLA

TorchDynamo (Dynamo) è un compilatore JIT a livello di Python progettato per rendere più veloci i programmi PyTorch non modificati. Fornisce un'API pulita per l'hook dei backend del compilatore. La sua funzionalità principale consiste nel modificare dinamicamente il bytecode Python appena prima dell'esecuzione. Nella versione PyTorch/XLA 2.0, viene fornito un backend sperimentale per Dynamo sia per l'inferenza che per l'addestramento.

Dynamo fornisce un grafico Torch FX (FX) quando riconosce un pattern del modello e PyTorch/XLA utilizza un approccio Lazy Tensor per compilare il grafico FX e restituire la funzione compilata. Per ulteriori informazioni relative ai dettagli tecnici sull'implementazione della dinamo di PyTorch/XLA, consulta il post della discussione sugli sviluppatori di Pytorch e la documentazione di TorchDynamo. Per ulteriori dettagli, consulta questo blog.

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

L'immagine Docker di Cloud TPU TorchServe consente di gestire il modello in modalità Eager di PyTorch utilizzando TorchServe su una VM Cloud TPU.

Puoi utilizzare l'immagine Docker torchserve-tpu fornita, pronta per gestire il modello Python archiviato su una VM Cloud 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 VM TPU:

CLOUD_TPU_TORCHSERVE_IMAGE_URL=us-docker.pkg.dev/cloud-tpu-images/inference/torchserve-tpu:v0.8.2-20230829
docker pull ${CLOUD_TPU_TORCHSERVE_IMAGE_URL}

Raccogli gli artefatti del modello

Per iniziare, devi fornire un gestore del modello, che indica al worker server del modello TorchServe di caricare il modello, elaborare i dati di input ed eseguire l'inferenza. Puoi utilizzare i gestori di inferenza predefiniti TorchServe (origine) o sviluppare il tuo gestore di modelli 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, utilizziamo gli artefatti del modello e il gestore predefinito del classificatore di immagini fornito da TorchServe:

La directory di lavoro è mostrata di seguito.

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 le ponderazioni 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 visualizzare i file e le directory mostrati di seguito:

>> 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 gestire il modello PyTorch con Cloud TPU TorchServe, devi pacchettizzare il gestore dei modelli e tutti gli artefatti dei modelli in un file di archivio 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 visualizzare il file di archivio del modello generato nella directory dell'archivio modelli:

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

Richieste di inferenza di elaborazione

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

Avvia il server 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 è in esecuzione, vedrai:

{
 "status": "Healthy"
}

Per eseguire 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 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 del server del modello

Utilizza i seguenti comandi per accedere ai log:
```
ls ${WORKDIR}/logs/
cat ${WORKDIR}/logs/model_log.log
```
Nel log dovresti visualizzare il seguente messaggio:

"Modello compilato con torchxla_trace_once di backend"

Esegui la pulizia

Arresta il container Docker:

rm -rf serve
rm -rf ${WORKDIR}

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

Gestione di modelli linguistici di grandi dimensioni (LLM)

SAX è un framework di gestione per supportare la gestione di modelli di grandi dimensioni che potrebbero richiedere TPU su più host per l'esecuzione con GSPMD, ad esempio modelli di linguaggio di grandi dimensioni basati su PAX. PAX è un framework integrato in JAX per l'addestramento di modelli su larga scala che consente sperimentazione e parallelizzazione avanzate e completamente configurabili.

La sezione sul cluster SAX descrive gli elementi chiave per comprendere il funzionamento di SAX. La sezione relativa alla pubblicazione del modello SAX illustra un esempio di pubblicazione di un modello host singolo con un modello GPTJ6B. SAX offre inoltre gestione multi-host sulle Cloud TPU e gli utenti possono eseguire modelli su topologie TPU più grandi per l'anteprima sperimentale di gestione multi-host. L'esempio seguente con un modello di test 175B mostra come sperimentare questa configurazione.

Cluster SAX (cella SAX)

Il server di amministrazione SAX e il server modello SAX sono due componenti essenziali che eseguono un cluster SAX.

Server di amministrazione SAX

Il server di amministrazione SAX monitora e coordina tutti i server modello SAX in un cluster SAX. In un cluster SAX, puoi avviare più server di amministrazione SAX in cui solo uno dei server di amministrazione SAX è attivo tramite l'elezione del leader, mentre gli altri sono server in standby. In caso di errore del server di amministrazione attivo, un server di amministrazione in standby diventerà attivo. Il server di amministrazione SAX attivo assegna repliche dei modelli e richieste di inferenza ai server dei modelli SAX disponibili.

Bucket di archiviazione amministratore SAX

Ogni cluster SAX richiede un bucket Cloud Storage per archiviare le configurazioni e le località dei server di amministrazione SAX e dei server modello SAX nel cluster SAX.

Server del modello SAX

Il server del modello SAX carica un checkpoint del modello ed esegue l'inferenza con GSPMD. Un server modello SAX viene eseguito su un singolo worker VM TPU. La gestione di modelli TPU con host singolo richiede un singolo server di modelli SAX su una VM TPU con host singolo. La gestione di modelli TPU multi-host richiede un gruppo di server modello SAX su una sezione TPU multi-host.

Pubblicazione del modello SAX

La sezione seguente illustra il flusso di lavoro per la gestione di modelli linguistici utilizzando SAX. Utilizza il modello GPT-J 6B come esempio per la pubblicazione di modelli con un solo host e un modello di test 175B per la pubblicazione di modelli multi-host.

Prima di iniziare, installa le immagini Docker SAX di Cloud TPU sulla VM TPU:

sudo usermod -a -G docker ${USER}
newgrp docker

gcloud auth configure-docker us-docker.pkg.dev

SAX_ADMIN_SERVER_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server"
SAX_MODEL_SERVER_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server"
SAX_UTIL_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/inference/sax-util"

SAX_VERSION=v1.0.0

export SAX_ADMIN_SERVER_IMAGE_URL=${SAX_ADMIN_SERVER_IMAGE_NAME}:${SAX_VERSION}
export SAX_MODEL_SERVER_IMAGE_URL=${SAX_MODEL_SERVER_IMAGE_NAME}:${SAX_VERSION}
export SAX_UTIL_IMAGE_URL="${SAX_UTIL_IMAGE_NAME}:${SAX_VERSION}"

docker pull ${SAX_ADMIN_SERVER_IMAGE_URL}
docker pull ${SAX_MODEL_SERVER_IMAGE_URL}
docker pull ${SAX_UTIL_IMAGE_URL}

Imposta altre variabili che utilizzerai in seguito:

export SAX_ADMIN_SERVER_DOCKER_NAME="sax-admin-server"
export SAX_MODEL_SERVER_DOCKER_NAME="sax-model-server"
export SAX_CELL="/sax/test"

Esempio di pubblicazione del modello con un solo host GPT-J 6B

La pubblicazione di modelli con host singolo è applicabile alla sezione TPU con host singolo, ovvero v5litepod-1, v5litepod-4 e v5litepod-8.

Crea un cluster SAX

Crea un bucket Cloud Storage per il cluster SAX:

SAX_ADMIN_STORAGE_BUCKET=${your_admin_storage_bucket}

gcloud storage buckets create gs://${SAX_ADMIN_STORAGE_BUCKET} \
--project=${PROJECT_ID}

Potrebbe essere necessario un altro bucket Cloud Storage Storage per archiviare il checkpoint.

SAX_DATA_STORAGE_BUCKET=${your_data_storage_bucket}

Connettiti alla VM TPU tramite SSH in un terminale per avviare il server di amministrazione SAX:

docker run \
--name ${SAX_ADMIN_SERVER_DOCKER_NAME} \
-it \
-d \
--rm \
--network host \
--env GSBUCKET=${SAX_ADMIN_STORAGE_BUCKET} \
${SAX_ADMIN_SERVER_IMAGE_URL}

Puoi controllare il log Docker in questo modo:

docker logs -f ${SAX_ADMIN_SERVER_DOCKER_NAME}

L'output nel log sarà simile al seguente:

I0829 01:22:31.184198       7 config.go:111] Creating config fs_root: "gs://test_sax_admin/sax-fs-root"
I0829 01:22:31.347883       7 config.go:115] Created config fs_root: "gs://test_sax_admin/sax-fs-root"
I0829 01:22:31.360837      24 admin_server.go:44] Starting the server
I0829 01:22:31.361420      24 ipaddr.go:39] Skipping non-global IP address 127.0.0.1/8.
I0829 01:22:31.361455      24 ipaddr.go:39] Skipping non-global IP address ::1/128.
I0829 01:22:31.361462      24 ipaddr.go:39] Skipping non-global IP address fe80::4001:aff:fe8e:fc8/64.
I0829 01:22:31.361469      24 ipaddr.go:39] Skipping non-global IP address fe80::42:bfff:fef9:1bd3/64.
I0829 01:22:31.361474      24 ipaddr.go:39] Skipping non-global IP address fe80::20fb:c3ff:fe5b:baac/64.
I0829 01:22:31.361482      24 ipaddr.go:56] IPNet address 10.142.15.200
I0829 01:22:31.361488      24 ipaddr.go:56] IPNet address 172.17.0.1
I0829 01:22:31.456952      24 admin.go:305] Loaded config: fs_root: "gs://test_sax_admin/sax-fs-root"
I0829 01:22:31.609323      24 addr.go:105] SetAddr /gcs/test_sax_admin/sax-root/sax/test/location.proto "10.142.15.200:10000"
I0829 01:22:31.656021      24 admin.go:325] Updated config: fs_root: "gs://test_sax_admin/sax-fs-root"
I0829 01:22:31.773245      24 mgr.go:781] Loaded manager state
I0829 01:22:31.773260      24 mgr.go:784] Refreshing manager state every 10s
I0829 01:22:31.773285      24 admin.go:350] Starting the server on port 10000
I0829 01:22:31.773292      24 cloud.go:506] Starting the HTTP server on port 8080

Avvia un server modello SAX con un solo host nel cluster SAX:

A questo punto, il cluster SAX contiene solo il server di amministrazione SAX. Puoi connetterti alla VM TPU tramite SSH in un secondo terminale per avviare un server modello SAX nel cluster SAX:

docker run \
    --privileged  \
    -it \
    -d \
    --rm \
    --network host \
    --name ${SAX_MODEL_SERVER_DOCKER_NAME} \
    --env SAX_ROOT=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
    ${SAX_MODEL_SERVER_IMAGE_URL} \
       --sax_cell=${SAX_CELL} \
       --port=10001 \
       --platform_chip=tpuv4 \
       --platform_topology=1x1

Puoi controllare il log Docker in questo modo:

docker logs -f ${SAX_MODEL_SERVER_DOCKER_NAME}

Converti checkpoint del modello:

Devi installare PyTorch e Transformers per scaricare il checkpoint GPT-J da EleutherAI:

pip3 install accelerate
pip3 install torch
pip3 install transformers

Per convertire il checkpoint in checkpoint SAX, devi installare paxml:

pip3 install paxml==1.1.0

Quindi, imposta la seguente variabile:

>>PT_CHECKPOINT_PATH=./fine_tuned_pt_checkpoint

Scarica il checkpoint PyTorch ottimizzato su ${PT_CHECKPOINT_PATH}, segui https://github.com/mlcommons/inference/blob/master/language/gpt-j/README.md#download-gpt-j-model ed esegui i seguenti comandi:

ls ${PT_CHECKPOINT_PATH}

che dovrebbe elencare quanto segue:

added_tokens.json  generation_config.json
pytorch_model.bin.index.json
pytorch_model-00001-of-00003.bin
pytorch_model-00002-of-00003.bin
pytorch_model-00003-of-00003.bin
special_tokens_map.json
trainer_state.json
config.json
merges.txt
tokenizer_config.json
vocab.json

Il seguente script converte il checkpoint GPT-J in checkpoint SAX:

python3 -m convert_gptj_ckpt --base EleutherAI/gpt-j-6b --pax pax_6b

wget https://raw.githubusercontent.com/google/saxml/main/saxml/tools/convert_gptj_ckpt.py
python3 -m convert_gptj_ckpt --base ${PT_CHECKPOINT_PATH} --pax .

Dovrebbero essere generati contenuti simili ai seguenti:

transformer.wte.weight (50401, 4096)
transformer.h.0.ln_1.weight (4096,)
transformer.h.0.ln_1.bias (4096,)
transformer.h.0.attn.k_proj.weight (4096, 4096)
.
.
.
transformer.ln_f.weight (4096,)
transformer.ln_f.bias (4096,)
lm_head.weight (50401, 4096)
lm_head.bias (50401,)
Saving the pax model to .
done

Al termine della conversione, inserisci il seguente comando:

ls checkpoint_00000000/

che dovrebbe elencare quanto segue:

metadate
state

Devi creare un file commit_success e collocarlo nelle directory secondarie:

CHECKPOINT_PATH=gs://${SAX_DATA_STORAGE_BUCKET}/path/to/checkpoint_00000000
gsutil -m cp -r checkpoint_00000000 ${CHECKPOINT_PATH}

touch commit_success.txt
gsutil cp commit_success.txt ${CHECKPOINT_PATH}/
gsutil cp commit_success.txt ${CHECKPOINT_PATH}/metadata/
gsutil cp commit_success.txt ${CHECKPOINT_PATH}/state/

Pubblica il modello nel cluster SAX

Ora puoi pubblicare GPT-J con il checkpoint convertito nel passaggio precedente.

MODEL_NAME=gptjtokenizedbf16bs32
MODEL_CONFIG_PATH=saxml.server.pax.lm.params.gptj.GPTJ4TokenizedBF16BS32
REPLICA=1

Per pubblicare il tag GPT-J (e i passaggi successivi), utilizza SSH per connetterti alla VM TPU in un terzo terminale:

docker run \
 ${SAX_UTIL_IMAGE_URL} \
   --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
   publish \
     ${SAX_CELL}/${MODEL_NAME} \
     ${MODEL_CONFIG_PATH} \
     ${CHECKPOINT_PATH} \
     ${REPLICA}

Vedrai molta attività dal log Docker del server del modello fino a quando non vedi qualcosa di simile al seguente per indicare che il modello è stato caricato correttamente:

I0829 01:33:49.287459 139865140229696 servable_model.py:697] loading completed.

Genera risultati di inferenza

Per GPT-J, i valori di input e output devono essere formattati come stringa ID token separata da virgole. Dovrai tokenizzare l'input di testo.

TEXT = ("Below is an instruction that describes a task, paired with "
"an input that provides further context. Write a response that "
"appropriately completes the request.\n\n### Instruction\:\nSummarize the "
"following news article\:\n\n### Input\:\nMarch 10, 2015 . We're truly "
"international in scope on Tuesday. We're visiting Italy, Russia, the "
"United Arab Emirates, and the Himalayan Mountains. Find out who's "
"attempting to circumnavigate the globe in a plane powered partially by the "
"sun, and explore the mysterious appearance of craters in northern Asia. "
"You'll also get a view of Mount Everest that was previously reserved for "
"climbers. On this page you will find today's show Transcript and a place "
"for you to request to be on the CNN Student News Roll Call. TRANSCRIPT . "
"Click here to access the transcript of today's CNN Student News program. "
"Please note that there may be a delay between the time when the video is "
"available and when the transcript is published. CNN Student News is "
"created by a team of journalists who consider the Common Core State "
"Standards, national standards in different subject areas, and state "
"standards when producing the show. ROLL CALL . For a chance to be "
"mentioned on the next CNN Student News, comment on the bottom of this page "
"with your school name, mascot, city and state. We will be selecting "
"schools from the comments of the previous show. You must be a teacher or a "
"student age 13 or older to request a mention on the CNN Student News Roll "
"Call! Thank you for using CNN Student News!\n\n### Response\:")

Puoi ottenere la stringa degli ID token tramite il tokenizzatore EleutherAI/gpt-j-6b:

from transformers import GPT2Tokenizer
tokenizer = GPT2Tokenizer.from_pretrained(EleutherAI/gpt-j-6b)

Tokenizza il testo di input:

encoded_example = tokenizer(TEXT)
input_ids = encoded_example.input_ids
INPUT_STR = ",".join([str(input_id) for input_id in input_ids])

Puoi aspettarti una stringa ID token simile alla seguente:

>>> INPUT_STR
'21106,318,281,12064,326,8477,257,4876,11,20312,351,281,5128,326,3769,2252,4732,13,19430,257,2882,326,20431,32543,262,2581,13,198,198,21017,46486,25,198,13065,3876,1096,262,1708,1705,2708,25,198,198,21017,23412,25,198,16192,838,11,1853,764,775,821,4988,3230,287,8354,319,3431,13,775,821,10013,8031,11,3284,11,262,1578,4498,24880,11,290,262,42438,22931,21124,13,9938,503,508,338,9361,284,2498,4182,615,10055,262,13342,287,257,6614,13232,12387,416,262,4252,11,290,7301,262,11428,5585,286,1067,8605,287,7840,7229,13,921,1183,635,651,257,1570,286,5628,41336,326,373,4271,10395,329,39311,13,1550,428,2443,345,481,1064,1909,338,905,42978,290,257,1295,329,345,284,2581,284,307,319,262,8100,13613,3000,8299,4889,13,48213,6173,46023,764,6914,994,284,1895,262,14687,286,1909,338,8100,13613,3000,1430,13,4222,3465,326,612,743,307,257,5711,1022,262,640,618,262,2008,318,1695,290,618,262,14687,318,3199,13,8100,13613,3000,318,2727,416,257,1074,286,9046,508,2074,262,8070,7231,1812,20130,11,2260,5423,287,1180,2426,3006,11,290,1181,5423,618,9194,262,905,13,15107,3069,42815,764,1114,257,2863,284,307,4750,319,262,1306,8100,13613,3000,11,2912,319,262,4220,286,428,2443,351,534,1524,1438,11,37358,11,1748,290,1181,13,775,481,307,17246,4266,422,262,3651,286,262,2180,905,13,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,0,6952,345,329,1262,8100,13613,3000,0,198,198,21017,18261,25'

Per generare un riepilogo per il tuo articolo:

docker run \
  ${SAX_UTIL_IMAGE_URL} \
    --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
    lm.generate \
      ${SAX_CELL}/${MODEL_NAME} \
      ${INPUT_STR}

Puoi aspettarti qualcosa di simile a:

+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+
|                                                                                                                                                    GENERATE                                                                                                                                                    |    SCORE     |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+
| 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,220,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,13,50256 | -0.023136413 |
| 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,220,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,0,50256  |  -0.91842502 |
| 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,13,50256     |   -1.1726116 |
| 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,220,921,1276,307,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,13,50256                            |   -1.2472695 |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+

Per annullare la tokenizzazione della stringa ID token di output:

output_token_ids = [int(token_id) for token_id in OUTPUT_STR.split(',')]
OUTPUT_TEXT = tokenizer.decode(output_token_ids, skip_special_tokens=True)

Il testo detokenizzato può essere visualizzato come:

>>> OUTPUT_TEXT
'This page includes the show Transcript.\nUse the Transcript to help
students with reading comprehension and vocabulary.\nAt the bottom of
the page, comment for a chance to be mentioned on CNN Student News.
You must be a teacher or a student age 13 or older to request a mention on the CNN Student News Roll Call.'

Pulisci i container Docker e i bucket di archiviazione Cloud Storage.

Anteprima della pubblicazione del modello multi-host da 175 miliardi

Alcuni degli LLM richiederanno una sezione TPU multi-host, ovvero v5litepod-16 e versioni successive. In questi casi, tutti gli host TPU multi-host devono disporre di una copia di un server modello SAX e tutti i server modello funzionano come gruppo di server modello SAX per gestire il modello di grandi dimensioni su una sezione TPU multi-host.

Crea un nuovo cluster SAX

Puoi seguire lo stesso passaggio della procedura di creazione di un cluster SAX nella procedura guidata su GPT-J per creare un nuovo cluster SAX e un server di amministrazione SAX.

In alternativa, se hai già un cluster SAX, puoi avviare un server modello multi-host nel tuo cluster SAX.

Avvia un server modello SAX multi-host in un cluster SAX

Utilizza lo stesso comando per creare una sezione TPU multi-host che utilizzi per una sezione TPU con host singolo, specifica semplicemente il tipo di acceleratore multi-host appropriato:

ACCELERATOR_TYPE=v5litepod-32
ZONE=us-east1-c

gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
  --node-id ${TPU_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE} \
  --accelerator-type ${ACCELERATOR_TYPE} \
  --runtime-version ${RUNTIME_VERSION} \
  --service-account ${SERVICE_ACCOUNT} \
  --reserved

Per eseguire il pull dell'immagine del server del modello SAX per tutti gli host/worker TPU e avviarli:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE} \
  --worker=all \
  --command="
    gcloud auth configure-docker \
      us-docker.pkg.dev
    # Pull SAX model server image
    docker pull ${SAX_MODEL_SERVER_IMAGE_URL}
    # Run model server
    docker run \
      --privileged  \
      -it \
      -d \
      --rm \
      --network host \
      --name ${SAX_MODEL_SERVER_DOCKER_NAME} \
      --env SAX_ROOT=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
      ${SAX_MODEL_SERVER_IMAGE_URL} \
        --sax_cell=${SAX_CELL} \
        --port=10001 \
        --platform_chip=tpuv4 \
        --platform_topology=1x1"

Pubblica il modello nel cluster SAX

In questo esempio viene utilizzato un modello LmCloudSpmd175B32Test:

MODEL_NAME=lmcloudspmd175b32test
MODEL_CONFIG_PATH=saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test
CHECKPOINT_PATH=None
REPLICA=1

Per pubblicare il modello di test:

docker run \
 ${SAX_UTIL_IMAGE_URL} \
   --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
   publish \
     ${SAX_CELL}/${MODEL_NAME} \
     ${MODEL_CONFIG_PATH} \
     ${CHECKPOINT_PATH} \
     ${REPLICA}

Genera risultati di inferenza

docker run \
  ${SAX_UTIL_IMAGE_URL} \
    --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
    lm.generate \
      ${SAX_CELL}/${MODEL_NAME} \
      "Q:  Who is Harry Porter's mother? A\: "

Tieni presente che, poiché questo esempio utilizza un modello di test con ponderazioni casuali, l'output potrebbe non essere significativo.

Esegui la pulizia

Interrompi i container Docker:
```
docker stop ${SAX_ADMIN_SERVER_DOCKER_NAME}
docker stop ${SAX_MODEL_SERVER_DOCKER_NAME}
```
Elimina il bucket di archiviazione dell'amministratore di Cloud Storage e qualsiasi bucket di archiviazione dati utilizzando gsutil, come mostrato di seguito.
```
gsutil rm -rf gs://${SAX_ADMIN_STORAGE_BUCKET}
gsutil rm -rf gs://${SAX_DATA_STORAGE_BUCKET}
```

Profilazione

Dopo aver configurato l'inferenza, i profiler possono essere utilizzati per analizzare le prestazioni e l'utilizzo di TPU. Di seguito sono riportati i riferimenti ad alcuni documenti correlati alla profilazione:

Assistenza e feedback

Tutti i feedback sono ben accetti. Per condividere feedback o richiedere assistenza, contattaci qui o invia un'email all'indirizzo cloudtpu-preview-support@google.com

Termini

Tutte le informazioni che Google ti ha fornito in merito a questa Anteprima pubblica sono informazioni riservate di Google e sono soggette alle disposizioni di riservatezza previste dai Termini di servizio di Google Cloud Platform (o da altri contratti che regolano l'utilizzo di Google Cloud Platform da parte dell'utente).

Inferenza di Cloud TPU v5e [anteprima pubblica]

Inizia

Panoramica e vantaggi

Concetti

Patatine

Core

Host

Inferenza batch

attiva

Host singolo e host multiplo

Risorsa in coda

VM Cloud TPU

Architettura di sistema

Chip TPU v5e

Tipi di acceleratori

Tipi di Cloud TPU v5e per la gestione

Tipi di VM

Confronto dei tipi di VM:

Inizia

Capacità di protezione

Prepara un progetto Google Cloud

Descrizioni delle variabili

Gestisci le VM TPU

Sviluppa ed esegui

Esecuzione di Inference su v5e

Stack software di inferenza

Inferenza e pubblicazione del modello JAX

Utilizza jax2tf per convertire il modello e salvarlo

Utilizza il convertitore di inferenza per convertire il modello salvato

Utilizzo della pubblicazione su TensorFlow

Esempio di pubblicazione di un modello JAX E2E

Prerequisito:

Esegui la demo di JAX BERT E2E

Esegui la demo di pubblicazione di JAX Stable Diffusion E2E

Inferenza e pubblicazione del modello TensorFlow

Convertitore di inferenza

Prerequisiti

Per iniziare

Shell VM TPU

Shell locale

Immagine convertitore

Pubblicazione su TensorFlow

Esempio di pubblicazione del modello TensorFlow di E2E:

Esegui la demo di pubblicazione di TensorFlow ResNet-50 E2E

Inferenza e pubblicazione del modello PyTorch

TorchDynamo e PyTorch/XLA

TorchServe

Raccogli gli artefatti del modello

Genera un file di archivio del modello

Richieste di inferenza di elaborazione

Esegui la pulizia

Gestione di modelli linguistici di grandi dimensioni (LLM)

Cluster SAX (cella SAX)

Server di amministrazione SAX

Bucket di archiviazione amministratore SAX

Server del modello SAX

Pubblicazione del modello SAX

Esempio di pubblicazione del modello con un solo host GPT-J 6B

Anteprima della pubblicazione del modello multi-host da 175 miliardi

Profilazione

Assistenza e feedback

Termini

Utilizza `jax2tf` per convertire il modello e salvarlo