מבוא לתוכניות הלימודים ב-Cloud

‫Pathways היא מערכת שנועדה לאפשר יצירה של מערכות למידת מכונה בקנה מידה גדול, עם ריבוי משימות והפעלה דלילה. הוא מאפשר להשתמש באלפי או בעשרות אלפי מאיצים, עם יכולת להקצות באופן דינמי כמויות שונות של משאבי מחשוב למשימות שונות, בהתאם לדרישות העיבוד שלהן.

‫Pathways מפשט חישובי למידת מכונה בקנה מידה גדול, ומאפשר ללקוח JAX יחיד לתזמר עומסי עבודה על פני כמה פרוסות TPU גדולות, שעשויות לכלול אלפי שבבי TPU.

אנחנו משתמשים ב-Pathways באופן פנימי ב-Google כדי לאמן מודלים גדולים כמו Gemini. הלקוחות של Google Cloud נהנים מאותם יתרונות של Pathways ב-Cloud.

לפני שמתחילים

ודאו שיש לכם:

• כלי Kubernetes מותקנים
• התקנתם את ה-CLI של gcloud
• הפעלתם את Google Kubernetes Engine API

במאמר הזה מוסבר איך להשתמש ביחידות TPU שמנוהלות על ידי Pathways ב-Google Kubernetes Engine ‏ (GKE) לעומסי עבודה באצווה, בזמן אמת ובאופן אינטראקטיבי. ההנחה היא שאתם כבר יודעים איך להשתמש ביחידות TPU עם GKE, כולל יחידות TPU עם פרוסה אחת או כמה פרוסות ב-Google Kubernetes Engine, ושיש לכם ניסיון כללי עם יחידות TPU עם כמה פרוסות.

נאמן מידע יחיד ונאמני מידע מרובים

יש שתי דרכים עיקריות לנהל ולתזמר חישובים בכמה מכשירים:

התכונה	בקר יחיד (Pathways)	Multi-controller (ברירת המחדל של JAX)
בקרה	נקודת בקרה אחת: תוכנת לקוח אחת משמשת כבקר המרכזי.	בקרה מבוזרת: כמה תהליכים משתתפים, ולכל אחד מהם יש מופע משלו של מפרש Python.
תצוגה	תצוגה מאוחדת: הלקוח רואה את כל המכשירים כמערכת אחת מאוחדת.	תצוגה מקומית: כל תהליך Python רואה רק את המכשירים שמחוברים אליו.
תכנות	תכנות פשוט יותר: המשתמשים מתקשרים עם לקוח אחד, כך שהמערכת נראית כמו מכונה אחת גדולה עם הרבה מאיצים מקומיים.	‫SPMD: משתמש בעיקר בפרדיגמת SPMD, שבה כל המכשירים צריכים להריץ את אותה התוכנית.
גמישות	תמיכה בתבניות חישוב מורכבות יותר מעבר ל-SPMD, כולל מקביליות אסימטרית של צינור עיבוד נתונים ודלילות חישובית.	יכול להיות פחות גמיש בניהול המשאבים, במיוחד בפרוסות TPU שונות.

רכיבי Pathways

בקטע הבא מתוארים הרכיבים העיקריים של ארכיטקטורת Pathways.

מנהל המשאבים של Pathways

זהו מישור הבקרה המרכזי של מערכת Pathways. הוא מנהל את כל משאבי המאיצים ואחראי על תיאום ההקצאה של המאיצים למשימות של המשתמשים. הוא עוקב אחרי התקינות של ה-workers ומטפל בתזמון, בהשהיה ובהמשך של משימות. הוא משמש כנקודת קשר אחת לשגיאות ולסטטוס המערכת. הרכיב הזה דורש רק משאבי מעבד (CPU).

לקוח Pathways

זהו יישום של Interim Framework Runtime (IFRT) שמשמש כנקודת הכניסה למערכת Pathways. הוא מקבל מהתוכנית פעולות ברמה גבוהה (HLO). הלקוח של Pathways אחראי על התיאום עם מנהל המשאבים של Pathways כדי לקבוע איפה למקם תוכנות שעברו קומפילציה להפעלה על סמך קוד המשתמש. היא מציגה תצוגה מאוחדת של המערכת ללקוח JAX נתון. הרכיב הזה דורש רק משאבי מעבד (CPU).

עובד ב-Pathways

אלה התהליכים שפועלים במכונות המאיצים (מכונות וירטואליות של TPU). הם מקבלים קובצי הפעלה מקומפלים של התוכנית משרת ה-proxy של IFRT ומבצעים את החישובים ב-TPU. ה-workers של Pathways שולחים נתונים בחזרה לתוכנית דרך שרת ה-proxy של IFRT. הרכיב הזה דורש משאבי מאיץ.

לקוח proxy של IFRT

זהו יישום OSS של ה-API של Interim Framework Runtime (IFRT), שמפריד בין קוד המשתמש לבין זמן הריצה הבסיסי, ומשפר את הניידות והשקיפות של הקוד. ‫JAX משתמש בהטמעה הזו כחלופה לסביבת זמן הריצה שמוגדרת כברירת מחדל עם כמה בקרים. לקוח ה-proxy של IFRT משמש כגשר תקשורת בין התוכנית שלכם לבין רכיבי Pathways. הוא שולח בקשות לשרת ה-proxy של IFRT ומקבל ממנו תוצאות. זהו יישום OSS של IFRT API. הרכיב הזה דורש רק משאבי מעבד (CPU).

שרת proxy של IFRT

שרת ה-gRPC הזה מקבל בקשות מלקוח ה-proxy של IFRT ומעביר אותן ללקוח Pathways, שמטפל בהפצה בפועל של העבודה. הרכיב הזה דורש רק משאבי מעבד (CPU).

שרת Sidecar

שרת ה-gRPC הזה ממוקם יחד עם ה-worker של Pathways ב-VM של המאיץ, כדי להריץ קוד Python שצוין על ידי המשתמש ישירות ב-VM של המאיץ, וכך לצמצם את זמן האחזור של העברת הנתונים מהבקר למאיצים. שרת ה-sidecar יוצר אינטראקציה עם ה-worker של Pathways באמצעות פרוטוקול מותאם אישית עם גרסאות בטרנספורט של gRPC.

רכיבי Pathways

רכיבי Pathways ב-GKE

בקטע הזה ממופים רכיבי Pathways לרכיבי Google Kubernetes Engine כמו קונטיינרים ו-pods.

אפשר למצוא קובצי אימג' של קונטיינרים של Pathways במיקומים הבאים.

סוג המאגר	Location
שרת proxy של IFRT	us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>
מנהל/ת או עובד/ת משאבים ב-Pathways	us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>

מנהל המשאבים של Pathways

אחרי שיוצרים אשכול GKE, אפשר להשתמש ב-containerSpec הבא כדי לפרוס את מנהל המשאבים של Pathways:

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
    imagePullPolicy: Always
    env:
    - name: HOST_ADDRESS
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    - name: TPU_SKIP_MDS_QUERY
      value: "true"
    args:
    - --server_port=29001
    - --node_type=resource_manager
    - --instance_count=WORKLOAD_NODEPOOL_COUNT
    - --instance_type=SLICE_TOPOLOGY
    - --gcs_scratch_location=gs://BUCKET_NAME

תיאורי הארגומנטים:

• ‫--server_port: מנהל המשאבים של Pathways משתמש ביציאה הזו כדי לתקשר עם רכיבים אחרים של Pathways.
• ‫--node_type: סוג הצומת. צריך להגדיר את זה ל-"resource_manager" בשביל מנהל המשאבים של Pathways, ואין צורך בזה בשביל מאגרים אחרים.
• ‫--instance_count: מספר פרוסות ה-TPU.
• ‫--instance_type: סוג ה-TPU והטופולוגיה של הפרוסה. בפורמט tpu{TPU type}:{TPU topology}, לדוגמה tpuv5e:4x4.
• ‫--gcs_scratch_location: קטגוריה של Cloud Storage שמשמשת לקבצים זמניים.

שרת proxy של IFRT

אפשר להשתמש ב-containerSpec הבא כדי לפרוס שרת proxy של IFRT:

 - name: pathways-proxy
    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
    imagePullPolicy: Always
    env:
    - name: PATHWAYS_HEAD
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    args:
    - --resource_manager_address=$(PATHWAYS_HEAD):29001
    - --server_port=29000
    - --gcs_scratch_location=gs://BUCKET_NAME
    ports:
    - containerPort: 29000

תיאורי הארגומנטים:

• ‫--resource_manager_address: שם המארח והיציאה שבהם שרת ה-proxy משתמש כדי לתקשר עם מנהל המשאבים של Pathways. היציאה צריכה להיות זהה לערך של --server_port שמשמש למאגר של מנהל המשאבים של Pathways.
• ‫--server_port: שרת ה-proxy של IFRT משתמש ביציאה הזו כדי לתקשר עם לקוח ה-proxy של IFRT.
• ‫--gcs_scratch_location: קטגוריה של Cloud Storage שמשמשת לקבצים זמניים.

עובד ב-Pathways

אפשר להשתמש ב-containerSpec הבא כדי לפרוס את ה-workers של Pathways:

- name: worker
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
  imagePullPolicy: Always
  env:
  - name: PATHWAYS_HEAD
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
  - name: MEGASCALE_NUM_SLICES
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']"
  - name: MEGASCALE_SLICE_ID
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/job-index']"
  - name: MEGASCALE_COORDINATOR_ADDRESS
    value: "$(PATHWAYS_HEAD)"
  args:
  - --server_port=29001
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --gcs_scratch_location=gs://BUCKET_NAME
  ports:
  - containerPort: 29001
  resources:
    limits:
      google.com/tpu: "4"

תיאורי הארגומנטים:

• ‫--resource_manager_address: שם המארח והיציאה שבהם משתמשים ה-TPU workers כדי לתקשר עם מנהל המשאבים של Pathways. היציאה צריכה להיות זהה לערך --server_port שמשמש את מאגר מנהל המשאבים של Pathways.
• ‫--server_port: ה-workers משתמשים ביציאה הזו כדי לתקשר עם שרת ה-proxy ועם מנהל המשאבים של Pathways.
• ‫--gcs_scratch_location: קטגוריה של Cloud Storage שמשמשת לקבצים זמניים.

למנהל המשאבים של Pathways, לשרת ה-proxy של IFRT ול-workers של Pathways יכולות להיות יציאות שונות, אבל בדוגמה הזו מנהל המשאבים של Pathways וה-worker של Pathways משתפים את אותה יציאה.

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

• יצירת אשכול GKE עם Pathways
• הרצת עומס עבודה באצווה באמצעות Pathways
• הסקה בכמה מארחים באמצעות Pathways
• הרצת עומס עבודה אינטראקטיבי באמצעות Pathways
• אימון עמיד באמצעות Pathways
• העברת עומסי עבודה של JAX ל-Pathways
• פתרון בעיות ב-Pathways בענן
• הרצת עומסי עבודה עם Pathways ב-Cluster Toolkit