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 tue applicazioni Cloud TPU. Questo documento descrive come configurare un contenitore Docker per ogni framework ML supportato da Cloud TPU.

Addestra un modello TensorFlow in un container Docker

Dispositivo TPU

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

    FROM 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
  2. crea un bucket Cloud Storage

    gcloud storage buckets create gs://your-bucket-name --location=europe-west4
  3. 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.18.0-pjrt
  4. Copia il Dockerfile nella VM TPU

    gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
  5. Accedi tramite SSH alla VM TPU

    gcloud compute tpus tpu-vm ssh your-tpu-name \
    --zone=europe-west4-a
  6. Crea l'immagine Docker

    sudo docker build -t your-image-name .
  7. Avvia il container Docker

    sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
  8. 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
  9. Addestramento di 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"

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

pod di TPU

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

    FROM 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
  2. 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.18.0-pod-pjrt
  3. Copia il Dockerfile nella VM TPU

    gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
  4. Accedi tramite SSH alla VM TPU

    gcloud compute tpus tpu-vm ssh your-tpu-name \
    --zone=europe-west4-a
  5. Crea l'immagine Docker

    sudo docker build -t your-image-name .
  6. Avvia un container Docker

    sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
  7. Addestramento di 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"

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

Addestra un modello PyTorch 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. Accedi tramite SSH alla VM TPU

    gcloud compute tpus tpu-vm ssh your-tpu-name \
    --zone=europe-west4-a
  3. 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
  4. 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"
  5. Clona il repository PyTorch XLA

    git clone --recursive https://github.com/pytorch/xla.git
  6. 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

pod di TPU

Quando esegui il codice PyTorch 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.

  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. 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  -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 rimuovi il contenitore al termine del processo.
    • --privileged espone il dispositivo TPU al contenitore.
    • --net=host lega 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

Addestrare un modello JAX in un container Docker

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. Accedi tramite SSH alla VM TPU

    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 -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
  5. Installa JAX

    pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
  6. Installa FLAX

    pip install --upgrade clu
    git clone https://github.com/google/flax.git
    pip install --user -e flax
  7. 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

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.

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

    FROM 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
  2. Crea l'immagine Docker

    docker build -t your-image-name .
  3. 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-west-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
  4. 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
  5. 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
  6. 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"
  7. 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"
  8. 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. Spegni 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

Addestrare un modello JAX in un container Docker utilizzando JAX Stable Stack

Puoi creare le immagini Docker MaxText e MaxDiffusion utilizzando l'immagine di base JAX Stable Stack.

JAX Stable Stack fornisce un ambiente coerente per MaxText e MaxDiffusion raggruppando JAX con pacchetti di base come orbax, flax e optax e libtpu.so. Queste librerie vengono testate per garantire la compatibilità e fornire una base stabile per creare ed eseguire MaxText e MaxDiffusion. In questo modo, vengono eliminati i potenziali conflitti dovuti a versioni del pacchetto incompatibili.

JAX Stable Stack include un libtpu.so completamente rilasciato e qualificato, la libreria di base che gestisce la compilazione, l'esecuzione e la configurazione della rete ICI dei programmi TPU. La release libtpu sostituisce la build notturna precedentemente utilizzata da JAX e garantisce la funzionalità coerente dei calcoli XLA su TPU con test di qualificazione a livello di PJRT negli IR HLO/StableHLO.

Per creare l'immagine Docker di MaxText e MaxDiffusion con JAX Stable Stack, quando esegui lo script docker_build_dependency_image.sh, imposta la variabile MODE su stable_stack e la variabile BASEIMAGE su l'immagine di base che vuoi utilizzare.

L'esempio seguente specifica us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.35-rev1 come immagine di base:

   bash docker_build_dependency_image.sh MODE=stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.35-rev1
   

Per un elenco delle immagini di base JAX Stable Stack disponibili, consulta Immagini JAX Stable Stack in Artifact Registry.

Passaggi successivi