Executar cargas de trabalho de TPU em um contêiner do Docker

Os contêineres do Docker facilitam a configuração de aplicativos combinando seu código e todas as dependências necessárias em um pacote distribuível. É possível executar contêineres do Docker em VMs de TPU para simplificar a configuração e o compartilhamento de aplicativos do Cloud TPU. Este documento descreve como configurar um contêiner do Docker para cada framework de ML compatível com o Cloud TPU.

Treinar um modelo do PyTorch em um contêiner do Docker

Dispositivo TPU

  1. Criar uma VM do Cloud TPU

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v2-8 \
    --version=tpu-ubuntu2204-base
  2. Conecte-se à VM da TPU usando SSH

    gcloud compute tpus tpu-vm ssh your-tpu-name \
    --zone=europe-west4-a
  3. Verifique se o usuário Google Cloud recebeu a função de leitor do Artifact Registry. Para mais informações, consulte Conceder papéis do Artifact Registry.

  4. Iniciar um contêiner na VM de TPU usando a imagem noturna do 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
  5. Configurar o ambiente de execução da TPU

    Há duas opções de tempo de execução do PyTorch/XLA: PJRT e XRT. Recomendamos usar a PJRT, a menos que você tenha um motivo para usar a XRT. Para saber mais sobre as diferentes configurações de execução, consulte a documentação de execução do PJRT.

    PJRT

    export PJRT_DEVICE=TPU

    XRT

    export XRT_TPU_CONFIG="localservice;0;localhost:51011"
  6. Clone o repositório PyTorch XLA

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

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

Quando o script de treinamento for concluído, limpe os recursos.

  1. Digite exit para sair do contêiner do Docker.
  2. Digite exit para sair da VM da TPU.
  3. Excluir a VM da TPU

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

Fração de TPU

Ao executar o código PyTorch em uma fração de TPU, é necessário executar o código em todos os workers de TPU ao mesmo tempo. Uma maneira de fazer isso é usar o comando gcloud compute tpus tpu-vm ssh com as flags --worker=all e --command. O procedimento a seguir mostra como criar uma imagem do Docker para facilitar a configuração de cada worker de TPU.

  1. Criar uma VM de TPU

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central2-b \
    --accelerator-type=v4-32 \
    --version=tpu-ubuntu2204-base
  2. Adicionar o usuário atual ao grupo do Docker

    gcloud compute tpus tpu-vm ssh your-tpu-name \
    --zone=us-central2-b \
    --worker=all \
    --command='sudo usermod -a -G docker $USER'
  3. Clone o repositório 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. Executar o script de treinamento em um contêiner em todos os workers de 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"

    Flags de comando do Docker:

    • --rm remove o contêiner após o término do processo.
    • --privileged expõe o dispositivo TPU ao contêiner.
    • --net=host vincula todas as portas do contêiner à VM da TPU para permitir a comunicação entre os hosts no pod.
    • -e define variáveis de ambiente.

Quando o script de treinamento for concluído, limpe os recursos.

Exclua a VM da TPU usando o comando a seguir:

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

Treinar um modelo JAX em um contêiner do Docker

Dispositivo TPU

  1. Criar a VM de TPU

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v2-8 \
    --version=tpu-ubuntu2204-base
  2. Conecte-se à VM da TPU usando SSH

    gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a
  3. Iniciar o daemon do Docker na VM de TPU

    sudo systemctl start docker
  4. Iniciar contêiner do 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. Instalar o JAX

    pip install jax[tpu]
  6. Instalar o FLAX

    pip install --upgrade clu
    git clone https://github.com/google/flax.git
    pip install --user -e flax
  7. Instalar pacotes tensorflow e tensorflow-dataset

    pip install tensorflow
    pip install tensorflow-datasets
  8. Execute o script de treinamento MNIST do FLAX

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

Quando o script de treinamento for concluído, limpe os recursos.

  1. Digite exit para sair do contêiner do Docker.
  2. Digite exit para sair da VM da TPU.
  3. Excluir a VM da TPU

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

Fração de TPU

Ao executar o código JAX em uma fração de TPU, é preciso executar o código JAX em todos os workers de TPU ao mesmo tempo. Uma maneira de fazer isso é usar o comando gcloud compute tpus tpu-vm ssh com as flags --worker=all e --command. O procedimento a seguir mostra como criar uma imagem do Docker para facilitar a configuração de cada worker da TPU.

  1. Crie um arquivo chamado Dockerfile no diretório atual e cole o texto a seguir

    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. Preparar um 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. Compilar a imagem Docker

    docker build -t your-image-name .
  4. Adicione uma tag à imagem do Docker antes de enviá-la ao Artifact Registry. Para mais informações sobre como trabalhar com o Artifact Registry, consulte Trabalhar com imagens de contêiner.

    docker tag your-image-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
  5. Envie a imagem do Docker para o Artifact Registry

    docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
  6. Criar uma VM de TPU

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v2-8 \
    --version=tpu-ubuntu2204-base
  7. Extraia a imagem do Docker do Artifact Registry em todos os trabalhadores de 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. Executar o contêiner em todos os workers da 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. Executar o script de treinamento em todos os workers da 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"

Quando o script de treinamento for concluído, limpe os recursos.

  1. Desativar o contêiner em todos os workers

    gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
    --zone=europe-west4-a \
    --command="docker kill your-container-name"
  2. Excluir a VM da TPU

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

A seguir