Esegui i carichi di lavoro TPU in un container Docker

Dispositivo TPU

1. Crea una VM Cloud TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

2. Connettiti alla VM TPU tramite SSH

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a

3. Assicurati che all' Google Cloud utente sia stato assegnato il ruolo di lettore del Registry degli elementi. Per ulteriori informazioni, consulta Concedere i ruoli di Artifact Registry.

4. Avvia un contenitore nella VM TPU utilizzando l'immagine PyTorch/XLA notturna.

   sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm bash

5. Configura il runtime TPU

   Esistono due opzioni di runtime PyTorch/XLA: PJRT e XRT. Ti consigliamo di utilizzare PJRT, a meno che tu non abbia un motivo per utilizzare XRT. Per scoprire di più sulle diverse configurazioni di runtime, consulta i motivi per cui utilizzare XRT. Per approfondire le diverse configurazioni del runtime, consulta la documentazione del runtime PJRT.

   PJRT

   export PJRT_DEVICE=TPU

   XRT

   export XRT_TPU_CONFIG="localservice;0;localhost:51011"

6. Clona il repository PyTorch XLA

   git clone --recursive https://github.com/pytorch/xla.git

7. Addestra ResNet50

   python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

Al termine dello script di addestramento, assicurati di ripulire le risorse.

1. Digita exit per uscire dal container Docker
2. Digita exit per uscire dalla VM TPU
3. Elimina la VM TPU

    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Sezione TPU

Quando esegui il codice PyTorch su una sezione TPU, devi eseguirlo su tutti i worker TPU contemporaneamente. Un modo per farlo è utilizzare il comando gcloud compute tpus tpu-vm ssh con i flag --worker=all e --command. La procedura riportata di seguito mostra come creare un'immagine Docker per semplificare la configurazione di ogni worker TPU.

1. Crea una VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=us-central2-b \
   --accelerator-type=v4-32 \
   --version=tpu-ubuntu2204-base

2. Aggiungi l'utente corrente al gruppo docker

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=us-central2-b \
   --worker=all \
   --command="sudo usermod -a -G docker $USER"

3. Clona il repository PyTorch XLA

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="git clone --recursive https://github.com/pytorch/xla.git"

4. Esegui lo script di addestramento in un contenitore su tutti i worker TPU.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="docker run --rm --privileged --net=host  -v ~/xla:/xla -e PJRT_DEVICE=TPU us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 python /xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"

   Flag dei comandi Docker:

   • --rm rimuovi il contenitore al termine del processo.
   • --privileged espone il dispositivo TPU al contenitore.
   • --net=host associa tutte le porte del contenitore alla VM TPU per consentire la comunicazione tra gli host del pod.
   • -e imposta le variabili di ambiente.

Al termine dello script di addestramento, assicurati di ripulire le risorse.

Elimina la VM TPU utilizzando il seguente comando:

$ gcloud compute tpus tpu-vm delete your-tpu-name \
  --zone=us-central2-b

Dispositivo TPU

1. Crea la VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

2. Connettiti alla VM TPU tramite SSH

   gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a

3. Avvia il daemon Docker nella VM TPU

   sudo systemctl start docker

4. Avvia il container Docker

   sudo docker run --net=host -ti --rm --name your-container-name \
    --privileged us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 \
    python:3.10 bash

5. Installa JAX

   pip install jax[tpu]

6. Installa FLAX

   pip install --upgrade clu
   git clone https://github.com/google/flax.git
   pip install --user -e flax

7. Installa i pacchetti tensorflow e tensorflow-dataset

   pip install tensorflow
   pip install tensorflow-datasets

8. Esegui lo script di addestramento FLAX MNIST

   cd flax/examples/mnist
   python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5

Al termine dello script di addestramento, assicurati di ripulire le risorse.

1. Digita exit per uscire dal container Docker
2. Digita exit per uscire dalla VM TPU

3. Elimina la VM TPU

   $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Sezione TPU

Quando esegui il codice JAX su una sezione TPU, devi eseguirlo su tutti i worker TPU contemporaneamente. Un modo per farlo è utilizzare il comando gcloud compute tpus tpu-vm ssh con i flag --worker=all e --command. La procedura riportata di seguito mostra come creare un'immagine Docker per semplificare la configurazione di ogni worker TPU.

1. Crea un file denominato Dockerfile nella directory corrente e incolla il seguente testo

   FROM python:3.10
   RUN pip install jax[tpu]
   RUN pip install --upgrade clu
   RUN git clone https://github.com/google/flax.git
   RUN pip install --user -e flax
   WORKDIR ./flax/examples/mnist

2. Preparare un Artifact Registry

   gcloud artifacts repositories create your-repo \
   --repository-format=docker \
   --location=europe-west4 --description="Docker repository" \
   --project=your-project
   
   gcloud artifacts repositories list \
   --project=your-project
   
   gcloud auth configure-docker europe-west4-docker.pkg.dev

3. Crea l'immagine Docker

   docker build -t your-image-name .

4. Aggiungi un tag all'immagine Docker prima di eseguirne il push in Artifact Registry. Per ulteriori informazioni sull'utilizzo di Artifact Registry, consulta Utilizzare le immagini container.

   docker tag your-image-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

5. Esegui il push dell'immagine Docker in Artifact Registry

   docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

6. Crea una VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

7. Estrai l'immagine Docker da Artifact Registry su tutti i worker TPU.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="sudo usermod -a -G docker ${USER}"

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker pull europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag"

8. Esegui il container su tutti i worker TPU.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag bash"

9. Esegui lo script di addestramento su tutti i worker TPU:

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker exec --privileged your-container-name python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5"

Al termine dello script di addestramento, assicurati di ripulire le risorse.

1. Arresta il contenitore su tutti i worker:

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker kill your-container-name"

2. Elimina la VM TPU utilizzando il seguente comando:

   $ gcloud compute tpus tpu-vm delete your-tpu-name \
   --zone=europe-west4-a