Addestramento con acceleratori TPU

Agent Platform supporta l'addestramento con vari framework e librerie utilizzando una VM TPU. Quando configuri le risorse di calcolo, puoi specificare TPU v2, TPU v3 o TPU v5e VM. 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 i dettagli sulla configurazione delle VM TPU per l'addestramento serverless, consulta Configurare le risorse di calcolo per l'addestramento serverless.

Addestramento TensorFlow

Nota: a partire da TPU v6e, TensorFlow non è più supportato.

Container predefinito

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

Nota: i container predefiniti per TPU v5e non sono disponibili.

Container personalizzato

Utilizza un container personalizzato in cui hai installato le 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.

Nota: Agent Platform supporta solo TensorFlow 2.11 e versioni successive per i container personalizzati e TensorFlow 2.12 e versioni successive per i container di addestramento predefiniti.

Seleziona la versione tensorflow che preferisci e la libreria libtpu corrispondente. Poi, installale 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

Nota: questo non è un Dockerfile completo. Se vuoi utilizzare queste istruzioni, aggiungile a un Dockerfile per un container di addestramento serverless. Puoi anche trovare la versione libtpu corrispondente per la versione di TensorFlow nella pagina delle configurazioni di Cloud TPU.

Pod di TPU

L'addestramento tensorflow su un TPU Pod richiede una configurazione aggiuntiva nel container di addestramento. Agent Platform gestisce un'immagine Docker di base che gestisce la configurazione iniziale.

URI immagine	Versione di 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

Ecco i passaggi per creare il tuo container personalizzato:

1. Scegli l'immagine di base per la versione di Python che preferisci. I wheel TPU 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 i wheel TensorFlow specifici, consulta Configurazioni di Cloud TPU.
2. Estendi l'immagine con il codice di addestramento 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/gemini-enterprise-agent-platform/reference/rest/v1/CustomJobSpec#containerspec)
CMD ["python3", "trainer.py"]

Per vedere un esempio di come addestrare un modello TensorFlow personalizzato utilizzando le TPU in Pipelines, esegui il notebook " Pipelines: TPU model train, upload, and deploy using google-cloud-pipeline-components" in uno dei seguenti ambienti:

Apri in Colab | Apri in Colab Enterprise | Apri in Agent Platform Workbench | Visualizza su GitHub

Addestramento PyTorch

Puoi utilizzare container predefiniti o personalizzati per PyTorch durante l'addestramento con le TPU.

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 la PyTorch libreria.

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 Eseguire il codice PyTorch nelle sezioni di pod di TPU).

Agent Platform attende una risposta da tutti gli host per decidere il completamento del job.

Addestramento JAX

Per eseguire l'addestramento JAX su Agent Platform, devi fornire un'immagine container con il tuo ambiente JAX.

Container predefinito

Non esistono container predefiniti per JAX.

Container personalizzato

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

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 di TPU (vedi Eseguire il codice JAX nelle sezioni di pod di TPU).

Agent Platform osserva 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 tabella seguente descrive in dettaglio le variabili di ambiente che puoi utilizzare all'interno del container:

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

Service account personalizzato

È possibile utilizzare un account di servizio personalizzato per l'addestramento TPU. 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

È possibile utilizzare un IP privato per l'addestramento TPU. Consulta la pagina su come utilizzare un IP privato per l'addestramento serverless.

Controlli di servizio VPC

I progetti con i 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 determinate regioni di Agent Platform.

Tipi di TPU

Per ulteriori informazioni sugli acceleratori TPU, come il limite di memoria, consulta Tipi di TPU.