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

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

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

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

אפשר להשתמש ב-TPU לפי דרישה, אבל אנחנו ממליצים מאוד להשתמש ב-TPU של Dataflow עם הזמנות שמטרגטות באופן ספציפי Google Cloud. כך תוכלו לוודא שיש לכם גישה למאיצים זמינים ולזמני הפעלה מהירים של העובדים. לצינורות שצורכים הזמנת 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	‫2 על 2	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	‫2 על 2	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

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

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, צריך להגביל את ההרצה המקבילית של העובדים. כדי להגביל את ההפעלה המקבילית של העובדים, פועלים לפי ההנחיות הבאות:

• אם בתרחיש לדוגמה שלכם נדרש להריץ מסקנות על מודל, השתמשו ב-BeamRunInference API. מידע נוסף זמין במאמר Large Language Model Inference in Beam.
• אם אי אפשר להשתמש ב-API של Beam RunInference, אפשר להשתמש באובייקטים משותפים multi-process shared של 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	‫2 על 2	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	‫2 על 2	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 וירטואלית, לדוגמה, 100.
• ‫IMAGE: הנתיב ב-Artifact Registry לקובץ אימג' של Docker.
• ‫NUMBER_OF_THREADS: אופציונלי. מספר ה-threads של ה-worker harness.

אימות של משימת 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