Questo documento fornisce una breve introduzione all'utilizzo di Pax su una TPU a un solo host (v2-8, v3-8, v4-8).
Pax è un framework per configurare e 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 modelli Praxis per la modularità.
Obiettivi
- Configura le risorse TPU per l'addestramento
- Installa Pax su una TPU a singolo host
- Addestra un modello SPMD basato su Transformer utilizzando Pax
Prima di iniziare
Esegui i seguenti comandi per configurare gcloud in modo che utilizzi il tuo progetto Cloud TPU e installa i componenti necessari per addestrare un modello che esegue Pax su una TPU a un solo 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, segui le istruzioni riportate in Installazione di Google Cloud CLI.
Configura il comando gcloud
(esegui gcloud auth list per visualizzare 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 attivarla dalla
console Google Cloud.
$ gcloud services enable tpu.googleapis.com
Esegui il seguente 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 tramite SSH direttamente alla VM TPU. Puoi eseguire codice arbitrario, installare pacchetti, visualizzare i log e eseguire il debug del codice direttamente sulla VM TPU.
Crea la VM TPU eseguendo il seguente comando da Cloud Shell o dal terminale del computer in cui è installato Google Cloud CLI.
Imposta zone in base alla disponibilità nel contratto. Se necessario, consulta Regioni e zone TPU.
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 tramite SSH alla VM TPU utilizzando il seguente comando:
$ gcloud compute tpus tpu-vm ssh tpu-name --zone zone
Una volta 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 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 tutto sia installato correttamente controllando che JAX veda 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 v2-8 o v3-8 oppure 4 se utilizzi v4-8.
Esecuzione del codice Pax su una VM TPU
Ora puoi eseguire qualsiasi codice Pax che preferisci. Gli esempi 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 trasformatori con 2 miliardi di parametri su dati sintetici.
I seguenti comandi mostrano l'output dell'addestramento per un modello linguistico SPMD. Si allena per 300 passi in circa 20 minuti.
(vm)$ python3 .local/lib/python3.10/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=job_log_dir
Nel caso del segmento v4-8, l'output deve includere:
Perdite e tempi di passaggio
tensore di riepilogo al passaggio=step_# loss = loss
tensore di riepilogo al passaggio=step_# Passaggi 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 dell'utilizzo della VM TPU, segui questi passaggi per ripulire le risorse.
Se non l'hai ancora fatto, scollega l'istanza Compute Engine:
(vm)$ exit
Elimina la Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name --zone zone
Passaggi successivi
Per ulteriori informazioni su Cloud TPU, consulta: