Esegui un calcolo su una VM Cloud TPU utilizzando JAX
Questo documento fornisce una breve introduzione all'utilizzo di JAX e Cloud TPU.
Prima di iniziare
Prima di eseguire i comandi in questo documento, devi creare un account Google Cloud, installare Google Cloud CLI e configurare il comando gcloud. Per maggiori informazioni, consulta Configurare l'ambiente Cloud TPU.
Crea una VM Cloud TPU utilizzando gcloud
Definisci alcune variabili di ambiente per semplificare l'utilizzo dei comandi.
export PROJECT_ID=your-project export ACCELERATOR_TYPE=v5p-8 export ZONE=us-east5-a export RUNTIME_VERSION=v2-alpha-tpuv5 export TPU_NAME=your-tpu-name
Descrizioni delle variabili di ambiente
PROJECT_ID- Il tuo Google Cloud ID progetto.
ACCELERATOR_TYPE- Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni di TPU.
ZONE- La zona in cui prevedi di creare la tua Cloud TPU.
RUNTIME_VERSION- La versione del runtime Cloud TPU. Per ulteriori informazioni, consulta Immagini VM TPU .
TPU_NAME- Il nome assegnato dall'utente alla tua Cloud TPU.
Crea la VM TPU eseguendo il seguente comando da Cloud Shell o dal terminale del computer su cui è installato Google Cloud CLI.
$ gcloud compute tpus tpu-vm create $TPU_NAME \ --project=$PROJECT_ID \ --zone=$ZONE \ --accelerator-type=$ACCELERATOR_TYPE \ --version=$RUNTIME_VERSION
Connettiti alla VM Cloud TPU
Connettiti alla VM TPU tramite SSH utilizzando il seguente comando:
$ gcloud compute tpus tpu-vm ssh $TPU_NAME \ --project=$PROJECT_ID \ --zone=$ZONE
Installa JAX sulla tua VM Cloud TPU
(vm)$ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Controllo del sistema
Verifica che JAX possa accedere alla TPU ed eseguire operazioni di base:
Avvia l'interprete Python 3:
(vm)$ python3>>> import jax
Mostra il numero di core TPU disponibili:
>>> jax.device_count()
Viene visualizzato il numero di core TPU. Il numero di core visualizzato dipende dalla versione della TPU in uso. Per saperne di più, consulta Versioni TPU.
Esegui un calcolo:
>>> jax.numpy.add(1, 1)
Viene visualizzato il risultato dell'addizione di NumPy:
Output del comando:
Array(2, dtype=int32, weak_type=true)Esci dall'interprete Python:
>>> exit()
Esecuzione del codice JAX su una VM TPU
Ora puoi eseguire qualsiasi codice JAX. Gli esempi di flax sono un ottimo punto di partenza per eseguire modelli di ML standard in JAX. Ad esempio, per addestrare una rete convoluzionale MNIST di base:
Installa le dipendenze degli esempi di Flax
(vm)$ pip install --upgrade clu (vm)$ pip install tensorflow (vm)$ pip install tensorflow_datasets
Installa FLAX
(vm)$ git clone https://github.com/google/flax.git (vm)$ pip install --user flax
Esegui lo script di addestramento FLAX MNIST
(vm)$ cd flax/examples/mnist (vm)$ python3 main.py --workdir=/tmp/mnist \ --config=configs/default.py \ --config.learning_rate=0.05 \ --config.num_epochs=5
Lo script scarica il set di dati e avvia l'addestramento. L'output dello script dovrebbe essere simile al seguente:
0214 18:00:50.660087 140369022753856 train.py:146] epoch: 1, train_loss: 0.2421, train_accuracy: 92.97, test_loss: 0.0615, test_accuracy: 97.88 I0214 18:00:52.015867 140369022753856 train.py:146] epoch: 2, train_loss: 0.0594, train_accuracy: 98.16, test_loss: 0.0412, test_accuracy: 98.72 I0214 18:00:53.377511 140369022753856 train.py:146] epoch: 3, train_loss: 0.0418, train_accuracy: 98.72, test_loss: 0.0296, test_accuracy: 99.04 I0214 18:00:54.727168 140369022753856 train.py:146] epoch: 4, train_loss: 0.0305, train_accuracy: 99.06, test_loss: 0.0257, test_accuracy: 99.15 I0214 18:00:56.082807 140369022753856 train.py:146] epoch: 5, train_loss: 0.0252, train_accuracy: 99.20, test_loss: 0.0263, test_accuracy: 99.18
Esegui la pulizia
Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questa pagina, segui questi passaggi.
Al termine dell'utilizzo della VM TPU, segui questi passaggi per ripulire le risorse.
Se non l'hai ancora fatto, disconnettiti dall'istanza Compute Engine:
(vm)$ exit
Elimina la Cloud TPU.
$ gcloud compute tpus tpu-vm delete $TPU_NAME \ --project=$PROJECT_ID \ --zone=$ZONE
Verifica che le risorse siano state eliminate eseguendo il seguente comando. Assicurati che la TPU non sia più inclusa nell'elenco. L'eliminazione può richiedere qualche minuto.
$ gcloud compute tpus tpu-vm list \ --zone=$ZONE
Note sul rendimento
Di seguito sono riportati alcuni dettagli importanti particolarmente pertinenti per l'utilizzo delle TPU in JAX.
Spaziatura interna
Una delle cause più comuni di prestazioni lente sulle TPU è l'introduzione di spaziatura involontaria:
- Gli array in Cloud TPU sono suddivisi in riquadri. Ciò comporta l'aggiunta di spazi iniziali a una delle dimensioni in modo che sia un multiplo di 8 e a un'altra dimensione in modo che sia un multiplo di 128.
- L'unità di moltiplicazione delle matrici ha il rendimento migliore con coppie di matrici di grandi dimensioni che riducono al minimo la necessità di spaziatura interna.
Tipo di dati bfloat16
Per impostazione predefinita, la moltiplicazione di matrici in JAX su TPU utilizza bfloat16 con accumulo float32. Questo può essere controllato con l'argomento precision sulle chiamate alle funzioni jax.numpy pertinenti (matmul, dot, einsum e così via). In particolare:
precision=jax.lax.Precision.DEFAULT: utilizza la precisione bfloat16 mista (più veloce)precision=jax.lax.Precision.HIGH: utilizza più passaggi MXU per ottenere una maggiore precisioneprecision=jax.lax.Precision.HIGHEST: utilizza ancora più passaggi MXU per ottenere una precisione completa di float32
JAX aggiunge anche il tipo di dati bfloat16, che puoi utilizzare per eseguire il casting esplicito degli array in bfloat16, ad esempio,jax.numpy.array(x, dtype=jax.numpy.bfloat16).
Passaggi successivi
Per ulteriori informazioni su Cloud TPU, consulta: