Como treinar com aceleradores de TPU

A Vertex AI é compatível com o treinamento com vários frameworks e bibliotecas usando uma VM de TPU. Ao configurar recursos de computação, é possível especificar VMs de TPU v2, TPU v3 ou TPU v5e. A TPU v5e é compatível com JAX 0.4.6+, TensorFlow 2.15+ e PyTorch 2.1+. Para detalhes sobre como configurar VMs de TPU para treinamento personalizado, consulte Configurar recursos de computação para treinamento personalizado.

Treinamento do TensorFlow

Contêiner pré-criado

Use um contêiner de treinamento pré-criado compatível com TPUs e crie um aplicativo de treinamento em Python.

Observação: contêineres pré-criados para a TPU v5e não estão disponíveis.

Contêiner personalizado

Use um contêiner personalizado em que você instalou versões de tensorflow e libtpu especialmente criadas para VMs de TPU. Essas bibliotecas são mantidas pelo serviço do Cloud TPU e estão listadas na documentação sobre Configurações de TPU compatíveis.

Observação: a Vertex AI é compatível apenas com as versões do TensorFlow 2.11 e mais recentes para contêineres personalizados e do TensorFlow 2.12 e versões mais recentes para contêineres de treinamento pré-criados.

Selecione a versão tensorflow e a biblioteca libtpu correspondente de sua escolha. Em seguida, instale-as na imagem do contêiner do Docker ao criar o contêiner.

Por exemplo, se você quiser usar o TensorFlow 2.12, inclua as seguintes instruções no Dockerfile:

  # Download and install `tensorflow`.
  RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.15.0/tensorflow-2.15.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

  # Download and install `libtpu`.
  # You must save `libtpu.so` in the '/lib' directory of the container image.
  RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.9.0/libtpu.so -o /lib/libtpu.so

  # TensorFlow training on TPU v5e requires the PJRT runtime. To enable the PJRT
  # runtime, configure the following environment variables in your Dockerfile.
  # For details, see https://cloud.google.com/tpu/docs/runtimes#tf-pjrt-support.
  # ENV NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  # ENV TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

Observação: este não é um Dockerfile completo. Se você quiser usar essas instruções, adicione-as a um Dockerfile para um contêiner de treinamento personalizado. Acesse a página de configurações do Cloud TPU para encontrar a versão correspondente do libtpu para a versão do TensorFlow.

Pod de TPU

O treinamento tensorflow em um TPU Pod requer configuração adicional no contêiner de treinamento. A Vertex AI mantém uma imagem do Docker básica que processa a configuração inicial.

URIs de imagem	Versão do Python e do TPU
us-docker.pkg.dev/vertex-ai/training/tf-tpu-pod-base-cp38:latest europe-docker.pkg.dev/vertex-ai/training/tf-tpu-pod-base-cp38:latest asia-docker.pkg.dev/vertex-ai/training/tf-tpu-pod-base-cp38:latest	3.8
us-docker.pkg.dev/vertex-ai/training/tf-tpu.2-15-pod-base-cp310:latest europe-docker.pkg.dev/vertex-ai/training/tf-tpu.2-15-pod-base-cp310:latest asia-docker.pkg.dev/vertex-ai/training/tf-tpu.2-15-pod-base-cp310:latest	3.10

Siga estas etapas para criar seu contêiner personalizado:

1. Escolha a imagem de base para a versão do Python de sua escolha. As rodas do TensorFlow de TPU para TensorFlow 2.12 e versões anteriores são compatíveis com Python 3.8. O TensorFlow 2.13 e versões mais recentes oferecem suporte ao Python 3.10 ou versões mais recentes. Para conhecer as rodas específicas do TensorFlow, consulte Configurações do Cloud TPU.
2. Amplie a imagem com seu código treinador e o comando de inicialização.

# Specifies base image and tag
FROM us-docker.pkg.dev/vertex-ai/training/tf-tpu-pod-base-cp38:latest
WORKDIR /root

# Download and install `tensorflow`.
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

# Download and install `libtpu`.
# You must save `libtpu.so` in the '/lib' directory of the container image.
RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so

# Copies the trainer code to the docker image.
COPY your-path-to/model.py /root/model.py
COPY your-path-to/trainer.py /root/trainer.py

# The base image is setup so that it runs the CMD that you provide.
# You can provide CMD inside the Dockerfile like as follows.
# Alternatively, you can pass it as an `args` value in ContainerSpec:
# (https://cloud.google.com/vertex-ai/docs/reference/rest/v1/CustomJobSpec#containerspec)
CMD ["python3", "trainer.py"]

Para conferir um exemplo de como treinar um modelo personalizado do TensorFlow usando TPUs no Vertex AI Pipelines, execute o notebook Jupyter "Vertex AI Pipelines: treinamento, upload e implantação de modelos de TPU usando google-cloud-pipeline-components" em um dos seguintes ambientes:

Abrir no Colab | Abrir no Colab Enterprise | Abrir nos blocos gerenciados pelo usuário do Vertex AI Workbench | Ver no GitHub

Treinamento do PyTorch

É possível usar contêineres pré-criados ou personalizados para o PyTorch ao treinar com TPUs.

Contêiner pré-criado

Use um contêiner de treinamento pré-criado compatível com TPUs e crie um aplicativo de treinamento em Python.

Contêiner personalizado

Use um contêiner personalizado em que você instalou a biblioteca PyTorch.

Por exemplo, seu Dockerfile pode ter esta aparência:

FROM python:3.10

# v5e specific requirement - enable PJRT runtime
ENV PJRT_DEVICE=TPU

# install pytorch and torch_xla
RUN pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0
 -f https://storage.googleapis.com/libtpu-releases/index.html

# Add your artifacts here
COPY trainer.py .

# Run the trainer code
CMD ["python3", "trainer.py"]

Pod de TPU

O treinamento é executado em todos os hosts do Pod de TPU. Consulte Executar o código PyTorch em frações do Pod de TPU.

A Vertex AI aguarda uma resposta de todos os hosts para decidir a conclusão do job.

Treinamento do JAX

Contêiner pré-criado

Não há contêineres pré-criados para JAX.

Contêiner personalizado

Use um contêiner personalizado em que você instalou a biblioteca JAX.

Por exemplo, seu Dockerfile pode ter esta aparência:

# Install JAX.
RUN pip install 'jax[tpu]>=0.4.6' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

# Add your artifacts here
COPY trainer.py trainer.py

# Set an entrypoint.
ENTRYPOINT ["python3", "trainer.py"]

Pod de TPU

O treinamento é executado em todos os hosts do Pod de TPU. Consulte Executar o código JAX em frações do Pod de TPU.

A Vertex AI monitora o primeiro host do Pod de TPU para decidir a conclusão do job. Use o seguinte snippet de código para garantir que todos os hosts saiam ao mesmo tempo:

# Your training logic
...

if jax.process_count() > 1:
  # Make sure all hosts stay up until the end of main.
  x = jnp.ones([jax.local_device_count()])
  x = jax.device_get(jax.pmap(lambda x: jax.lax.psum(x, 'i'), 'i')(x))
  assert x[0] == jax.device_count()

Variáveis de ambiente

A tabela a seguir detalha as variáveis de ambiente que podem ser usadas no contêiner:

Nome	Valor
TPU_NODE_NAME	my-first-tpu-node
TPU_CONFIG	{"project": "tenant-project-xyz", "zone": "us-central1-b", "tpu_node_name": "my-first-tpu-node"}

Conta de serviço personalizada

Uma conta de serviço personalizada pode ser usada para o treinamento da TPU. Para saber como usar uma conta de serviço personalizada, consulte a página sobre como usar uma conta de serviço personalizada.

IP particular (peering de rede VPC) para treinamento

Um IP privado pode ser usado no treinamento da TPU. Consulte a página sobre como usar um IP privado para treinamento personalizado.

VPC Service Controls

Os projetos ativados pelo VPC Service Controls podem enviar jobs de treinamento de TPU.

Limitações

As limitações a seguir se aplicam quando você treina usando uma VM de TPU:

• As TPUs só estão disponíveis em determinadas regiões da Vertex AI.

Tipos de TPU

Consulte Tipos de TPU para mais informações sobre aceleradores de TPU, como limite de memória.