מבוא לתוכניות הלימודים ב-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 VMs). הם מקבלים קובצי הפעלה מהודרים של התוכנית משרת ה-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.
רכיבי תוכניות הלימודים

PathwaysJob API

‫PathwaysJob API הוא API מקורי של Kubernetes ב-OSS, שמשמש לפריסת עומסי עבודה של אימון ML והסקת מסקנות באצווה. הבקר של PathwaysJob משתמש ב-JobSet API כדי לנהל את מחזור החיים ואת התיאום של כל רכיבי Pathways. ההגדרה המותאמת אישית של משאבים (CRD) מספקת ממשק ברמה גבוהה להגדרת עומסי העבודה של Pathways, ומבטלת את הצורך לנהל ישירות מפרטים של Pod בנסיבות נפוצות. רשימה מקיפה של כל הפרמטרים והמשמעויות הספציפיות שלהם מופיעה במאמרי העזרה של ה-API של PathwaysJob ב-GitHub.

apiVersion: pathways-job.pathways.domain/v1
kind: PathwaysJob
metadata:
  name: pathways-USER
spec:
  maxRestarts: MAX_RESTARTS
  pathwaysVersion: jax-JAX_VERSION
  workers:
  - type: $(TPU_MACHINE_TYPE)
    topology: $(TOPOLOGY)
    numSlices: $(WORKLOAD_NODEPOOL_COUNT)
    maxSliceRestarts:  # Optional
  customComponents: # This section is completely optional
  - componentType: proxy_server
    image: CUSTOM_PROXY_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: pathways_server
    image: CUSTOM_PATHWAYS_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: worker
    image: CUSTOM_WORKER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: colocated_python_sidecar
    image: CUSTOM_SIDECAR_IMAGE
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  pathwaysDir: "gs://BUCKET_NAME" # Pre-create this bucket.
  controller:
    deploymentMode: default # Default mode deploys pathways cpu resources (resource
# manager and proxy server) on a dedicated CPU node, recommended for training
    elasticSlices: ELASTIC_SLICES
    template:
      spec:
        containers:
        - name: main
          image: python:3.11
        command:
          - bash
          - -c
          - |
            pip install pathwaysutils
            python3 -c 'import pathwaysutils; import jax; pathwaysutils.initialize(); print(jax.devices())'

בטבלה הבאה מפורטות ההגדרות של PathwaysJob API:

מאפיין תיאור
apiVersion מציינת את גרסת ה-API של PathwaysJob Custom Resource Definition (CRD): pathways-job.pathways.domain/v1.
kind מזהה את אובייקט Kubernetes כ-PathwaysJob.
metadata.name השם של אובייקט PathwaysJob ב-Kubernetes, בדרך כלל לפי התבנית pathways-.
spec ההגדרה הזו מגדירה את המצב והתצורה הרצויים של PathwaysJob.
spec.maxRestarts המספר המקסימלי של פעמים שבהן המערכת יכולה להפעיל מחדש את PathwaysJob באופן אוטומטי אם היא נתקלת בכשלים.
spec.pathwaysVersion (אופציונלי) מציינים את הגרסה הרצויה של ה-framework של JAX לשימוש בסביבת Pathways עבור המשימה הזו (לדוגמה, jax-0.5.3).
spec.workers מערך שמגדיר את התצורה של מאגר העובדים של PathwaysJob, בדרך כלל באמצעות משאבי TPU.
spec.workers[].type סוג מכונת ה-TPU שבה רוצים להשתמש לצמתי העובדים (לדוגמה, $TPU_MACHINE_TYPE יכול להיות ct6e-standard-4t)
spec.workers[].topology הטופולוגיה של פרוסות ה-TPU שהוקצו לעובדים (לדוגמה, $TOPOLOGY יכול להיות 2x2, ‏ 4x4, ‏ 2x2x2).
spec.workers[].numSlices מספר פרוסות ה-TPU להקצאה למאגר העובדים (לדוגמה, הערך של $WORKLOAD_NODEPOOL_COUNT יכול להיות 2).
spec.workers[].maxSliceRestarts (אופציונלי) המספר המקסימלי של פעמים שאפשר להפעיל מחדש עובד בודד בתוך פרוסה אם הוא נכשל.
spec.customComponents (אופציונלי) מערך שמאפשר להגדיר ולפרוס רכיבים בהתאמה אישית (כמו שרתי proxy, שרתי Pathways או עובדים נוספים) לצד המשימה הראשית.
spec.customComponents[].componentType מציין את הסוג של הרכיב המותאם אישית שמוגדר (לדוגמה, proxy_server, ‏ pathways_server, ‏ worker, ‏ colocated_python_sidecar).
spec.customComponents[].image קובץ האימג' של Docker שבו ישתמש הקונטיינר של הרכיב המותאם אישית הזה.
spec.customComponents[].customFlags מערך של דגלים מותאמים אישית של שורת הפקודה שיועברו לקונטיינר כשהוא יתחיל לפעול.
spec.customComponents[].customEnv מערך של משתני סביבה מותאמים אישית שיוגדרו בקונטיינר. לכל רכיב יש שם וערך.
spec.pathwaysDir קטגוריית Cloud Storage שמשמשת את PathwaysJob לאחסון של ארטיפקטים של קומפילציה ונתונים זמניים אחרים. צריך ליצור את הקטגוריה הזו לפני שמריצים את עומס העבודה.
spec.controller הגדרות התצורה של בקרת הנתיבים, שמנהלת את הביצוע הכולל של העבודה.
spec.controller.deploymentMode ההגדרה הזו מציינת איך משאבי ה-CPU של בקר Pathways (מנהל משאבי Pathways ושרת Proxy) נפרסים. במצב ברירת המחדל, המודלים האלה נפרסים בצומת CPU ייעודי, ואילו במצב colocate_head_with_workers הם נפרסים לצד עובד TPU.
spec.controller.elasticSlices (אופציונלי) מספר חלקי ה-TPU המקסימלי שיכולים להיות לא זמינים במהלך ההפעלה של המשימה לפני שהיא תיחשב כלא תקינה.
spec.controller.template (אופציונלי) מגדיר את תבנית ה-Pod לעבודת המשתמש. הפעולה הזו נדרשת לעומסי עבודה באצווה, אבל לא לעומסי עבודה אינטראקטיביים.
spec.controller.template.spec המפרט של ה-Pod למשימת המשתמש.
spec.controller.template.spec.containers מערך שמגדיר את הקונטיינרים שיפעלו בעבודת המשתמש
spec.controller.template.spec.containers[].name השם של הקונטיינר בתוך עבודת המשתמש (בדוגמה הזו, זה main).
spec.controller.template.spec.containers[].image קובץ האימג' של Docker שבו יש להשתמש עבור הקונטיינר בקונטיינר הראשי (בדוגמה הזו, python:3.11).
spec.controller.template.spec.containers[].command הפקודה שתופעל כשהקונטיינר הראשי יתחיל. בדוגמה הזו, הספרייה `pathwaysutils` מותקנת, Pathways מאותחל ומודפסים מכשירי JAX.

רכיבי 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>

מנהל משאבים/עובד של נתיבי למידה

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 חולקים את אותה יציאה.

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