הפעלת צינור עיבוד נתונים עם מעבדי TPU

בדף הזה מוסבר איך להריץ צינור (pipeline) של Apache Beam ב-Dataflow עם TPU. על משימות שמשתמשות ב-TPU חלים חיובים כמו שמפורט בדף התמחור של Dataflow.

מידע נוסף על שימוש ב-TPU עם Dataflow זמין במאמר תמיכה ב-TPU ב-Dataflow.

אופציונלי: הזמנה ספציפית לשימוש במאיצים

אפשר להשתמש ב-TPU על פי דרישה, אבל אנחנו ממליצים מאוד להשתמש ב-TPU של Dataflow עם הזמנות ספציפיות ב-Google Cloud Platform. כך תוכלו לוודא שיש לכם גישה למאיצים זמינים ולזמני הפעלה מהירים של העובדים. לצינורות שצורכים הזמנת TPU לא נדרשת מכסת TPU נוספת.

אם לא מבצעים הזמנה ובוחרים להשתמש ב-TPU על פי דרישה, צריך להקצות מכסת TPU לפני שמריצים את צינור הנתונים.

אופציונלי: הקצאת מכסת TPU

אפשר להשתמש ב-TPU לפי דרישה או באמצעות הזמנה. אם רוצים להשתמש ב-TPU לפי דרישה, צריך להקצות מכסת TPU לפני שמתחילים. אם משתמשים בהזמנה ממוקדת ספציפית, אפשר לדלג על הקטע הזה.

כדי להשתמש ב-TPU על פי דרישה ללא הזמנה, צריך לבדוק את המגבלה ואת השימוש הנוכחי במכסת Compute Engine API ל-TPU באופן הבא:

המסוף

1. נכנסים לדף Quotas במסוף Google Cloud :

   לפתיחת הדף Quotas

2. בתיבה Filter, מבצעים את הפעולות הבאות:

   1. משתמשים בטבלה הבאה כדי לבחור ולהעתיק את המאפיין של המכסה, בהתאם לגרסת ה-TPU ולסוג המכונה. לדוגמה, אם אתם מתכננים ליצור צמתים של TPU v5e על פי דרישה, שסוג המכונה שלהם מתחיל ב-ct5lp-, צריך להזין Name: TPU v5 Lite PodSlice chips.

      גרסת TPU, סוג המכונה מתחיל ב-	המאפיין והשם של המכסה עבור מופעים על פי דרישה
      ‫TPU v5e,‏ ct5lp-	Name: TPU v5 Lite PodSlice chips
      ‫TPU v5p, ct5p-	Name: TPU v5p chips
      ‫TPU v6e,‏ ct6e-	Dimensions (e.g. location): tpu_family:CT6E

   2. בוחרים במאפיין מאפיינים (למשל מיקום) ומזינים region: ואחריו את שם האזור שבו מתכננים להפעיל את צינור הנתונים. לדוגמה, מזינים region:us-west4 אם מתכננים להשתמש באזור us-west4-a. מכסת ה-TPU היא אזורית, ולכן כל האזורים באותו אזור צורכים את אותה מכסת TPU.

הגדרה של קובץ אימג' מותאם אישית של קונטיינר

כדי לבצע אינטראקציה עם TPU בצינורות עיבוד נתונים של Dataflow, צריך לספק תוכנה שיכולה לפעול במכשירי XLA בסביבת זמן הריצה של צינור עיבוד הנתונים. לשם כך צריך להתקין ספריות TPU בהתאם לצרכים של צינור העיבוד ולהגדיר משתני סביבה בהתאם למכשיר ה-TPU שבו משתמשים.

כדי להתאים אישית את קובץ האימג' של הקונטיינר, מתקינים את Apache Beam בקובץ בסיס לאימג' מוכן לשימוש שמכיל את ספריות ה-TPU הנדרשות. לחלופין, אפשר להתקין את תוכנת ה-TPU בתמונות שפורסמו עם מהדורות של Apache Beam SDK.

כדי לספק קובץ אימג' של קונטיינר בהתאמה אישית, משתמשים באפשרות sdk_container_image pipeline. מידע נוסף זמין במאמר בנושא שימוש במאגרי תגים מותאמים אישית ב-Dataflow.

כשמשתמשים במאיץ TPU, צריך להגדיר את משתני הסביבה הבאים בקובץ אימג' של קונטיינר.

ENV TPU_SKIP_MDS_QUERY=1 # Don't query metadata
ENV TPU_HOST_BOUNDS=1,1,1 # There's only one host
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0 # Always 0 for single-host TPUs

בהתאם למאיץ שבו אתם משתמשים, צריך להגדיר גם את המשתנים שבטבלה הבאה.

סוג	טופולוגיה	הגדרות חובה ב-Dataflow worker_machine_type	משתני סביבה נוספים
tpu-v5-lite-podslice	‫1x1	ct5lp-hightpu-1t	TPU_ACCELERATOR_TYPE=v5litepod-1 TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v5-lite-podslice	2x2	ct5lp-hightpu-4t	TPU_ACCELERATOR_TYPE=v5litepod-4 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v5-lite-podslice	2x4	ct5lp-hightpu-8t	TPU_ACCELERATOR_TYPE=v5litepod-8 TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v6e-slice	‫1x1	ct6e-standard-1t	TPU_ACCELERATOR_TYPE=v6e-1 TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v6e-slice	2x2	ct6e-standard-4t	TPU_ACCELERATOR_TYPE=v6e-4 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v6e-slice	2x4	ct6e-standard-8t	TPU_ACCELERATOR_TYPE=v6e-8 TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v5p-slice	2x2x1	ct5p-hightpu-4t	TPU_ACCELERATOR_TYPE=v5p-8 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1

קובץ Docker לדוגמה עבור קובץ אימג' של קונטיינר מותאם אישית יכול להיראות כמו בדוגמה הבאה:

FROM python:3.11-slim

COPY --from=apache/beam_python3.11_sdk:2.66.0 /opt/apache/beam /opt/apache/beam

# Configure the environment to access TPU device

ENV TPU_SKIP_MDS_QUERY=1
ENV TPU_HOST_BOUNDS=1,1,1
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0

# Configure the environment for the chosen accelerator.
# Adjust according to the accelerator you use.
ENV TPU_ACCELERATOR_TYPE=v5litepod-1
ENV TPU_CHIPS_PER_HOST_BOUNDS=1,1,1

# Install TPU software stack.
RUN pip install jax[tpu] apache-beam[gcp]==2.66.0 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

ENTRYPOINT ["/opt/apache/beam/boot"]

הרצת המשימה עם מעבדי TPU

השיקולים להרצת משימת Dataflow עם TPU כוללים את הדברים הבאים:

• מאחר שמאגרי TPU יכולים להיות גדולים, כדי להימנע ממצב שבו נגמר מקום בכונן, צריך להגדיל את גודל דיסק האתחול שמוגדר כברירת מחדל ל-50 ג'יגה-בייט או לגודל מתאים אחר לפי הצורך של קובץ אימג' של קונטיינר, באמצעות --disk_size_gbאפשרות הצינור.
• הגבלת המקביליות הפנימית של העובדים.

מעבדי TPU ומקביליות של עובדים

בהגדרת ברירת המחדל, צינורות עיבוד נתונים ב-Python מפעילים תהליך אחד של Apache Beam SDK לכל ליבת מכונה וירטואלית. לסוגי מכונות TPU יש מספר גדול של ליבות vCPU, אבל רק תהליך אחד יכול לבצע חישובים במכשיר TPU. בנוסף, יכול להיות שתהליך ישמור לעצמו מכשיר TPU למשך משך החיים של התהליך. לכן, כשמריצים צינור Dataflow TPU, צריך להגביל את ההקבלה בתוך העובד. כדי להגביל את ההפעלה המקבילית של העובדים, פועלים לפי ההנחיות הבאות:

• אם בתרחיש לדוגמה שלכם נדרש להריץ מסקנות במודל, צריך להשתמש ב-Beam RunInference API. מידע נוסף זמין במאמר Large Language Model Inference in Beam.
• אם אי אפשר להשתמש ב-API של Beam RunInference, אפשר להשתמש באובייקטים משותפים מרובי-תהליכים של Beam כדי להגביל פעולות מסוימות לתהליך יחיד.
• אם אתם לא יכולים להשתמש בהמלצות הקודמות ומעדיפים להפעיל רק תהליך Python אחד לכל עובד, צריך להגדיר את אפשרות הצינור --experiments=no_use_multiple_sdk_containers.
• אפשר לצמצם עוד יותר את מספר השרשורים באמצעות האפשרות --number_of_worker_harness_threads של צינור העיבוד, אם היא משפרת את הביצועים.

בטבלה הבאה מפורטים משאבי המחשוב הכוללים לכל עובד בכל הגדרת TPU.

סוג ה-TPU	טופולוגיה	סוג מכונה	צ'יפים של TPU	vCPU	‫RAM ‏ (GB)
tpu-v5-lite-podslice	‫1x1	ct5lp-hightpu-1t	1	24	48
tpu-v5-lite-podslice	2x2	ct5lp-hightpu-4t	4	112	192
tpu-v5-lite-podslice	2x4	ct5lp-hightpu-8t	8	224	384
tpu-v6e-slice	‫1x1	ct6e-standard-1t	1	44	176
tpu-v6e-slice	2x2	ct6e-standard-4t	4	180	720
tpu-v6e-slice	2x4	ct6e-standard-8t	8	360	1440
tpu-v5p-slice	2x2x1	ct5p-hightpu-4t	4	208	448

הפעלת צינור עיבוד נתונים עם מעבדי TPU

כדי להריץ משימת Dataflow עם TPU, משתמשים בפקודה הבאה.

python PIPELINE \
  --runner "DataflowRunner" \
  --project "PROJECT" \
  --temp_location "gs://BUCKET/tmp" \
  --region "REGION" \
  --dataflow_service_options "worker_accelerator=type:TPU_TYPE;topology:TPU_TOPOLOGY" \
  --worker_machine_type "MACHINE_TYPE" \
  --disk_size_gb "DISK_SIZE_GB" \
  --sdk_container_image "IMAGE" \
  --number_of_worker_harness_threads NUMBER_OF_THREADS

מחליפים את מה שכתוב בשדות הבאים:

• ‫PIPELINE: קובץ קוד המקור של צינור העיבוד.
• ‫PROJECT: שם הפרויקט ב- Google Cloud .
• ‫BUCKET: הקטגוריה של Cloud Storage.
• ‫REGION: אזור Dataflow, לדוגמה, us-central1.
• ‫TPU_TYPE: סוג TPU נתמך, לדוגמה, tpu-v5-lite-podslice. רשימה מלאה של הסוגים והטופולוגיות זמינה במאמר מאיצי TPU נתמכים.
• ‫TPU_TOPOLOGY: טופולוגיית ה-TPU, לדוגמה, 1x1.
• ‫MACHINE_TYPE: סוג המכונה המתאים, לדוגמה, ct5lp-hightpu-1t.
• ‫DISK_SIZE_GB: הגודל של דיסק האתחול של כל worker VM, לדוגמה, 100.
• ‫IMAGE: הנתיב ב-Artifact Registry לקובץ אימג' של Docker.
• ‫NUMBER_OF_THREADS: אופציונלי. מספר ה-threads של רתמת העובדים.

אימות של משימת Dataflow

כדי לוודא שהעבודה משתמשת במכונות וירטואליות של Worker עם TPU, פועלים לפי השלבים הבאים:

1. במסוף Google Cloud , נכנסים לדף Dataflow‏ > Jobs.

   מעבר לדף Jobs

2. בוחרים משרה.

3. לוחצים על הכרטיסייה מדדי משרות.

4. בקטע Autoscaling, מוודאים שיש לפחות מכונה וירטואלית אחת של Current workers.

5. בחלונית הצדדית פרטי המשרה, בודקים שהתג machine_type מתחיל ב-ct. לדוגמה, ct6e-standard-1t. ההודעה הזו מציינת שימוש ב-TPU.

פתרון בעיות במשימת Dataflow

אם נתקלתם בבעיות בהרצת משימת Dataflow עם TPU, כדאי לעיין במאמר פתרון בעיות במשימת Dataflow TPU.

המאמרים הבאים

• כדאי לנסות את הדוגמאות במדריך למתחילים: הפעלת Dataflow ב-TPU.
• מידע נוסף על תמיכה ב-TPU ב-Dataflow
• מידע נוסף על הסקת מסקנות ממודלים גדולים ב-Beam