Training mit TPU-Beschleunigern

Vertex AI unterstützt das Training mit verschiedenen Frameworks und Bibliotheken mit einer TPU-VM. Beim Konfigurieren von Rechenressourcen können Sie TPU V2- oder TPU V3-VMs angeben. Weitere Informationen finden Sie unter Computing-Ressourcen für benutzerdefiniertes Training konfigurieren.

TensorFlow-Training

Vordefinierter Container

Verwenden Sie einen vordefinierten Trainingscontainer, der TPUs unterstützt, und erstellen Sie eine Python-Trainingsanwendung.

Benutzerdefinierter Container

Verwenden Sie einen benutzerdefinierten Container, in dem Sie Versionen von tensorflow und libtpu installiert haben, die speziell für TPU-VMs erstellt wurden. Diese Bibliotheken werden vom Cloud TPU-Dienst verwaltet und in der Dokumentation Unterstützte TPU-Konfigurationen aufgeführt.

Hinweis: Vertex AI unterstützt nur TensorFlow 2.11 und neuere Versionen für benutzerdefinierte Container und TensorFlow 2.12 und neuere Versionen für vordefinierte Trainingscontainer.

Wählen Sie die gewünschte tensorflow-Version und die entsprechende libtpu-Bibliothek aus. Installieren Sie diese dann beim Erstellen des Containers im Docker-Container-Image.

Wenn Sie beispielsweise TensorFlow 2.12 verwenden möchten, fügen Sie Ihrem Dockerfile die folgende Anleitung hinzu:

# 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

Hinweis: Dies ist kein vollständiges Dockerfile. Wenn Sie diese Anweisungen verwenden möchten, fügen Sie sie einem Dockerfile für einen benutzerdefinierten Trainingscontainer hinzu. Die entsprechende libtpu-Version für die TensorFlow-Version finden Sie auch auf der Cloud TPU-Konfigurationsseite.

TPU-Pod

Das Training von tensorflow auf einem TPU Pod erfordert eine zusätzliche Einrichtung im Trainingscontainer. Vertex AI verwaltet ein Basis-Docker-Image, das die Ersteinrichtung übernimmt.

Image-URIs	Python-Version
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

So erstellen Sie Ihren benutzerdefinierten Container:

1. Wählen Sie das Basis-Image für die Python-Version Ihrer Wahl aus. TPU TensorFlow Wheels unterstützen Python 3.8 für Version 2.12 und niedriger. Ab Version 2.13 wird Python 3.10 unterstützt. Informationen zu den jeweiligen TensorFlow-Räumen finden Sie auf der Seite für die Cloud TPU-Konfigurationen.
2. Erweitern Sie das Image mit Ihrem Trainercode und dem Startbefehl.

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

Wenn Sie ein Beispiel für das Trainieren eines benutzerdefinierten TensorFlow-Modells mithilfe von TPUs in Vertex AI Pipelines sehen möchten, führen Sie das Jupyter-Notebook "Vertex AI Pipelines: TPU model model training, upload and deploy with google-cloud-pipeline-components" in einer der folgenden Umgebungen aus:

In Colab öffnen | In Colab Enterprise öffnen | In nutzerverwalteten Vertex AI Workbench-Notebooks öffnen | Auf GitHub ansehen

PyTorch-Training

Sie können für das Training mit TPUs vordefinierte oder benutzerdefinierte Container für PyTorch verwenden.

Vordefinierter Container

Verwenden Sie einen vordefinierten Trainingscontainer, der TPUs unterstützt, und erstellen Sie eine Python-Trainingsanwendung.

Benutzerdefinierter Container

Verwenden Sie einen benutzerdefinierten Container, in dem Sie die PyTorch-Bibliothek installiert haben.

Ihr Dockerfile könnte beispielsweise so aussehen:

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

TPU-Pod

Das Training wird auf allen Hosts des TPU-Pods ausgeführt (siehe PyTorch-Code auf TPU-Pod-Slices ausführen).

Vertex AI wartet auf eine Antwort von allen Hosts, um den Abschluss des Jobs zu entscheiden.

JAX-Training

Vordefinierter Container

Es gibt keine vordefinierten Container für JAX.

Benutzerdefinierter Container

Verwenden Sie einen benutzerdefinierten Container, in dem Sie die JAX-Bibliothek installiert haben.

Ihr Dockerfile könnte beispielsweise so aussehen:

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

TPU-Pod

Das Training wird auf allen Hosts des TPU-Pods ausgeführt (siehe JAX-Code auf TPU-Pod-Slices ausführen).

Vertex AI überwacht den ersten Host des TPU-Pods, um über den Abschluss des Jobs zu entscheiden. Mit dem folgenden Code-Snippet können Sie dafür sorgen, dass alle Hosts gleichzeitig beendet werden:

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

Umgebungsvariablen

In der folgenden Tabelle werden die Umgebungsvariablen beschrieben, die Sie im Container verwenden können:

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

Benutzerdefiniertes Dienstkonto

Ein benutzerdefiniertes Dienstkonto kann für das TPU-Training verwendet werden. Informationen zur Verwendung eines benutzerdefinierten Dienstkontos finden Sie auf der Seite zur Verwendung eines benutzerdefinierten Dienstkontos.

Private IP-Adresse (VPC-Netzwerk-Peering) für das Training

Eine private IP-Adresse kann für das TPU-Training verwendet werden. Informationen dazu finden Sie auf der Seite Private IP-Adresse für benutzerdefiniertes Training verwenden.

VPC Service Controls

VPC Service Controls-fähige Projekte können TPU-Trainingsjobs senden.

Beschränkungen

Beim Trainieren mit einer TPU-VM gelten die folgenden Einschränkungen:

• TPUs sind nur in bestimmten Vertex AI-Regionen verfügbar.

• Wenn Sie eine TPU-VM verwenden, können Sie keine Materialien mit einem vom Kunden verwalteten Verschlüsselungsschlüssel (Customer-Managed Encryption Key, CMEK) verschlüsseln.

TPU-Typen

Weitere Informationen zu TPU-Beschleunigern wie das Arbeitsspeicherlimit finden Sie unter TPU-Typen.