Questo documento fornisce una breve introduzione al lavoro con Pax su una TPU a singolo host (v2-8, v3-8, v4-8).
Pax è un framework per configurare ed eseguire esperimenti di machine learning su JAX. Pax si concentra sulla semplificazione del machine learning su larga scala condividendo i componenti dell'infrastruttura con i framework ML esistenti e utilizzando la libreria di modellazione Praxis per la modularità.
Obiettivi
- Configurazione delle risorse TPU per l'addestramento
- Installa Pax su una TPU con singolo host
- Addestra un modello SPMD basato su Transformer utilizzando Pax
Prima di iniziare
Esegui i comandi seguenti per configurare gcloud in modo da utilizzare il progetto Cloud TPU e installare i componenti necessari per addestrare un modello che esegue Pax su una TPU con singolo host.
Installa Google Cloud CLI
Google Cloud CLI contiene strumenti e librerie per interagire con i prodotti e i servizi Google Cloud CLI. Se non l'hai ancora installato, installalo ora seguendo le istruzioni riportate in Installazione di Google Cloud CLI.
Configura il comando gcloud
Esegui gcloud auth list per vedere gli account disponibili.
$ gcloud config set account account
$ gcloud config set project project-id
Abilita l'API Cloud TPU
Abilita l'API Cloud TPU utilizzando il seguente comando gcloud in Cloud Shell.
(puoi anche abilitarlo dalla console Google Cloud).
$ gcloud services enable tpu.googleapis.com
Esegui questo comando per creare un'identità di servizio (un account di servizio).
$ gcloud beta services identity create --service tpu.googleapis.com
Crea una VM TPU
Con le VM Cloud TPU, il modello e il codice vengono eseguiti direttamente sulla VM TPU. Accedi direttamente alla VM TPU tramite SSH. Puoi eseguire codice arbitrario, installare pacchetti, visualizzare i log e il codice di debug direttamente sulla VM TPU.
Crea la VM TPU eseguendo questo comando da Cloud Shell o dal terminale del computer in cui è installata Google Cloud CLI.
Imposta zone in base alla disponibilità nel contratto,
fai riferimento alle regioni e zone TPU,
se necessario.
Imposta la variabile accelerator-type su v2-8, v3-8 o v4-8.
Imposta la variabile version su tpu-vm-base
per le versioni TPU v2 e v3 o su tpu-vm-v4-base per le TPU v4.
$ gcloud compute tpus tpu-vm create tpu-name \
--zone zone \
--accelerator-type accelerator-type \
--version version
Connettiti alla VM Google Cloud TPU
Accedi alla VM TPU utilizzando il seguente comando:
$ gcloud compute tpus tpu-vm ssh tpu-name --zone zone
Quando hai eseguito l'accesso alla VM, il prompt della shell cambia da username@projectname a username@vm-name:
Installa Pax sulla VM Google Cloud TPU
Installa Pax, JAX e libtpu sulla tua VM TPU utilizzando i seguenti comandi:
(vm)$ python3 -m pip install -U pip \
python3 -m pip install paxml jax[tpu]
-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Controllo del sistema
Verifica che sia installato correttamente controllando che JAX rilevi i core TPU:
(vm)$ python3 -c "import jax; print(jax.device_count())"
Viene visualizzato il numero di core TPU, che dovrebbe essere 8 se utilizzi una versione v2-8 o v3-8 o 4 se utilizzi una versione v4-8.
Esecuzione di codice Pax su una VM TPU
Ora puoi eseguire qualsiasi codice Pax tu voglia. Gli esempi di lm_cloud sono un ottimo punto di partenza per iniziare a eseguire modelli in Pax. Ad esempio, i seguenti comandi addestrano un modello linguistico SPMD basato su parametri 2B su dati sintetici.
I seguenti comandi mostrano l'output di addestramento per un modello linguistico SPMD. Viene addestrato per 300 passi in circa 20 minuti.
(vm)$ python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=job_log_dir
LmCloudSpmd2BLimitSteps
Nella sezione v4-8, l'output deve includere:
Perdite e tempi di passi
tensore di riepilogo al passaggio=step_# loss = loss
tensore di riepilogo al passaggio=step_# passi al secondo x
Esegui la pulizia
Per evitare che al tuo Account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.
Al termine delle operazioni della VM TPU, segui questi passaggi per la pulizia delle risorse.
Disconnettiti dall'istanza Compute Engine, se non l'hai ancora fatto:
(vm)$ exit
Elimina la tua Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name --zone zone
Passaggi successivi
Per ulteriori informazioni su Cloud TPU, vedi: