Vertex AI supporta l'addestramento con vari framework e librerie utilizzando una VM TPU. Quando configuri le risorse di computing, puoi specificare TPU VM v2, TPU v3 o TPU v5e. TPU v5e supporta JAX 0.4.6+, TensorFlow 2.15+ e PyTorch 2.1+. Per Per maggiori dettagli sulla configurazione delle VM TPU per l'addestramento personalizzato, consulta Configurare le risorse per l'addestramento personalizzato.
Addestramento TensorFlow
Container predefinito
Utilizza un container di addestramento predefinito supporta le TPU e crea un'istanza Applicazione di addestramento Python.
Container personalizzato
Utilizza un container personalizzato in cui
hanno installato versioni di tensorflow e libtpu appositamente
per le VM TPU. Queste librerie sono gestite dalla Cloud TPU
e sono elencati nel
Configurazioni TPU supportate
documentazione.
Seleziona la versione tensorflow di tua scelta e la versione libtpu corrispondente
libreria. Poi, installale nell'immagine container Docker quando crei
nel container.
Ad esempio, se vuoi utilizzare TensorFlow 2.12, includi 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 tensorflow su un TPU Pod richiede una configurazione aggiuntiva durante l'addestramento
containerizzato. Vertex AI gestisce un'immagine Docker di base che gestisce
la configurazione iniziale.
| URI immagine | Versione Python e versione TPU |
|---|---|
|
3,8 |
|
3,10 |
Per creare un container personalizzato:
- Scegli l'immagine di base per la versione Python che preferisci. TPU Supporto per TensorFlow 2.12 e versioni precedenti di TensorFlow Wheels Python 3.8. TensorFlow 2.13 e versioni successive supportano Python 3.10 o maggiori. Per gli specifici modelli TensorFlow, consulta Configurazioni di Cloud TPU.
- 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-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"]
Formazione 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
hanno installato la libreria PyTorch.
Ad esempio, il Dockerfile potrebbe avere il seguente aspetto:
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 di TPU
L'addestramento viene eseguito su tutti gli host del pod di TPU (vedi Esegui il codice PyTorch sulle pod di TPU TPU).
Vertex AI attende una risposta da tutti gli host per decidere il completamento del job.
formazione JAX
Container predefinito
Non ci sono container predefiniti per JAX.
Container personalizzato
Utilizza un container personalizzato in cui
hanno 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 di TPU (vedi Esegui codice JAX sulle pod di TPU TPU).
Vertex AI controlla il primo host del pod di TPU per decidere il completamento del job. Puoi usare il seguente snippet di codice per verificare in modo che tutti gli host si chiudano 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 il 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"} |
Account di servizio personalizzato
Un account di servizio personalizzato può essere utilizzato per l'addestramento delle TPU. Per informazioni su come utilizzare un account di servizio personalizzato, consulta la pagina su come utilizzare un account di servizio personalizzato.
IP privato (peering di rete VPC) per l'addestramento
Un IP privato può essere utilizzato per l'addestramento delle TPU. Fai riferimento alla pagina su come utilizzare un IP privato per l'addestramento personalizzato.
Controlli di servizio VPC
I progetti abilitati per i Controlli di servizio VPC possono inviare job di addestramento TPU.
Limitazioni
Quando esegui l'addestramento utilizzando una VM TPU:
Tipi di TPU
Consulta la sezione Tipi di TPU per Ulteriori informazioni sugli acceleratori TPU, come il limite di memoria.