Pelatihan dengan akselerator TPU

Vertex AI mendukung pelatihan dengan berbagai framework dan library menggunakan VM TPU. Saat mengonfigurasi resource komputasi, Anda dapat menentukan VM TPU v2, TPU v3, atau TPU v5e. TPU v5e mendukung JAX 0.4.6+, TensorFlow 2.15+, dan PyTorch 2.1+. Untuk mengetahui detail tentang cara mengonfigurasi VM TPU untuk pelatihan kustom, lihat Mengonfigurasi resource komputasi untuk pelatihan kustom.

Pelatihan TensorFlow

Container bawaan

Menggunakan container pelatihan bawaan yang mendukung TPU, dan membuat aplikasi pelatihan Python.

Catatan: Container bawaan untuk TPU v5e tidak tersedia.

Container kustom

Gunakan container kustom tempat Anda telah menginstal versi tensorflow dan libtpu yang dibuat khusus untuk VM TPU. Library ini dikelola oleh layanan Cloud TPU dan tercantum di dokumentasi Konfigurasi TPU yang didukung.

Catatan: Vertex AI hanya mendukung TensorFlow 2.11 dan versi yang lebih baru untuk penampung kustom dan TensorFlow 2.12 dan versi yang lebih baru untuk penampung pelatihan bawaan.

Pilih versi tensorflow pilihan Anda dan library libtpu yang sesuai. Selanjutnya, instal file tersebut di image container Docker saat Anda membuat container.

Misalnya, jika Anda ingin menggunakan TensorFlow 2.12, sertakan petunjuk berikut di Dockerfile Anda:

  # 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

Catatan: Ini bukan Dockerfile yang lengkap. Jika Anda ingin menggunakan petunjuk ini, tambahkan petunjuk ke Dockerfile untuk container pelatihan kustom. Anda juga dapat menemukan versi libtpu yang sesuai untuk versi TensorFlow dari halaman konfigurasi Cloud TPU.

Pod TPU

Pelatihan tensorflow pada TPU Pod memerlukan penyiapan tambahan di container pelatihan. Vertex AI mengelola image Docker dasar yang menangani penyiapan awal.

URI Image	Versi Python dan Versi 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

Berikut adalah langkah-langkah untuk membuat container kustom:

1. Pilih image dasar untuk versi Python pilihan Anda. Wheel TensorFlow TPU untuk TensorFlow 2.12 dan yang lebih rendah mendukung Python 3.8. TensorFlow 2.13 dan yang lebih baru mendukung Python 3.10 atau yang lebih baru. Untuk wheel TensorFlow tertentu, lihat Konfigurasi Cloud TPU.
2. Perluas gambar dengan kode pelatih dan perintah startup Anda.

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

Untuk melihat contoh cara melatih model TensorFlow kustom menggunakan TPU di Vertex AI Pipelines, jalankan notebook Jupyter "Vertex AI Pipelines: TPU model train, upload, and deploy using google-cloud-pipeline-components" di salah satu lingkungan berikut:

Buka di Colab | Buka di Colab Enterprise | Buka di notebook Vertex AI Workbench yang dikelola pengguna | Lihat di GitHub

Pelatihan PyTorch

Anda dapat menggunakan container bawaan atau kustom untuk PyTorch saat melakukan pelatihan dengan TPU.

Container bawaan

Menggunakan container pelatihan bawaan yang mendukung TPU, dan membuat aplikasi pelatihan Python.

Container kustom

Gunakan container kustom tempat Anda menginstal library PyTorch.

Misalnya, Dockerfile Anda dapat terlihat seperti berikut:

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 TPU

Pelatihan ini berjalan pada semua host Pod TPU (lihat Menjalankan kode PyTorch pada slice Pod TPU).

Vertex AI menunggu respons dari semua host untuk memutuskan penyelesaian tugas.

Pelatihan JAX

Container bawaan

Tidak ada container bawaan untuk JAX.

Container kustom

Gunakan container kustom tempat Anda menginstal library JAX.

Misalnya, Dockerfile Anda dapat terlihat seperti berikut:

# 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 TPU

Pelatihan ini berjalan pada semua host Pod TPU (lihat Menjalankan kode JAX pada slice Pod TPU).

Vertex AI mengamati host pertama Pod TPU untuk memutuskan penyelesaian tugas. Anda dapat menggunakan cuplikan kode berikut untuk memastikan bahwa semua host keluar pada waktu yang sama:

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

Variabel lingkungan

Tabel berikut menjelaskan variabel lingkungan yang dapat Anda gunakan di container:

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

Akun Layanan Kustom

Akun layanan kustom dapat digunakan untuk pelatihan TPU. Untuk mengetahui cara menggunakan akun layanan kustom, lihat halaman tentang cara menggunakan akun layanan kustom.

IP pribadi (peering jaringan VPC) untuk pelatihan

IP pribadi dapat digunakan untuk pelatihan TPU. Lihat halaman tentang cara menggunakan IP pribadi untuk pelatihan kustom.

Kontrol Layanan VPC

Project yang mengaktifkan Kontrol Layanan VPC dapat mengirimkan tugas pelatihan TPU.

Batasan

Keterbatasan berikut berlaku saat Anda melakukan pelatihan menggunakan VM TPU:

• TPU hanya tersedia di wilayah Vertex AI tertentu.

Jenis TPU

Lihat Jenis TPU untuk mengetahui informasi selengkapnya tentang akselerator TPU seperti batas memori.