Questo documento fornisce una breve introduzione al lavoro 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 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 con host singolo
- Addestra un modello SPMD basato su trasformatore 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 l'addestramento di un modello che esegue Pax su una TPU con 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, installalo ora utilizzando le istruzioni in Installazione di Google Cloud CLI.
Configura il comando gcloud
Esegui gcloud auth list per vedere gli account attualmente 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 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. Se necessario, fai riferimento alle 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 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:
Installare 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 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 2-8 o v3-8 oppure 4 se utilizzi una versione v4-8.
Esecuzione del codice Pax su una VM TPU
Ora puoi eseguire qualsiasi codice Pax che desideri. 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 di linguaggio SPMD basato su Transformer di 2 miliardi di parametri su dati sintetici.
I seguenti comandi mostrano l'output di addestramento per un modello linguistico SPMD. 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 deve includere:
Sconfitte e tempi di passi
tensore di riepilogo al passaggio=step_# loss = loss
tensore di riepilogo al passaggio=step_# passi/sec 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 della gestione della VM TPU, segui questi passaggi per eseguire la pulizia delle risorse.
Disconnettiti dall'istanza Compute Engine, se non l'hai ancora fatto:
(vm)$ exit
Elimina la Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name --zone zone
Passaggi successivi
Per ulteriori informazioni su Cloud TPU, vedi: