אימונים גמישים עם Pathways

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

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

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

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

השהיה והמשך

בדרך כלל, GKE שולח הודעת קדימות ל-pod של מאיץ, לפני שה-pod נדחק. כברירת מחדל, הסבילות להקדמת תהליכים ב-Pathways מופעלת בכל פריסות הענן, ומשימות של Pathways Accelerator מאזינות להודעות האלה.

כשמתקבלת הודעה על קדימות, המערכת של Pathways קודם קובעת אם אפשר לשחזר את עומס העבודה הנוכחי – כלומר, אם אפשר לשמור ולשחזר את עומס העבודה באופן שקוף. אם כן, המערכת מנסה להשהות בשקיפות את עומס העבודה של ה-ML על ידי כתיבת המצב הנוכחי שלו לאחסון קבוע, כמו Cloud Storage, לפני ש-GKE מוציא את משימות ההאצה. כש-GKE מתזמן מחדש את העבודות שלכם מאוחר יותר, Pathways מחדש את עומס העבודה של ה-ML על ידי קריאה חוזרת של המצב שנשמר.

אם אי אפשר לשחזר את עומס העבודה, Pathways משבית את משימת ההאצה ומעביר את הכשל למשימה שלכם אם מוגדר אימון גמיש. אם לא מוגדרת הדרכה גמישה, ‏ GKE מפעיל מחדש את כל עומס העבודה על סמך מדיניות ההפעלה מחדש של JobSet.

עומסי עבודה אופייניים של למידת מכונה שמוגדרים באמצעות JAX מסתמכים על רכיבי Pathways XLA חסרי מצב שאפשר לשחזר באמצעות תמונת מצב של זיכרון ברוחב פס גבוה (HBM). עומסי עבודה מסוימים של ML, כמו אלה שמוגדרים באמצעות JAX colocated python API, מסתמכים על רכיבי Pathways עם שמירת מצב. אי אפשר לשחזר אותם.

אימון עם גומיות

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

  • זיהוי של כשל: כשמתרחש כשל בחומרה (לדוגמה, קורס תהליך עבודה של TPU), מערכת Pathways מזהה את זה ומודיעה על כך למשימת האימון של המשתמש באמצעות חריגה בפעם הבאה שמתבצעת גישה לנתונים שנמצאים בחומרה הזו. ההתראה הזו לא גורמת לקריסה של עומס העבודה. היא מאפשרת לקוד לטפל בהתראה ולהגדיר מחדש את המשאבים כדי להמשיך את העיבוד או לצאת בצורה תקינה.
  • מטפל בגמישות שהוגדר על ידי המשתמש: קוד המודל של המשתמש צריך להיות מסוגל לטפל בחריגה הזו. זו הסיבה שקוראים לזה 'שחזור ספציפי לדגם'.
    • יצירת תמונת מצב: הגישה הנפוצה ביותר היא לשמור מדי פעם תמונות מצב של מצב המודל. אם מתרחשת שגיאה, אפשר לטעון את התמונה העדכנית ביותר כדי להמשיך את האימון.
    • הגדרה מחדש: סביר להניח שתצטרכו להגדיר מחדש את משימת האימון כדי להתאים אותה למספר הפרוסות הזמינות. לדוגמה, אם אחד מהפלחים מפסיק לפעול, אפשר להקטין את מספר הפלחים הפעילים באחד עד שיהיה פלח חלופי. מידע נוסף מופיע במאמר בנושא Elastic Handler.
    • עדכונים של גרף הנתונים או החישוב: הקוד צריך לטפל בכל שינוי במספר המכשירים שזמינים לחישוב, על ידי יצירה מחדש של גרף החישוב לפי הצורך. יכול להיות שתצטרכו לחלק מחדש את הנתונים או לקמפל מחדש את המודל.
  • התפקיד של Pathways בתהליך השחזור: Pathways מספק את הפרימיטיבים לתמיכה בהגדרה מחדש שהוגדרה על ידי המשתמש:
    • החלפת פרוסה: אם פרוסה שנכשלה מוחלפת, אפשר לעדכן את הלקוח כשהפרוסה החדשה זמינה. לאחר מכן, הקוד יכול להגדיר מחדש את השימוש בפלח החדש.
    • שחזור שקוף: Pathways מטפל בפרטים ברמה הנמוכה של השחזור, כמו חיבור מחדש לחלקים התקינים של האשכול.
  • ‫Utilities in pathwaysutils: קבוצה של כלי Pathways שמוגדרים ב-pathways-utils.

הטמעה של handler גמיש

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

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

def elastic_handler(elastic_utils, *args, **kwargs):
  mesh = initialize_mesh(**kwargs["mesh_kwargs"])
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **kwargs["initialize_training_loop_kwargs"])

  step, snapshot = elastic_utils.get_next_snapshot()
  state = initial_state.replace(**snapshot)

  return state, step, mesh, jitted_train_step, other_variables

עדכון לולאת האימון

צריך לבצע את השינויים הבאים בלולאת האימון:

  1. יצירת חשבון ניהול גמיש
  2. עוטפים את לולאת האימון בבלוקים של try-except שמטפלים ב-jax.errors.JaxRuntimeError
  3. בתוך רכיב ה-jax.errors.JaxRuntimeError handler, קוראים ל-maybe_reshard_down. אם השגיאה קשורה לאירוע גמיש, מנהל הגמישות יבצע שוב את חלוקת הנתונים, אחרת הוא יציג אותה שוב.
  4. הפעלת הפונקציות Call maybe_snapshot ו-maybe_reshard_up בסוף לולאת האימון
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()

def initialize_mesh(**kwargs):
  ...


def initialize_training_loop(**kwargs):
  ...


def train_loop(
    final_step,
    elastic_manager,
    mesh_kwargs,
    initialize_training_loop_kwargs,
):
  mesh = initialize_mesh(**mesh_kwargs)
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **initialize_training_loop_kwargs)

  step = initial_step
  while step < final_step:
    try:
      state = jitted_train_step(state)

      elastic_manager.maybe_snapshot(step=step, snapshot=state)
      handler_returns = elastic_manager.maybe_reshard_up(
          step=step,
          snapshot=state,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns
      step += 1
    except jax.errors.JaxRuntimeError as error:
      handler_returns = elastic_manager.maybe_reshard_down(
          error=error,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns

  return state


def main():
  elastic_manager = manager.Manager(
      devices=jax.devices(),
      snapshot_period=10,
      snapshot_buffer_size=1,
      reshard_check_period=5,
      max_elastic_down_event_count=10,
      max_reshard_retry_count=3,
  )

  train_loop(100, elastic_manager, {}, {})

הגדרת מנהל הגמישות

אפשר להגדיר את המנהל הגמיש בכמה דרכים שונות. התדירות של יצירת תמונת מצב נקבעת לפי תקופת תמונת המצב. התקופה של התמונה המייצגת משפיעה על המספר הממוצע של השלבים שאבדו בגלל אירוע אלסטי. התקופה לבדיקת חלוקה מחדש קובעת באיזו תדירות לולאת האימון תבדוק את זמינות הפרוסות. הפרמטר max_elastic_down_event_count מאפשר להגדיר את מספר האירועים הגמישים שהלולאה של האימון תתמוך בהם בגלל אובדן של נתחי נתונים. המאפיין max_reshard_retry_count מציין את מספר הפעמים שמנהל ה-Elastic צריך לנסות שוב לבצע חלוקה מחדש. האובייקט manager הוא אובייקט יחידני וצריך ליצור אותו רק פעם אחת.

Snapshots

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

הפחתת ה-sharding

אחרי שמתרחשת jax.errors.JaxRuntimeError, התכונה 'תוכניות לימודים' בודקת אם השגיאה נובעת מאירוע גמיש בגלל פרוסה שאבדה. אם כן, הפונקציה תפעיל את elastic handler בלולאה עד שהפעולה תצליח או עד שיגיע למספר המקסימלי של ניסיונות חוזרים. אם השגיאה לא נובעת מאירוע גמיש, השגיאה תופיע שוב. הערכים המוחזרים של ה-handler הגמיש מועברים למתקשר.

הגדלת מספר הרסיסים

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

החלפה חמה

החלפה מהירה (Hot-Swap) היא תכונה ב-API של GKE JobSet שבה עבודה עם עדיפות גבוהה יכולה להשתלט במהירות על משאבים מעבודה עם עדיפות נמוכה יותר, וכך למזער את זמן ההשבתה ולהבטיח התאוששות מהירה יותר.

כשיוצרים JobSet, ‏ GKE מתזמן את עומס העבודה בכמה פרוסות, בהתאם להגדרות של JobSet. אם מתרחש כשל בחומרה באחת או יותר מהפרוסות, הפודים המושפעים מסומנים ככשל. כשמתזמנים מחדש את Jobset, אם בחרתם להשאיר פרוסת זמן פנויה באשכול GKE שאפשר להשתמש בה למשימה בעדיפות נמוכה יותר, מערכת Jobset תמפה מחדש את עומס העבודה של פרוסת הזמן שנכשלה של המשימה בעדיפות גבוהה יותר לפרוסת הזמן הפנויה שבה נעשה שימוש על ידי המשימה בעדיפות נמוכה יותר באותו אשכול GKE. המיפוי מחדש הזה בדרך כלל נמשך פחות מדקה.

אחרי הפעלה מחדש של JobSet, החלפה בזמן ריצה יכולה להתרחש במצבים הבאים:

  1. מצב ברירת מחדל: אם יש פרוסות TPU פנויות במצב המתנה באותו אשכול, מתזמן Kubernetes ייתן עדיפות לתזמון של העבודות שהופעלו מחדש בפרוסות האלה, במקום לחכות לתיקון של הפרוסות שנכשלו. כך אפשר לשחזר את החשבון מהר יותר.
  2. עומסי עבודה הטרוגניים: באשכולות שבהם פועלים כמה עומסי עבודה עם PriorityClass מוגדר של Kubernetes, הפעלה מחדש של JobSet יכולה להפעיל החלפה מהירה. אם ההתאמה של האפליקציה שהופעלה מחדש תואמת למשאבים של אפליקציה עם עדיפות נמוכה יותר, מערכת Kubernetes תבצע קדימה (preemption) של האפליקציה עם העדיפות הנמוכה יותר, ותאפשר לאפליקציה עם העדיפות הגבוהה יותר להתחיל מיד. לדוגמה, אפשר להגדיר את ה-pods של העובדים ב-Pathways עם עדיפויות שונות באמצעות PriorityClass.

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

kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

מחילים את קובץ ה-YAML הזה על אשכול GKE:

kubectl apply -f high-prior-job.yaml

לאחר מכן, מצרפים את מחלקת העדיפות החדשה לעבודת ה-worker של Pathways על ידי הוספת הטקסט הבא ל-podspec של pathways-worker Pod.

priorityClassName: high-prior-job

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