Esegui carichi di lavoro TPU in un container Docker

Attenzione:questa pagina si applica all'API Cloud TPU, che non è più in fase di sviluppo attivo. L'API Cloud TPU riceverà solo correzioni di bug e aggiornamenti della sicurezza. Per le funzionalità più recenti e il supporto delle versioni TPU più recenti, ti consigliamo di utilizzare Compute Engine o Google Kubernetes Engine per gestire le risorse TPU. Per saperne di più, consulta Risorse TPU in Compute Engine e Informazioni sulle TPU in GKE. L'API Cloud TPU non è supportata su TPU7x e versioni TPU successive.

I container Docker semplificano la configurazione delle applicazioni combinando il codice e tutte le dipendenze necessarie in un unico pacchetto distribuibile. Puoi eseguire container Docker all'interno delle VM TPU per semplificare la configurazione e la condivisione delle tue applicazioni Cloud TPU. Questo documento descrive come configurare un container Docker per ogni framework ML supportato da Cloud TPU.

Importante: puoi eseguire un massimo di 256 pod Docker su una VM TPU.

Addestra un modello PyTorch in un container Docker

Prima di eseguire i comandi seguenti, crea le variabili di ambiente seguenti:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Descrizioni delle variabili di ambiente

• PROJECT_ID: il tuo ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
• TPU_NAME: il nome della TPU.
• ZONE: La zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
• ACCELERATOR_TYPE: Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per maggiori informazioni sui tipi di acceleratore supportati per ogni versione di TPU, consulta la sezione Versioni di TPU.
• RUNTIME_VERSION: la versione software di Cloud TPU.

Dispositivo TPU

1. Crea VM Cloud TPU

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. Connettiti alla VM TPU utilizzando SSH

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
   

3. Assicurati che all'utente Google Cloud sia stato concesso il ruolo Lettore Artifact Registry. Per saperne di più, consulta Concessione dei ruoli di Artifact Registry.

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

   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 \
   bash
   

   Nota: dopo aver eseguito questo comando, il prompt dei comandi cambierà per indicare che il terminale è connesso al container in esecuzione.

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 saperne di più sulle diverse configurazioni di runtime, consulta la documentazione sul 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, libera spazio dalle 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 $TPU_NAME --zone=$ZONE
   

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 seguente procedura 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 $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. Aggiungi l'utente corrente al gruppo Docker

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --zone=$ZONE \
   --worker=all \
   --command='sudo usermod -a -G docker $USER'
   

3. Clona il repository PyTorch XLA

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="git clone --recursive https://github.com/pytorch/xla.git"
   

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --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 del comando Docker:

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

Al termine dello script di addestramento, libera spazio dalle risorse.

Elimina la VM TPU utilizzando il seguente comando:

gcloud compute tpus tpu-vm delete $TPU_NAME \
--zone=$ZONE

Addestra un modello JAX in un container Docker

Prima di eseguire i comandi seguenti, crea le variabili di ambiente seguenti:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Descrizioni delle variabili di ambiente

• PROJECT_ID: il tuo ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
• TPU_NAME: il nome della TPU.
• ZONE: La zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
• ACCELERATOR_TYPE: Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per maggiori informazioni sui tipi di acceleratore supportati per ogni versione di TPU, consulta la sezione Versioni di TPU.
• RUNTIME_VERSION: la versione software di Cloud TPU.

Dispositivo TPU

1. Crea la VM TPU

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. Connettiti alla VM TPU utilizzando SSH

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
   

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 \
   bash
   

5. Installare JAX

   pip install jax[tpu]
   

6. Installare 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-datasets

   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, libera spazio dalle 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 $TPU_NAME --zone=$ZONE
   

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 seguente procedura mostra come creare un'immagine Docker per semplificare la configurazione di ogni worker TPU.

1. Crea un file denominato Dockerfile nella directory attuale 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
   RUN pip install tensorflow
   RUN pip install tensorflow-datasets
   WORKDIR ./flax/examples/mnist
   

2. Prepara un Artifact Registry

   gcloud artifacts repositories create your-repo \
   --repository-format=docker \
   --location=europe-west4 --description="Docker repository" \
   --project=$PROJECT_ID
   
   gcloud artifacts repositories list \
   --project=$PROJECT_ID
   
   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, vedi Utilizzare le immagini container.

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

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

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

6. Crea una VM TPU

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command='sudo usermod -a -G docker ${USER}'
   

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
   

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

8. Esegui il container su tutti i worker TPU

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/$PROJECT_ID/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 $TPU_NAME --worker=all \
   --zone=$ZONE \
   --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, libera spazio dalle risorse.

1. Arresta il container su tutti i worker

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="docker kill your-container-name"
   

2. Elimina la VM TPU

   gcloud compute tpus tpu-vm delete $TPU_NAME \
   --zone=$ZONE
   

Passaggi successivi

• Tutorial su Cloud TPU
• Gestire le TPU
• Architettura Cloud TPU
• Esecuzione del codice JAX nelle sezioni TPU