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 atau TPU V3. Untuk mengetahui detailnya, lihat Mengonfigurasi resource komputasi untuk pelatihan kustom.

Pelatihan TensorFlow

Container bawaan

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

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.

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.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 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
  • 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

Berikut adalah langkah-langkah untuk membuat container kustom:

  1. Pilih image dasar untuk versi Python pilihan Anda. Roda TensorFlow TPU mendukung Python 3.8 untuk versi 2.12 dan yang lebih lama. Mulai 2.13, Python 3.10 akan didukung. Untuk roda TensorFlow tertentu, lihat halaman konfigurasi Cloud TPU.
  2. Perluas gambar dengan kode pelatih dan perintah startup.
# 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"]

Pelatihan PyTorch

Anda dapat menggunakan container kustom atau bawaan 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.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 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.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 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. Tentang cara menggunakan akun layanan kustom, lihat halaman 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 mendukung Kontrol Layanan VPC dapat mengirimkan tugas pelatihan TPU.

Batasan

Keterbatasan berikut berlaku saat Anda melakukan pelatihan menggunakan VM TPU:

Jenis TPU

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