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
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
Conecte-se à VM da TPU usando SSH
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
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.
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
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"
Clone o repositório PyTorch XLA
git clone --recursive https://github.com/pytorch/xla.git
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.
- Digite
exit
para sair do contêiner do Docker. - Digite
exit
para sair da VM da TPU. 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.
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
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'
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"
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
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
Conecte-se à VM da TPU usando SSH
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Iniciar o daemon do Docker na VM de TPU
sudo systemctl start docker
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
Instalar o JAX
pip install jax[tpu]
Instalar o FLAX
pip install --upgrade clu git clone https://github.com/google/flax.git pip install --user -e flax
Instalar pacotes
tensorflow
etensorflow-dataset
pip install tensorflow pip install tensorflow-datasets
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.
- Digite
exit
para sair do contêiner do Docker. - Digite
exit
para sair da VM da TPU. 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.
Crie um arquivo chamado
Dockerfile
no diretório atual e cole o texto a seguirFROM 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
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
Compilar a imagem Docker
docker build -t your-image-name .
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
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
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
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"
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"
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.
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"
Excluir a VM da TPU
gcloud compute tpus tpu-vm delete your-tpu-name \ --zone=europe-west4-a
A seguir
- Tutoriais da Cloud TPU
- Gerenciar TPUs
- Arquitetura do Cloud TPU
- Executar o código JAX em frações de TPU