Questo documento fornisce una breve introduzione a lavorare con Pax. su una TPU a host singolo (v2-8, v3-8, v4-8).
Pax è un framework per configurare ed eseguire esperimenti di machine learning su JAX. Pax sulla semplificazione del machine learning su larga scala condividendo i componenti dell'infrastruttura con i framework ML esistenti e l'utilizzo Libreria di modelli Prassi per la modularità.
Obiettivi
- Configura le risorse TPU per l'addestramento
- Installa Pax su una TPU con host singolo
- Addestra un modello SPMD basato su trasformatore utilizzando Pax
Prima di iniziare
Esegui questi comandi per configurare gcloud da utilizzare
il tuo progetto Cloud TPU e installare i componenti necessari per
un modello che esegue Pax su una TPU con host singolo.
Installa Google Cloud CLI
Google Cloud CLI contiene strumenti e librerie per interagire con i prodotti e i servizi Google Cloud CLI. Se non hai installato in precedenza, installalo ora seguendo 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 abilitarla 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, modello e codice vengono eseguiti direttamente sulla VM TPU. Accedi tramite SSH direttamente alla VM TPU. Puoi eseguire in modo arbitrario installare pacchetti, visualizzare i log e eseguire il debug del codice direttamente sulla VM TPU.
Crea la tua VM TPU eseguendo questo comando da un Cloud Shell o il terminale del computer su cui Google Cloud CLI sia installato.
Imposta zone in base alla disponibilità nel tuo contratto.
riferimento a 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 tpu-vm-v4-base per 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
Dopo aver eseguito l'accesso alla VM, il prompt della shell cambia da
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 sia tutto installato correttamente controllando che JAX vede 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 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. lm_cloud esempi sono un ottimo punto di partenza per eseguire modelli in Pax. Ad esempio, i seguenti comandi addestrano Modello linguistico SPMD basato su trasformatore 2B Parameter su dati sintetici.
I seguenti comandi mostrano l'output di addestramento per una lingua SPMD model. Si addestra 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
Nella sezione v4-8, l'output dovrebbe includere:
Perdite e tempi di passi
tensore di riepilogo al passaggio=step_# loss = loss
tensore di riepilogo al passo=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 con la VM TPU, segui questi passaggi per pulire le risorse.
Disconnettiti dall'istanza Compute Engine, se non lo hai già fatto:
(vm)$ exit
Elimina il tuo Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name --zone zone
Passaggi successivi
Per ulteriori informazioni su Cloud TPU, vedi: