Entraînement avec des accélérateurs TPU

Vertex AI vous permet d'effectuer des entraînements avec plusieurs frameworks et bibliothèques à l'aide d'une VM TPU. Lors de la configuration des ressources de calcul, vous pouvez spécifier des VM TPU v2, TPU v3 ou TPU v5e. TPU v5e est compatible avec JAX 0.4.6+, TensorFlow 2.15+ et PyTorch 2.1+. Pour en savoir plus sur la configuration des VM TPU pour l'entraînement personnalisé, consultez Configurer des ressources de calcul pour l'entraînement personnalisé.

Entraînement TensorFlow

Conteneur prédéfini

Utilisez un conteneur d'entraînement prédéfini compatible avec les TPU et créez une application d'entraînement Python.

Conteneur personnalisé

Utilisez un conteneur personnalisé dans lequel vous avez installé des versions de tensorflow et de libtpu spécialement conçues pour les VM TPU. Ces bibliothèques sont gérées par le service Cloud TPU et sont répertoriées dans la documentation sur les configurations TPU compatibles.

Sélectionnez la version tensorflow de votre choix et la bibliothèque libtpu correspondante. Installez-les ensuite dans votre image de conteneur Docker lors de sa création.

Par exemple, si vous souhaitez utiliser TensorFlow 2.12, incluez les instructions suivantes dans votre fichier Dockerfile :

  # Download and install `tensorflow`.
  RUN pip install http://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 http://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 http://cloud.go888ogle.com.fqhub.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 TPU

L'entraînement tensorflow sur un TPU Pod nécessite une configuration supplémentaire dans le conteneur d'entraînement. Vertex AI conserve une image Docker de base qui gère la configuration initiale.

URI d'images	Version Python et version de TPU
`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`	3.8
`us-docker.pkg.dev/vertex-ai/training/tf-tpu.2-15-pod-base-cp310:latest` `europe-docker.pkg.dev/vertex-ai/training/tf-tpu.2-15-pod-base-cp310:latest` `asia-docker.pkg.dev/vertex-ai/training/tf-tpu.2-15-pod-base-cp310:latest`	3,10

Pour créer votre conteneur personnalisé, procédez comme suit :

Choisissez l'image de base pour la version Python de votre choix. Les roues TPU pour les versions 2.12 et antérieures de TensorFlow sont compatibles avec Python 3.8. Les versions 2.13 et ultérieures de TensorFlow sont compatibles avec Python 3.10 ou version ultérieure. Pour en savoir plus sur les roues TensorFlow spécifiques, consultez Configurations de Cloud TPU.
Développez l'image avec votre code d'entraînement et la commande de démarrage.

# 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 http://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 http://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:
# (http://cloud.go888ogle.com.fqhub.com/vertex-ai/docs/reference/rest/v1/CustomJobSpec#containerspec)
CMD ["python3", "trainer.py"]

Entraînement PyTorch

Vous pouvez utiliser des conteneurs prédéfinis ou personnalisés pour PyTorch lors de l'entraînement avec des TPU.

Conteneur prédéfini

Utilisez un conteneur d'entraînement prédéfini compatible avec les TPU et créez une application d'entraînement Python.

Conteneur personnalisé

Utilisez un conteneur personnalisé dans lequel vous avez installé la bibliothèque PyTorch.

Par exemple, votre Dockerfile peut se présenter sous la forme suivante :

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 http://storage.googleapis.com/libtpu-releases/index.html

# Add your artifacts here
COPY trainer.py .

# Run the trainer code
CMD ["python3", "trainer.py"]

pod TPU

L'entraînement s'exécute sur tous les hôtes du pod TPU (consultez la section Exécuter du code PyTorch sur des tranches de pod TPU).

Vertex AI attend une réponse de tous les hôtes pour décider de la fin du job.

Entraînement JAX

Conteneur prédéfini

Aucun conteneur prédéfini n'est disponible pour JAX.

Conteneur personnalisé

Utilisez un conteneur personnalisé dans lequel vous avez installé la bibliothèque JAX.

Par exemple, votre Dockerfile peut se présenter sous la forme suivante :

# Install JAX.
RUN pip install 'jax[tpu]>=0.4.6' -f http://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

L'entraînement s'exécute sur tous les hôtes du pod TPU (consultez la section Exécuter du code JAX sur des tranches de pods TPU).

Vertex AI surveille le premier hôte du pod TPU pour déterminer la fin de la tâche. Vous pouvez utiliser l'extrait de code suivant pour vous assurer que tous les hôtes se ferment en même temps :

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

Variables d'environnement

Le tableau suivant détaille les variables d'environnement que vous pouvez utiliser dans le conteneur :

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

Compte de service personnalisé

Un compte de service personnalisé peut être utilisé pour l'entraînement TPU. Pour utiliser un compte de service personnalisé, consultez la page Comment utiliser un compte de service personnalisé.

Adresse IP privée (appairage de réseaux VPC) pour l'entraînement

Vous pouvez utiliser une adresse IP privée pour l'entraînement TPU. Pour plus d'informations, reportez-vous à la page Utiliser une adresse IP privée pour l'entraînement personnalisé.

VPC Service Controls

Les projets pour lesquels VPC Service Controls est activé peuvent envoyer des jobs d'entraînement TPU.

Limites

Les limites suivantes s'appliquent lorsque vous effectuez un entraînement à l'aide d'une VM TPU :

Les TPU ne sont disponibles que dans certaines régions Vertex AI.
Si vous utilisez une VM TPU, vous ne pouvez pas chiffrer des ressources avec une clé de chiffrement gérée par le client (CMEK).

Types de TPU

Pour en savoir plus sur les accélérateurs TPU, tels que la limite de mémoire, consultez la page Types de TPU.