Addestramento con acceleratori TPU

Vertex AI supporta l'addestramento con vari framework e librerie utilizzando una VM TPU. Quando configuri le risorse di calcolo, puoi specificare VM TPU v2, TPU v3 o TPU v5e. TPU v5e supporta JAX 0.4.6+, TensorFlow 2.15+ e PyTorch 2.1+. TPU v6e supporta Python 3.10+, JAX 0.4.37+ e PyTorch 2.1+ utilizzando PJRT come runtime predefinito.

Per informazioni dettagliate sulla configurazione delle VM TPU per l'addestramento personalizzato, vedi Configurazione delle risorse di computing per l'addestramento personalizzato.

Addestramento TensorFlow

Container predefinito

Utilizza un container di addestramento predefinito che supporta le TPU e crea un'applicazione di addestramento Python.

Container personalizzato

Utilizza un container personalizzato in cui hai installato versioni di tensorflow e libtpu create appositamente per le VM TPU. Queste librerie vengono gestite dal servizio Cloud TPU e sono elencate nella documentazione relativa alle configurazioni TPU supportate.

Seleziona la versione tensorflow che preferisci e la relativa raccolta libtpu. Successivamente, installali nell'immagine container Docker quando crei il container.

Ad esempio, se vuoi utilizzare TensorFlow 2.15, includi le seguenti istruzioni nel 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

pod di TPU

L'addestramento di tensorflow su un TPU Pod richiede una configurazione aggiuntiva nel contenitore di addestramento. Vertex AI gestisce un'immagine Docker di base che gestisce la configurazione iniziale.

Ecco i passaggi per creare il tuo contenitore personalizzato:

1. Scegli l'immagine di base per la versione Python che preferisci. TPU Le ruote TensorFlow per TensorFlow 2.12 e versioni precedenti supportano Python 3.8. TensorFlow 2.13 e versioni successive supportano Python 3.10 o versioni successive. Per le ruote TensorFlow specifiche, vedi Configurazioni TPU.
2. Estendi l'immagine con il codice dell'addestratore e il comando di avvio.

# 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"]

Addestramento PyTorch

Quando esegui l'addestramento con le TPU, puoi utilizzare container predefiniti o personalizzati per PyTorch.

Container predefinito

Utilizza un container di addestramento predefinito che supporti le TPU e crea un'applicazione di addestramento Python.

Container personalizzato

Utilizza un container personalizzato in cui hai installato la libreria PyTorch.

Ad esempio, il Dockerfile potrebbe avere il seguente aspetto:

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

L'addestramento viene eseguito su tutti gli host del pod di TPU (vedi Esegui il codice PyTorch nelle sezioni di pod di TPU TPU).

Vertex AI attende una risposta da tutti gli host per decidere il completamento del job.

Addestramento JAX

Container predefinito

Non esistono container predefiniti per JAX.

Container personalizzato

Utilizza un container personalizzato in cui hai installato la libreria JAX.

Ad esempio, il Dockerfile potrebbe avere il seguente aspetto:

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

L'addestramento viene eseguito su tutti gli host del pod TPU (vedi Esecuzione del codice JAX nelle sezioni di pod di TPU).

Vertex AI monitora il primo host del pod di TPU per decidere il completamento del job. Puoi utilizzare il seguente snippet di codice per assicurarti che tutti gli host escano contemporaneamente:

# 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()

Variabili di ambiente

La seguente tabella descrive in dettaglio le variabili di ambiente che puoi utilizzare all'interno del container:

Service account personalizzato

Per l'addestramento TPU può essere utilizzato un account di servizio personalizzato. Per informazioni su come utilizzare un account di servizio personalizzato, consulta la pagina su come utilizzare un service account personalizzato.

IP privato (peering di rete VPC) per l'addestramento

Per l'addestramento della TPU può essere utilizzato un IP privato. Consulta la pagina su come utilizzare un IP privato per l'addestramento personalizzato.

Controlli di servizio VPC

I progetti con Controlli di servizio VPC abilitati possono inviare job di addestramento TPU.

Limitazioni

Quando esegui l'addestramento utilizzando una VM TPU, si applicano le seguenti limitazioni:

• Le TPU sono disponibili solo in alcune regioni di Vertex AI.

Tipi di TPU

Consulta la sezione Tipi di TPU per maggiori informazioni sugli acceleratori TPU, ad esempio il limite di memoria.