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 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 Transformer 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 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, 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 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 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 v2-8 o v3-8 oppure 4 se utilizzi v4-8.
Esecuzione di codice Pax su una VM TPU
Ora puoi eseguire qualsiasi codice Pax. Il componente lm_cloud esempi sono un ottimo punto di partenza per 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
Nella sezione v4-8, l'output dovrebbe 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 il tuo Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name --zone zone
Passaggi successivi
Per ulteriori informazioni su Cloud TPU, vedi: