Pelatihan dengan akselerator TPU

Agent Platform mendukung pelatihan dengan berbagai framework dan library menggunakan VM TPU. Saat mengonfigurasi resource komputasi, Anda dapat menentukan TPU v2, TPU v3, atau TPU v5e VM. TPU v5e mendukung JAX 0.4.6+, TensorFlow 2.15+, dan PyTorch 2.1+. TPU v6e mendukung Python 3.10+, JAX 0.4.37+, dan PyTorch 2.1+ menggunakan PJRT sebagai runtime default.

Untuk mengetahui detail tentang cara mengonfigurasi VM TPU untuk pelatihan tanpa server, lihat Mengonfigurasi resource komputasi untuk pelatihan tanpa server.

Pelatihan TensorFlow

Catatan: Mulai dari TPU v6e, TensorFlow tidak lagi didukung.

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: Agent Platform hanya mendukung TensorFlow 2.11 dan versi yang lebih baru untuk container kustom serta TensorFlow 2.12 dan versi yang lebih baru untuk container 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.15, 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 tanpa server. 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. Agent Platform 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	Python 3.8
us-docker.pkg.dev/vertex-ai/training/tf-tpu-pod-base-cp310:latest europe-docker.pkg.dev/vertex-ai/training/tf-tpu-pod-base-cp310:latest asia-docker.pkg.dev/vertex-ai/training/tf-tpu-pod-base-cp310:latest	Python 3.10

Berikut adalah langkah-langkah untuk membuat container kustom:

1. Pilih image dasar untuk versi Python pilihan Anda. TPU TensorFlow wheel untuk TensorFlow 2.12 dan yang lebih rendah mendukung Python 3.8. TensorFlow 2.13 dan yang lebih tinggi mendukung Python 3.10 atau yang lebih tinggi. Untuk mengetahui TensorFlow wheel 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-cp310:latest
WORKDIR /root

# 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

# 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.
# Use CMD, not ENTRYPOINT, to avoid accidentally overriding the pod base image's ENTRYPOINT.
# Alternatively, you can pass it as an `args` value in ContainerSpec:
# (https://cloud.google.com/gemini-enterprise-agent-platform/reference/rest/v1/CustomJobSpec#containerspec)
CMD ["python3", "trainer.py"]

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

Buka di Colab | Buka di Colab Enterprise | Buka di Agent Platform Workbench | Lihat di GitHub

Pelatihan PyTorch

Anda dapat menggunakan container bawaan atau kustom untuk PyTorch saat melatih 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, v6e 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).

Agent Platform menunggu respons dari semua host untuk memutuskan penyelesaian tugas.

Pelatihan JAX

Untuk melakukan pelatihan JAX di Agent Platform, Anda harus menyediakan image container dengan lingkungan 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).

Agent Platform 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 tanpa server.

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 Agent Platform tertentu.

Jenis TPU

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