Esegui i workload 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 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.

Addestra un modello PyTorch in un container Docker

PyTorch su Cloud TPU descrive le librerie PyTorch per TPU e rimanda alla documentazione per ottenere i pacchetti Python e le immagini Docker. Puoi utilizzare TorchTPU direttamente (consulta Inizia a utilizzare GKE per l'immagine container della release pubblica) o utilizzare la libreria di riferimento TorchTitan per creare ed eseguire workload di addestramento containerizzati.

Addestra un modello JAX in un container Docker

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

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: l'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 ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta 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 tramite 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. 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-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, esegui la pulizia delle 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 procedura seguente 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 ad Artifact Registry. Per ulteriori informazioni sull'utilizzo di Artifact Registry, consulta Utilizza 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, esegui la pulizia delle 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