Preparação com aceleradores de TPU

O Vertex AI suporta a preparação com várias frameworks e bibliotecas através de uma VM de TPU. Ao configurar recursos de computação, pode especificar VMs de TPU v2, TPU v3 ou TPU v5e. A TPU v5e suporta JAX 0.4.6+, TensorFlow 2.15+ e PyTorch 2.1+. A TPU v6e suporta Python 3.10+, JAX 0.4.37+ e PyTorch 2.1+ com PJRT como o tempo de execução predefinido.

Para ver detalhes sobre a configuração de VMs de TPUs para o treino personalizado, consulte o artigo Configure recursos de computação para o treino personalizado.

Formação do TensorFlow

Nota: a partir da TPU v6e, o TensorFlow já não é suportado.

Contentor pré-criado

Use um contentor de preparação pré-criado que suporte TPUs e crie uma aplicação de preparação Python.

Nota: os contentores pré-criados para a TPU v5e não estão disponíveis.

Contentor personalizado

Use um contentor personalizado no qual instalou versões do tensorflow e libtpu criadas especialmente para VMs de TPU. Estas bibliotecas são mantidas pelo serviço Cloud TPU e estão listadas na documentação Configurações de TPU suportadas.

Nota: o Vertex AI só suporta o TensorFlow 2.11 e versões mais recentes para contentores personalizados e o TensorFlow 2.12 e versões mais recentes para contentores de preparação pré-criados.

Selecione a versão tensorflow da sua escolha e a respetiva biblioteca libtpu. Em seguida, instale-os na imagem do contentor Docker quando criar o contentor.

Por exemplo, se quiser usar o TensorFlow 2.15, inclua as seguintes instruções no seu 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

Nota: este não é um Dockerfile completo. Se quiser usar estas instruções, adicione-as a um Dockerfile para um contentor de preparação personalizado. Também pode encontrar a versão libtpu correspondente à versão do TensorFlow na página de configurações do Cloud TPU.

pod TPU

A formação tensorflow num TPU Pod requer uma configuração adicional no contentor de formação. O Vertex AI mantém uma imagem Docker base que processa a configuração inicial.

URIs de imagens	Versão do Python
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	Python 3.8
us-docker.pkg.dev/vertex-ai/training/tf-tpu-pod-base-cp310:latest europe-docker.pkg.dev/vertex-ai/training/tf-tpu-pod-base-cp310:latest asia-docker.pkg.dev/vertex-ai/training/tf-tpu-pod-base-cp310:latest	Python 3.10

Seguem-se os passos para criar o seu contentor personalizado:

1. Escolha a imagem base para a versão do Python que preferir. As wheels do TensorFlow para o TensorFlow 2.12 e inferior suportam o Python 3.8. O TensorFlow 2.13 e superior suporta o Python 3.10 ou superior. Para os wheels específicos do TensorFlow, consulte as configurações do Cloud TPU.
2. Expanda a imagem com o seu código de treinador e o comando de arranque.

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

# 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

# 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.
# Use CMD, not ENTRYPOINT, to avoid accidentally overriding the pod base image's ENTRYPOINT.
# 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 ver um exemplo de como preparar um modelo TensorFlow personalizado usando TPUs em pipelines da Vertex AI, execute o bloco de notas "Vertex AI Pipelines: TPU model train, upload, and deploy using google-cloud-pipeline-components" num dos seguintes ambientes:

Abrir no Colab | Abrir no Colab Enterprise | Abrir no Vertex AI Workbench | Ver no GitHub

Formação do PyTorch

Pode usar contentores pré-criados ou personalizados para o PyTorch quando fizer a preparação com as TPUs.

Contentor pré-criado

Use um contentor de preparação pré-criado que suporte TPUs e crie uma aplicação de preparação Python.

Contentor personalizado

Use um contentor personalizado no qual instalou a biblioteca PyTorch.

Por exemplo, o seu Dockerfile pode ter o seguinte aspeto:

FROM python:3.10

# v5e, v6e 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 TPU

A preparação é executada em todos os anfitriões do TPU Pod (consulte o artigo Execute código PyTorch em fatias do TPU Pod).

O Vertex AI aguarda uma resposta de todos os anfitriões para decidir a conclusão da tarefa.

Formação em JAX

Contentor pré-criado

Não existem contentores pré-criados para o JAX.

Contentor personalizado

Use um contentor personalizado no qual instalou a biblioteca JAX.

Por exemplo, o seu Dockerfile pode ter o seguinte aspeto:

# 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 TPU

A preparação é executada em todos os anfitriões do TPU Pod (consulte o artigo Execute código JAX em fatias do TPU Pod).

O Vertex AI monitoriza o primeiro anfitrião do TPU Pod para decidir a conclusão da tarefa. Pode usar o seguinte fragmento de código para garantir que todos os anfitriões saem 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 seguinte detalha as variáveis de ambiente que pode usar no contentor:

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

Pode usar uma conta de serviço personalizada para o treino de TPUs. Para saber como usar uma conta de serviço personalizada, consulte a página sobre como usar uma conta de serviço personalizada.

IP privado (intercâmbio da rede da VPC) para a preparação

Um IP privado pode ser usado para o treino de TPUs. Consulte a página sobre como usar um IP privado para o treino personalizado.

VPC Service Controls

Os projetos com o VPC Service Controls ativado podem enviar tarefas de preparação de TPUs.

Limitações

Aplicam-se as seguintes limitações quando faz a preparação com uma VM de TPU:

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

Tipos de TPUs

Consulte os tipos de TPUs para ver mais informações sobre os aceleradores de TPUs, como o limite de memória.