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 o TPU V3. Per maggiori dettagli, consulta Configurare risorse di computing per l'addestramento personalizzato.
Addestramento TensorFlow
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 le versioni di tensorflow
e libtpu
create appositamente per le VM TPU. Queste librerie sono gestite dal servizio Cloud TPU e sono elencate nella documentazione delle configurazioni TPU supportate.
Seleziona la versione tensorflow
che preferisci e la libreria libtpu
corrispondente. e installali nell'immagine del container Docker
al momento della creazione del container.
Ad esempio, se vuoi utilizzare TensorFlow 2.12, includi le seguenti istruzioni nel tuo Dockerfile:
# 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
pod di TPU
L'addestramento di tensorflow
su un TPU Pod
richiede una configurazione aggiuntiva nel container di addestramento. Vertex AI mantiene un'immagine Docker di base che gestisce
la configurazione iniziale.
URI immagine | Versione Python |
---|---|
|
3,8 |
Ecco i passaggi per creare un container personalizzato:
- Scegli l'immagine di base per la versione Python di tua scelta. Le ruote TPU TensorFlow supportano Python 3.8 per le versioni 2.12 e precedenti. A partire dalla versione 2.13, sarà supportato Python 3.10. Per le ruote TensorFlow specifiche, consulta la pagina delle configurazioni di Cloud TPU.
- Estendi l'immagine con il tuo codice trainer 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"]
Addestramento PyTorch
Puoi utilizzare container predefiniti o personalizzati per PyTorch durante l'addestramento con TPU.
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 tuo Dockerfile potrebbe avere il seguente aspetto:
FROM python:3.8
# install pytorch and torch_xla
RUN pip3 install torch~=2.0.0 https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-2.0-cp38-cp38-linux_x86_64.whl
# 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 su sezioni di 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
hai installato la libreria JAX
.
Ad esempio, il tuo Dockerfile potrebbe avere il seguente aspetto:
# Install JAX.
RUN pip install 'jax[tpu]>=0.2.16' -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 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 assicurarti 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 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"} |
Account di servizio personalizzato
Per l'addestramento delle TPU è possibile usare un account di servizio personalizzato. 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
È possibile usare un IP privato per l'addestramento di TPU. Consulta la 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, si applicano le seguenti limitazioni:
Le TPU sono disponibili solo in determinate regioni di Vertex AI.
Se utilizzi una VM TPU, non puoi criptare nessun materiale con una chiave di crittografia gestita dal cliente (CMEK).
Tipi di TPU
Per ulteriori informazioni sugli acceleratori TPU, come il limite di memoria, consulta la sezione Tipi di TPU.