Esegui applicazioni Cloud TPU in un container Docker
I container Docker semplificano la configurazione delle applicazioni combinando il codice e tutte le dipendenze necessarie in un unico pacchetto distribuibile. Puoi eseguire i container Docker all'interno delle VM TPU per semplificare la configurazione e la condivisione delle applicazioni Cloud TPU. Questo documento descrive come configurare un container Docker per ogni framework ML supportato da Cloud TPU.
Addestra un modello TensorFlow in un container Docker
dispositivo TPU
Crea un file denominato
Dockerfile
nella directory attuale e incolla il seguente testoFROM python:3.8 RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so RUN git clone https://github.com/tensorflow/models.git WORKDIR ./models RUN pip install -r official/requirements.txt ENV PYTHONPATH=/models
Crea bucket Cloud Storage
gsutil mb -c standard -l europe-west4 gs://your-bucket-name
Crea una VM TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-tf-2.14.0
Copia il Dockerfile nella VM TPU
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
SSH nella VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Crea l'immagine Docker
sudo docker build -t your-image-name .
Avvia il container Docker
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
Imposta le variabili di ambiente
export STORAGE_BUCKET=gs://your-bucket-name export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x
Addestra ResNet
python3 official/vision/train.py \ --tpu=local \ --experiment=resnet_imagenet \ --mode=train_and_eval \ --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \ --model_dir=${MODEL_DIR} \ --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,trainer.train_steps=100"
Una volta completato lo script di addestramento, assicurati di eseguire la pulizia delle risorse.
- Digita
exit
per uscire dal container Docker - Digita
exit
per uscire dalla VM TPU - Elimina la VM TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
pod di TPU
Crea un file denominato
Dockerfile
nella directory attuale e incolla il seguente testoFROM python:3.8 RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so RUN git clone https://github.com/tensorflow/models.git WORKDIR ./models RUN pip install -r official/requirements.txt ENV PYTHONPATH=/models
Crea una VM TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v3-32 \ --version=tpu-vm-tf-2.14.0-pod
Copia il Dockerfile nella VM TPU
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
SSH nella VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Crea l'immagine Docker
sudo docker build -t your-image-name .
Avvia un container Docker
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
Addestra ResNet
python3 official/vision/train.py \ --tpu=local \ --experiment=resnet_imagenet \ --mode=train_and_eval \ --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \ --model_dir=${MODEL_DIR} \ --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
Una volta completato lo script di addestramento, assicurati di eseguire la pulizia delle risorse.
- Digita
exit
per uscire dal container Docker - Digita
exit
per uscire dalla VM TPU - Elimina la VM TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Addestra un modello PyTorch in un container Docker
dispositivo TPU
Crea VM Cloud TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-ubuntu2204-base
SSH nella VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Avvia un container nella VM TPU utilizzando l'immagine PyTorch/XLA di notte.
sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm bash
Configura runtime TPU
Esistono due opzioni di runtime PyTorch/XLA: PJRT e XRT. Ti consigliamo di usare PJRT a meno che tu non abbia un motivo per usare XRT. Per scoprire di più sulle diverse configurazioni di runtime, vedi "Hai un motivo per utilizzare XRT". Per scoprire di più sulle diverse configurazioni di runtime, consulta la documentazione sul runtime PJRT.
PJRT
export PJRT_DEVICE=TPU
X
export XRT_TPU_CONFIG="localservice;0;localhost:51011"
Clona il repository PyTorch XLA
git clone --recursive https://github.com/pytorch/xla.git
Addestra ResNet50
python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
Una volta completato lo script di addestramento, assicurati di eseguire la pulizia delle risorse.
- Digita
exit
per uscire dal container Docker - Digita
exit
per uscire dalla VM TPU - Elimina la VM TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
pod di TPU
Quando esegui il codice PyTorch su un pod di TPU, devi eseguirlo contemporaneamente su tutti i worker TPU. Un modo per farlo è utilizzare il comando gcloud compute tpus tpu-vm ssh
con i flag --worker=all
e --command
. La procedura seguente mostra come creare un'immagine Docker per semplificare la configurazione di ogni worker TPU.
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
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"
Esegui lo script di addestramento in un container 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 -e PJRT_DEVICE=TPU gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm python /pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"
Flag dei comandi 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
ha impostato le variabili di ambiente.
Una volta completato lo script di addestramento, assicurati di eseguire la pulizia delle risorse.
Elimina la VM TPU utilizzando il seguente comando:
$ gcloud compute tpus tpu-vm delete your-tpu-name \
--zone=us-central2-b
Addestra un modello JAX in un container Docker
Dispositivo TPU
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
SSH nella VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Avvia il daemon Docker nella VM TPU
sudo systemctl start docker
Avvia container Docker
sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
Installa JAX
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Installa FLAX
pip install --upgrade clu git clone https://github.com/google/flax.git pip install --user -e flax
Esegui lo script di addestramento di 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
Una volta completato lo script di addestramento, assicurati di eseguire la pulizia delle risorse.
- Digita
exit
per uscire dal container Docker - Digita
exit
per uscire dalla VM TPU Elimina la VM TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
pod di TPU
Quando esegui il codice JAX su un pod di 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.
Crea un file denominato
Dockerfile
nella directory attuale e incolla il seguente testoFROM python:3.8 RUN pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html RUN pip install --upgrade clu RUN git clone https://github.com/google/flax.git RUN pip install --user -e flax WORKDIR ./flax/examples/mnist
Crea l'immagine Docker
docker build -t your-image-name .
Aggiungi un tag all'immagine Docker prima di eseguirne il push ad Artifact Registry. Per ulteriori informazioni sull'utilizzo di Artifact Registry, consulta Utilizzare le immagini container.
docker tag your-image-name europe-west-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
Esegui il push dell'immagine Docker ad Artifact Registry
docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
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
Esegui il pull dell'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"
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:your-tag bash"
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"
Una volta completato lo script di addestramento, assicurati di eseguire la pulizia delle risorse.
Arresta il container 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"
Elimina la VM TPU utilizzando il seguente comando:
$ gcloud compute tpus tpu-vm delete your-tpu-name \ --zone=europe-west4-a
Passaggi successivi
- Tutorial su Cloud TPU
- Gestisci le TPU
- Architettura del sistema Cloud TPU
- Esegui il codice TensorFlow sulle sezioni dei pod di TPU
- Esegui codice JAX su sezioni di pod di TPU