Agent Platform תומכת באימון באמצעות מסגרות וספריות שונות, באמצעות TPU VM. כשמגדירים משאבי מחשוב, אפשר לציין מכונות וירטואליות של TPU v2, TPU v3 או TPU v5e. TPU v5e תומך ב-JAX 0.4.6+, TensorFlow 2.15+ ו-PyTorch 2.1+. TPU v6e תומך ב-Python 3.10+, JAX 0.4.37+ ו-PyTorch 2.1+ באמצעות PJRT כזמן הריצה שמוגדר כברירת מחדל.
לפרטים על הגדרת מכונות וירטואליות של TPU לאימון בלי שרת (serverless), ראו הגדרת משאבי מחשוב לאימון בלי שרת (serverless).
אימון TensorFlow
מאגר מוכן מראש
משתמשים בקונטיינר אימון מוכן מראש שתומך ב-TPU ויוצרים אפליקציית אימון ב-Python.
מאגר תגים בהתאמה אישית
שימוש בקונטיינר בהתאמה אישית שבו מותקנות גרסאות של tensorflow ו-libtpu שנבנו במיוחד עבור מכונות TPU VM. הספריות האלה מתוחזקות על ידי שירות Cloud TPU והן מפורטות במסמכי התיעוד בנושא תצורות TPU נתמכות.
בוחרים את גרסת tensorflow הרצויה ואת הספרייה התואמת libtpu
שלה. לאחר מכן, מתקינים אותם בקובץ האימג' של קונטיינר Docker כשיוצרים את הקונטיינר.
לדוגמה, אם רוצים להשתמש ב-TensorFlow 2.15, צריך לכלול את ההוראות הבאות ב-קובץ Docker:
# 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
TPU Pod
כדי לאמן את tensorflow על TPU Pod צריך לבצע הגדרה נוספת במאגר ההדרכה. Agent Platform מתחזקת תמונת Docker בסיסית שמטפלת בהגדרה הראשונית.
| מזהי URI של תמונות | גרסת Python |
|---|---|
|
Python 3.8 |
|
Python 3.10 |
כדי ליצור מאגר תגים בהתאמה אישית:
- בוחרים את תמונת הבסיס לגרסת Python הרצויה. TPU TensorFlow wheels for TensorFlow 2.12 and lower support Python 3.8. TensorFlow 2.13 ואילך תומך ב-Python 3.10 ואילך. למידע על חבילות TensorFlow Wheels ספציפיות, אפשר לעיין במאמר בנושא הגדרות של Cloud TPU.
- מרחיבים את התמונה באמצעות קוד ההדרכה ופקודת ההפעלה.
# 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"]
אימון ב-PyTorch
כשמבצעים אימון באמצעות TPU, אפשר להשתמש במאגרי תגים מוכנים מראש או במאגרי תגים בהתאמה אישית עבור PyTorch.
מאגר מוכן מראש
משתמשים בקונטיינר אימון מוכן מראש שתומך ב-TPU ויוצרים אפליקציית אימון ב-Python.
מאגר תגים בהתאמה אישית
משתמשים במאגר תגים בהתאמה אישית שבו התקנתם את הספרייה PyTorch.
לדוגמה, קובץ ה-Dockerfile יכול להיראות כך:
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"]
TPU Pod
האימון מופעל בכל המארחים של TPU Pod (ראו הפעלת קוד PyTorch בפרוסות של TPU Pod).
פלטפורמת ה-Agent ממתינה לתשובה מכל המארחים כדי להחליט אם העבודה הושלמה.
אימון ב-JAX
כדי לבצע אימון של JAX ב-Agent Platform, צריך לספק קובץ אימג' של קונטיינר עם סביבת JAX.
מאגר מוכן מראש
אין מאגרי תגים מוכנים מראש ל-JAX.
מאגר תגים בהתאמה אישית
משתמשים במאגר תגים בהתאמה אישית שבו התקנתם את הספרייה JAX.
לדוגמה, קובץ ה-Dockerfile יכול להיראות כך:
# 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"]
TPU Pod
האימון מופעל בכל המארחים של TPU Pod (ראו הרצת קוד JAX בפרוסות של TPU Pod).
Agent Platform עוקבת אחרי המארח הראשון של TPU Pod כדי להחליט אם העבודה הושלמה. אפשר להשתמש בקטע הקוד הבא כדי לוודא שכל המארחים יצאו בו-זמנית:
# 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()
משתני סביבה
בטבלה הבאה מפורטים משתני הסביבה שבהם אפשר להשתמש בקונטיינר:
| שם | ערך |
|---|---|
| TPU_NODE_NAME | my-first-tpu-node |
| TPU_CONFIG | {"project": "tenant-project-xyz", "zone": "us-central1-b", "tpu_node_name": "my-first-tpu-node"} |
חשבון שירות בהתאמה אישית
אפשר להשתמש בחשבון שירות בהתאמה אישית לאימון TPU. איך משתמשים בחשבון שירות בהתאמה אישית
כתובת IP פרטית (קישור בין רשתות VPC שכנות) לאימון
אפשר להשתמש בכתובת IP פרטית לאימון TPU. אפשר לעיין בדף בנושא איך משתמשים בכתובת IP פרטית לאימון בלי שרת.
VPC Service Controls
בפרויקטים שמופעל בהם VPC Service Controls אפשר לשלוח משימות אימון של TPU.
מגבלות
ההגבלות הבאות חלות כשמבצעים אימון באמצעות TPU VM:
סוגי TPU
מידע נוסף על סוגי TPU, כמו מגבלת הזיכרון