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

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

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

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

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

חשוב לוודא שיש לכם:

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

בקר יחיד ובקרים מרובים

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

התכונה

בקר יחיד (Pathways)

Multi-controller (ברירת המחדל של JAX)

בקרה

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

בקרה מבוזרת: משתתפים כמה תהליכים, שלכל אחד מהם יש מופע משלו של מתורגמן Python.

תצוגה

תצוגה מאוחדת: הלקוח רואה את כל המכשירים כמערכת אחת מאוחדת.

תצוגה מקומית: כל תהליך Python רואה רק את המכשירים שמחוברים אליו.

Programming

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

SPMD: בעיקר משתמשים בפרדיגמת SPMD, שבה כל המכשירים צריכים להריץ את אותה התוכנית.

גמישות

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

יכול להיות שיהיה פחות גמיש בניהול משאבים, במיוחד בין חלקי TPU שונים.

רכיבים של מסלולי למידה

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

כלי לניהול משאבים של מסלולים

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

לקוח Pathways

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

עובד במסלול

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

לקוח proxy של IFRT

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

שרת proxy של IFRT

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

שרת Sidecar

שרת gRPC הזה ממוקם יחד עם תהליך העבודה של Pathways במכונה הווירטואלית של המאיץ, כדי להריץ קוד Python שצוין על ידי המשתמש במכונה הווירטואלית של המאיץ ישירות, וכך להפחית את זמן האחזור של העברת הנתונים מהבקר למאיצים. שרת ה-sidecar מבצע אינטראקציה עם תהליך העבודה של 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>

כלי לניהול משאבים של מסלולים

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

  - 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 שמשמש למאגר התגים של הכלי לניהול משאבי נתיבי המרות.
  • ‫--server_port: שרת ה-proxy של IFRT משתמש ביציאה הזו כדי לתקשר עם לקוח ה-proxy של IFRT.
  • ‫--gcs_scratch_location: קטגוריה של Cloud Storage שמשמשת לקבצים זמניים.

עובד במסלול

אפשר להשתמש ב-containerSpec הבאים כדי לפרוס עובדים ב-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 כדי לתקשר עם מנהל המשאבים של Pathways. היציאה צריכה להיות זהה לערך --server_port שמשמש למאגר התגים של הכלי לניהול משאבים של נתיבי המרה
  • ‫--server_port: העובדים משתמשים ביציאה הזו כדי לתקשר עם שרת ה-proxy ומנהל המשאבים של Pathways.
  • ‫--gcs_scratch_location: קטגוריה של Cloud Storage שמשמשת לקבצים זמניים.

למנהל המשאבים של Pathways, לשרת ה-proxy של IFRT ולעובדי Pathways יכולים להיות יציאות שונות, אבל בדוגמה הזו, מנהל המשאבים של Pathways ועובד Pathways חולקים את אותה יציאה.

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