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

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

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

  • השהיה והמשך: מאפשרת לעומס העבודה להתמודד עם הפרעות מתוכננות, כמו הודעות על הפקעה, בלי צורך בקוד מותאם אישית לטיפול בהפקעה.
  • אימון גמיש: מאפשר לעומס העבודה להתמודד עם כשלים לא מתוכננים בחומרה בלי להקריס את הלקוח, באמצעות לוגיקת שחזור בקוד האימון.
  • נקודות בקרה בכמה רמות: מאפשרות לעומס העבודה לשמור נקודות בקרה מקומיות ומשוכפלות בתדירות גבוהה בין ה-TPU workers באמצעות Python באותו מיקום, עם גיבויים אסינכרוניים ב-Cloud Storage כדי לצמצם את זמן השחזור.
  • החלפה חמה: מאפשרת לעומס עבודה של GKE JobSet בעדיפות גבוהה יותר להשתלט במהירות על פרוסות TPU פנויות או בעדיפות נמוכה יותר, במקרה של כשל בחומרה.

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

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

אם אתם משתמשים באימון גמיש, התקינו את החבילות הבאות בקובץ האימג' של הקונטיינר של עומס העבודה:

  • pathwaysutils
  • ‫Orbax‏ (orbax-checkpoint): >= 0.12.6 (נדרש לתמונות מצב בזיכרון עם ocp.v1.training.pathways.Snapshotter)
  • גרעיניות (grain): חובה אם משתמשים ב-grain.experimental.ElasticIterator

אם אתם משתמשים ב-Multi-Tier Checkpointing, אתם צריכים גם:

  • כדי להגדיר את אשכול GKE ל-Multi-Tier Checkpointing, פועלים לפי ההוראות במאמר אימון מודלים של למידת מכונה בקנה מידה גדול ב-GKE באמצעות Multi-Tier Checkpointing:
    • מפעילים באשכול את התוספים Multi-Tier Checkpointing ו-Cloud Storage FUSE (GcsFuseCsiDriver), יחד עם איחוד זהויות של עומסי עבודה ל-GKE.
    • יוצרים קטגוריה של Cloud Storage עם מרחב שמות היררכי.
    • מקצים את התפקידים roles/storage.objectUser ו-roles/storage.insightsCollectorService בניהול זהויות והרשאות גישה (IAM) לחשבון השירות של Multi-Tier Checkpointing ב-Kubernetes.
    • פורסים משאב בהתאמה אישית של CheckpointConfiguration שמציין את הקטגוריה של Cloud Storage, את בורר הצמתים של ה-TPU ואת גודל הנפח בזיכרון (inMemoryVolumeSize).
  • יוצרים קובץ אימג' של קונטיינר משני (sidecar) ב-Python ומעבירים אותו בדחיפה לאותו מיקום – לדוגמה, באמצעות קובץ Docker לדוגמה של קונטיינר משני (sidecar). בתמונות של עומס העבודה ושל ה-sidecar, צריך להשתמש בגרסאות תואמות של jax ו-jaxlib שתואמות לגרסת Pathways ועומדות בדרישות המינימליות הבאות:
    • גרסת Orbax‏: >= 0.12.4
    • גרסת JAX (כולל Python באותו מיקום): >= 0.10.1

השהיה והמשך

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

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

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

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

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

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

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

  • זיהוי כשלים: כשמתרחש כשל בחומרה או הפקעה (preemption) בפרוסת TPU, ‏Pathways מבטלת את ההקצאה של הפרוסה המושפעת ומעלה את השגיאה jax.errors.JaxRuntimeError בפעם הבאה שעומס העבודה ניגש לנתונים בפרוסה הזו. תהליך הלקוח נשאר פעיל, ולכן הקוד יכול לזהות את השגיאה הזו, להגדיר מחדש את המשאבים ולהמשיך באימון.
  • מצבי שחזור: המחלקה pathwaysutils.elastic.manager.Manager והדקורטור elastic_retry שלה תומכים בשתי אסטרטגיות שחזור:
    • השהיה והמשך: אם מגדירים את minimum_slice_count ל-None או למספר הפרוסות הכולל, האימון מושהה כשפרוסה נכשלת וממשיך אחרי שכל הפרוסות תקינות שוב. מכיוון שמספר הפרוסות לא משתנה בין הניסיונות, עומס העבודה יכול להשתמש מחדש ברשת המכשירים הקיימת ובקובצי ההפעלה המקומפלים.
    • שינוי גודל הרפליקה: אם מגדירים את minimum_slice_count לערך קטן ממספר הפרוסות הכולל, האימון ממשיך בפרוסות הפעילות שנותרו כשפרוסה נכשלת, על ידי בנייה מחדש של רשת המכשירים וקימפול מחדש לטופולוגיה הקטנה יותר. שרשור של מעקב ברקע בודק מדי פעם פרוסות לא פעילות ומתעד פרוסות תקינות חדשות ב-elastic_manager.available_inactive_slices. לאחר מכן, לולאת האימון יכולה לבנות מחדש את רשת המכשירים עם הפרוסות החלופיות, להפעיל מחדש את הטופולוגיה המורחבת ולהציב את המצב הנוכחי ברשת החדשה.
  • שחזור מצב: כדי להמשיך את האימון אחרי אירוע גמיש, עומס העבודה משחזר את הפרמטרים של המודל, את מצב האופטימיזציה ואת מצב האיטרטור של הנתונים באמצעות תמונות מצב בזיכרון. המחלקה Snapshotter ב-orbax.checkpoint מעתיקה באופן אסינכרוני את המצב המשוכפל מ-TPU HBM לזיכרון המעבד (CPU) של pinned_host בכל פרוסה. במהלך השחזור, Snapshotter משחזר את המצב מהרפליקות של המארח ששרדו ומחלק אותו מחדש ישירות לרשת ה-TPU הפעילה בלי לקרוא מהאחסון המרוחק.

הפעלת גמישות בשרת ה-proxy של Pathways

כדי להפעיל קבוצות גמישות של מכשירים ב-Pathways ב-GKE, מוסיפים את הדגל --num_elastic_instances לשדה args של הקונטיינר pathways-proxy במניפסט של JobSet, ומציינים את מספר פרוסות ה-TPU הגמישות בעומס העבודה:

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:latest
  args:
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --server_port=29000
  - --gcs_scratch_location=gs://BUCKET_NAME
  - --num_elastic_instances=NUM_SLICES

מחליפים את מה שכתוב בשדות הבאים:

  • ‫BUCKET_NAME: הקטגוריה של Cloud Storage שמשמשת לקבצים זמניים של Pathways.
  • ‫NUM_SLICES: מספר פרוסות ה-TPU בעומס העבודה.

סינון מכשירים פעילים ויצירת רשת המכשירים

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

import jax
import numpy as np
from pathwaysutils.elastic import manager

elastic_manager = manager.Manager()


def create_mesh(
    elastic_manager: manager.Manager, devices_per_slice: int
) -> jax.sharding.Mesh:
  device_array = np.asarray(elastic_manager.active_devices).reshape(
      (elastic_manager.active_slice_count, devices_per_slice)
  )
  return jax.sharding.Mesh(device_array, axis_names=("replica", "model"))

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

לולאת אימון אלסטית

לולאת אימון גמישה מתאמת בין מעקב אחר פרוסות, שחזור מכשלים ויצירת קובץ Snapshot בזיכרון, כמו שמוצג בדוגמת הקוד הבאה:

  • הפונקציה train יוצרת מכונה אחת של manager.Manager() ו-Orbax Snapshotter מחוץ ל-wrapper של הניסיון החוזר, כדי שהמעקב אחרי הפרוסות ותמונות המצב של זיכרון המארח יישמרו בין הניסיונות.
  • אם מחילים את @elastic_manager.elastic_retry על run_training_loop, הפונקציה מוגדרת להמתין לפרוסות תקינות, להגדיר את jax.default_device לפרוסה פעילה, להשאיר את השרשור של מעקב הרקע פעיל ולנסות שוב ב-jax.errors.JaxRuntimeError.
  • בתחילת ניסיון השחזור, הקוד צריך לבנות מחדש את רשת המכשירים מ-elastic_manager.active_devices ולקרוא ל-snapshotter.load כדי לשחזר את קובץ ה-snapshot האחרון ל-sharding המעודכן של הרשת. מאפשרים למאגרי הנתונים הזמניים של המכשיר מהניסיון שנכשל לצאת מההיקף כדי ש-JAX יוכל לפנות את הזיכרון שלהם.
  • בכל שלב בלולאה, הקוד קורא מדי פעם ל-snapshotter.save כדי לגבות את מצב האימון, ובודק את elastic_manager.available_inactive_slices כדי לזהות פרוסות חלופיות תקינות. כשהפרוסות החלופיות זמינות, הקוד מוסיף אותן ל-elastic_manager.active_slice_indices, בונה מחדש את הרשת, מאתחל מחדש את לולאת האימון ומציב את המצב ברשת המורחבת.

בדוגמה הבאה מוסבר איך להטמיע לולאת אימון גמישה באמצעות שילוב של manager.Manager,‏ elastic_retry ו-Snapshotter בזיכרון של Orbax:

import jax
import orbax.checkpoint as ocp
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()


def train(
    final_step: int,
    snapshot_interval: int,  # in steps
    devices_per_slice: int,
):
  elastic_manager = manager.Manager()
  snapshotter = ocp.v1.training.pathways.Snapshotter(replica_axis_index=0)

  @elastic_manager.elastic_retry(
      minimum_slice_count=1,
      poll_interval=10,  # in seconds
      timeout=300,  # in seconds
      retry_policy=manager.ElasticRetryLimit(max_attempts=10),
  )
  def run_training_loop():
    mesh = create_mesh(elastic_manager, devices_per_slice)
    state, abstract_state, jitted_train_step, data_iterator = (
        initialize_training_loop(mesh)
    )
    step = 0

    # Restore from the latest in-memory snapshot after a slice failure.
    if snapshotter.latest is not None:
      snapshotter.join()
      restored = snapshotter.load({
          "state": abstract_state,
          "step": step,
          "iterator_state": data_iterator.get_state(),
      })
      state = restored["state"]
      step = int(restored["step"]) + 1
      data_iterator.set_state(restored["iterator_state"])

    while step < final_step:
      batch = next(data_iterator)
      state = jitted_train_step(state, batch)

      if step % snapshot_interval == 0:
        snapshotter.save(
            step,
            {
                "state": state,
                "step": step,
                "iterator_state": data_iterator.get_state(),
            },
        )

      # If replacement slices are healthy, reinitialize and place state onto the new mesh.
      if elastic_manager.available_inactive_slices:
        elastic_manager.active_slice_indices |= (
            elastic_manager.available_inactive_slices
        )
        mesh = create_mesh(elastic_manager, devices_per_slice)
        iterator_state = data_iterator.get_state()
        _, abstract_state, jitted_train_step, data_iterator = (
            initialize_training_loop(mesh)
        )
        data_iterator.set_state(iterator_state)
        state_shardings = jax.tree.map(lambda x: x.sharding, abstract_state)
        state = jax.device_put(state, state_shardings)

      step += 1

    return state

  state = run_training_loop()
  snapshotter.join()
  return state

הגדרת מנהל הגמישות וההתנהגות של הניסיון החוזר

מפעילים את manager.Manager(devices=None) פעם אחת לכל תהליך. אם משמיטים את devices, המנהל משתמש ב-jax.devices() ומקבץ מכשירים לפי המאפיין slice_index שלהם.

המכונה Manager חושפת את המאפיינים הבאים כדי לבדוק את הטופולוגיה של האשכול במהלך האימון:

  • ‫slice_to_devices: מיפוי מכל slice_index לרצף של אובייקטים מסוג jax.Device ששייכים לפרוסה הזו.
  • ‫all_slice_indices: frozenset של כל האינדקסים של הפרוסות בהקצאה.
  • ‫active_slice_indices: frozenset של אינדקסי פרוסות שמשתתפים ברשת החישוב הפעילה.
  • ‫active_devices: רשימה של אובייקטים מסוג jax.Device ששייכים לפרוסות הפעילות, מסודרים לפי slice_index.
  • ‫inactive_slice_indices: frozenset של אינדקסי פרוסות בהקצאה שלא משתתפים ברשת החישוב.
  • ‫available_inactive_slices: frozenset של אינדקסים של פרוסות לא פעילות ששרשור המעקב ברקע זיהה כתקינים ומוכנים להצטרף מחדש.
  • ‫total_slice_count ו-active_slice_count: המספר הכולל של הפרוסות שהוקצו ומספר הפרוסות הפעילות.
  • ‫default_device: המכשיר 0 מפרוסה פעילה, ש-elastic_retry מגדיר באופן אוטומטי כמכשיר ברירת המחדל של JAX לכל ניסיון.

כדי להתאים אישית את ההתנהגות של הניסיונות החוזרים וההתאמה לעומס, אפשר להעביר את הארגומנטים הבאים אל @elastic_manager.elastic_retry:

  • ‫minimum_slice_count: מספר הפרוסות התקינות המינימלי שנדרש לפני שמתחילים את ההפעלה הראשונית או ניסיון חוזר. אם המדיניות מוגדרת לערך None (ברירת המחדל), elastic_retry ממתין לכל הפרוסות בהקצאה (total_slice_count), ופועל במצב השהיה והמשך. אם מגדירים מספר שלם קטן יותר (כמו 1), elastic_retry פועל במצב שינוי גודל של רפליקה.
  • ‫poll_interval: המרווח בשניות בין בדיקות התקינות של הפרוסות שמנהל האימון הגמיש (Manager) מבצע במהלך האימון. ברירת המחדל היא 10.
  • ‫timeout: מספר השניות המקסימלי להמתנה עד שפרוסות minimum_slice_count יהפכו לפעילות לפני העלאת TimeoutError. ערך ברירת המחדל הוא None (אין זמן קצוב לתפוגה).
  • ‫pre_callback: קריאה חוזרת (callback) אופציונלית שפועלת בתחילת כל ניסיון אחרי שמנהל המערכת מעדכן את active_slice_indices ולפני שפונקציית לולאת האימון העטופה פועלת.
  • ‫on_elastic_event_callback: פונקציית קריאה חוזרת (callback) אופציונלית שפועלת אחרי ש-elastic_retry תופסת את jax.errors.JaxRuntimeError, לפני שמתחיל הניסיון הבא. אפשר להשתמש בקריאה החוזרת הזו כדי לנקות נקודות ביקורת לא שלמות או לאפס איטרטורים של נתונים.
  • ‫retry_policy: פונקציה שניתן להפעיל ((attempt: int, error: Exception) -> bool) שקובעת אם להמשיך לנסות אחרי אירוע גמיש. אפשר להשתמש ב-manager.ElasticRetryLimit(max_attempts) כדי להגביל את מספר הניסיונות הכולל. אם משמיטים את הפרמטר הזה, elastic_retry מנסה שוב ללא הגבלה.

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

תמונות מצב בזיכרון באמצעות Orbax

המחלקה ocp.v1.training.pathways.Snapshotter (זמינה ב-orbax-checkpoint מגרסה 0.12.6 ואילך) שומרת גיבוי בזיכרון של מצב האימון בזיכרון pinned_host בפרוסות:

  • ‫Snapshotter(replica_axis_index=0): מאתחלת את ה-snapshotter ומציינת איזה ציר של הרשת מייצג רפליקות של פרוסות מקבילות של נתונים. כל מערכי ה-JAX במצב השמור חייבים להיות משוכפלים לאורך ציר הרשת הזה, כדי שהפרוסות ששרדו ישמרו עותק מלא של המודל ושל מצב האופטימיזציה כשפרוסה נכשלת. אם המצב מכיל מערכים של מכשיר יחיד, צריך לשכפל אותם ברשת לפני שקוראים ל-save.
  • ‫save(step, state): העברה אסינכרונית של state לזיכרון של pinned_host ופיצול שלו לעצים לכל רפליקה לאורך replica_axis_index. אם העברת קובץ snapshot קודמת עדיין מתבצעת, save מדלג על הבקשה החדשה כדי למנוע עיכוב בלולאת האימון.
  • ‫load(abstract_state, reset_snapshot_state=True): בודק את הרפליקות של המארח שנשמרו כדי לזהות את הפרוסות ששרדו, משרשר את הרפליקות התקינות ומחלק מחדש את המצב לפרוסות היעד שצוינו על ידי abstract_state.
  • ‫join(): חסימה עד להשלמת כל פעולת קובץ snapshot בתור.
  • ‫latest: מחזירה מטא-נתונים (כולל step) של קובץ ה-snapshot האחרון שהושלם, או None אם לא הושלם אף קובץ snapshot.

הפחתת הפיצול במקרה של כשל בפרוסה

אם פרוסת TPU נכשלת או מופקעת במהלך הביצוע, Pathways מעלה jax.errors.JaxRuntimeError (כמו DATA_LOSS,‏ DEADLINE_EXCEEDED,‏ NOT_FOUND או INTERNAL). בפונקציה שמעוטרת ב-@elastic_manager.elastic_retry,‏ elastic_retry תופס את החריגה הזו ובודק את elastic.is_error_due_to_slice_down(error):

  • אם השגיאה לא נגרמה בגלל אירוע של צמצום, elastic_retry מעלה מחדש את החריגה באופן מיידי.
  • אם השגיאה נגרמת בגלל אירוע של הפחתת פרוסות, elastic_retry מפסיק את כל המעקבים הפעילים של JAX Profiler, מריץ את on_elastic_event_callback ומעריך את retry_policy(attempt, error).
  • לפני שמתחילים את הניסיון הבא, elastic_retry קורא ל-elastic.wait_for_slices פעם אחת כדי לדגום את תקינות הפרוסות עד שלפחות minimum_slice_count פרוסות פעילות. אחר כך הוא מעדכן את elastic_manager.active_slice_indices ומריץ את פונקציית לולאת האימון בפרוסות ששרדו.

הגדלת הפיצול כשפרוסות מתחברות מחדש

כשעומס העבודה פועל על מספר מופחת של פרוסות במצב שינוי גודל של רפליקה, השרשור של מעקב הרקע מריץ בדיקות תקינות תקופתיות ב-elastic_manager.inactive_slice_indices. כשהפרוסות שתוקנו או תוזמנו מחדש עוברות את בדיקת התקינות, ה-thread של המעקב מוסיף את האינדקסים של הפרוסות ל-elastic_manager.available_inactive_slices.

כשהערך של elastic_manager.available_inactive_slices לא ריק בסוף שלב אימון, עומס העבודה גדל בחזרה בלי לאבד את ההתקדמות באימון:

  • הוספת האינדקסים של הפרוסות ב-elastic_manager.available_inactive_slices אל elastic_manager.active_slice_indices משלבת את הפרוסות שחוברו מחדש ב-elastic_manager.active_devices ומנקה אותן מ-available_inactive_slices.
  • כדי להתאים את החישוב לטופולוגיה המורחבת, צריך לבנות מחדש את רשת המכשירים מ-elastic_manager.active_devices ולהפעיל מחדש את שלב האימון ואת האיטרטור של הנתונים.
  • קריאה ל-jax.device_put מעבירה את מצב האימון הנוכחי למקטעי הרשת החדשים בלי לשחזר מתמונת מצב.

Multi-Tier Checkpointing

משימות אימון בקנה מידה גדול נתקלות בהפרעות תכופות, וכתיבה או שחזור של נקודות ביקורת ישירות אל Cloud Storage וממנו עלולים לעכב את התקדמות האימון. Multi-Tier Checkpointing מפחית את התקורה של נקודות הבדיקה ואת זמן השחזור על ידי תיאום של שלוש שכבות אחסון ב-GKE:

  • דיסק RAM מקומי: ה-TPU workers כותבים נקודות ביקורת בתדירות גבוהה לנפח CSI בזיכרון (multitier-checkpoint.csi.storage.gke.io) שמותקן בכל צומת עובד. אמצעי האחסון הזה נשמר גם אם מפעילים מחדש את ה-Pod באותו צומת.
  • רפליקציה בין עמיתים באשכול: שירות הרפליקציה של GKE משכפל את חלקי נקודות הבדיקה המקומיות בין עובדי TPU עמיתים באשכול.
  • גיבוי ב-Cloud Storage: שירות השכפול מגבה באופן אסינכרוני נקודות ביקורת מקומיות לקטגוריה של Cloud Storage שמופעל בה מרחב שמות היררכי, בלי לחסום את שלבי האימון של ה-TPU.

בניגוד ל-JAX עם כמה בקרים, שבו כל מארח TPU מריץ סקריפט אימון משלו, ב-Pathways מופעל בקר יחיד ב-Pod של pathways-head והעבודה נשלחת ל-Pods של pathways-worker. כדי לקרוא ולכתוב קטעי נקודות ביקורת ישירות בדיסק ה-RAM המקומי של כל מארח TPU, בלי לנתב את נתוני נקודות הביקורת דרך הבקר היחיד, Pathways משתמש ב-Python באותו מיקום. כל Pod של worker מריץ קונטיינר של colocated-python-sidecar לצד הקונטיינר של pathways-worker, ושני הקונטיינרים טוענים את הנפח של Multi-Tier Checkpointing CSI וספריית זיכרון משותפת (/tmp/sidecar). במהלך האתחול והאימון, Orbax שולח את הגדרת נקודת הבקרה, שומר ומשחזר מבקר Pathways ל-sidecars של ה-worker בכל מארח TPU. ב-Pathways, נקודות הבדיקה בכמה רמות תומכות גם באימון גמיש כשמגדילים או מקטינים את מספר פרוסות ה-TPU הפעילות.

‫Pathways תומך ב-Multi-Tier Checkpointing עם MaxText. הוראות להגדרה ולתצורה מפורטות במדריך לנקודות בקרה בכמה רמות ב-MaxText.

הגדרת עומס העבודה של Pathways

כדי להגדיר עומס עבודה של Pathways ל-Multi-Tier Checkpointing, אתם יכולים להשתמש ב-Cluster Toolkit (gcluster job submit) או להגדיר מניפסט JobSet בהתאמה אישית של Pathways.

שימוש ב-Cluster Toolkit

כששולחים משימה באמצעות gcluster job submit --pathways, מעבירים את הדגלים --gke-mtc-enabled,‏ --gke-mtc-ramdisk-dir ו---pathways-colocated-python-sidecar-image:

gcluster job submit \
    --pathways \
    --gke-mtc-enabled \
    --gke-mtc-ramdisk-dir="/tmp/mtc_checkpoints" \
    --pathways-colocated-python-sidecar-image="SIDECAR_IMAGE" \
    ...

מחליפים את SIDECAR_IMAGE ב-URI של קובץ האימג' של הקונטיינר המשני (sidecar) של Python שנמצא באותו מיקום.

כשמגדירים את הדגלים האלה, gcluster מגדיר באופן אוטומטי את הפרטים הבאים:

  • מעביר את --sidecar_name=external למאגר pathways-proxy.
  • מריץ את מאגר הנתונים colocated-python-sidecar בכל Pod של worker ומעלה ספריית /tmp/sidecar משותפת (--cloud_pathways_sidecar_shm_directory=/tmp/sidecar) גם ב-pathways-worker וגם ב-colocated-python-sidecar.
  • הפקודה הזו טוענת את אמצעי האחסון של multitier-checkpoint.csi.storage.gke.io CSI בנתיב שצוין על ידי --gke-mtc-ramdisk-dir גם ב-pathways-worker וגם ב-colocated-python-sidecar בפודים של worker, אבל לא בפוד של pathways-head.

שימוש במניפסט מותאם אישית של Pathways JobSet

אם אתם מנהלים מניפסט משלכם של Pathways JobSet במקום להשתמש ב-gcluster, מעבירים את --sidecar_name=external אל pathways-proxy. במפרט של Pod worker, טוענים את נפח הזיכרון המשותף של ה-sidecar ואת נפח ה-CSI multitier-checkpoint.csi.storage.gke.io בקונטיינרים pathways-worker ו-colocated-python-sidecar:

containers:
- name: pathways-worker
  args:
  - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
- name: colocated-python-sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
volumes:
- name: shared-tmp
  emptyDir:
    medium: Memory
- name: checkpoint-mount
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

שימוש ב-Orbax Multi-Tier Checkpointing בעומס עבודה מותאם אישית של JAX

בסקריפט אימון מותאם אישית של JAX שפועל ב-Pathways, משתמשים בחבילה orbax.checkpoint.experimental.emergency.multi_tier_checkpointing ומגדירים את use_colocated_python=True גם ב-initialize_multi_tier_checkpointing וגם ב-ReplicatorCheckpointManagerOptions:

from etils import epath
import jax
import orbax.checkpoint as ocp
from orbax.checkpoint.experimental.emergency.multi_tier_checkpointing import (
    initialization,
    replicator_checkpoint_manager,
)
import pathwaysutils

pathwaysutils.initialize()

local_checkpoint_dir = epath.Path("/tmp/mtc_checkpoints")

# Initialize multi-tier checkpointing across worker sidecars.
# When using elastic training, pass the active slice devices to `devices`
# on initial setup and whenever the active slice count changes.
initialization.initialize_multi_tier_checkpointing(
    local_checkpoint_directory=local_checkpoint_dir,
    backup_interval_minutes=10,
    num_slices=num_slices,
    run_name=run_name,
    data_parallelism=data_parallelism,
    use_colocated_python=True,
    devices=jax.devices(),
)

options = replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions(
    save_interval_steps=20,
    use_colocated_python=True,
)
checkpoint_manager = (
    replicator_checkpoint_manager.ReplicatorCheckpointManager(
        local_checkpoint_dir,
        options=options,
        global_mesh=mesh,
    )
)

# Restore from the latest checkpoint if one exists.
start_step = 1
latest_step = checkpoint_manager.latest_step()
if latest_step is not None:
  restore_args = jax.tree.map(
      lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
      state,
  )
  restored = checkpoint_manager.restore(
      latest_step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeRestore(
              item=state,
              restore_args=restore_args,
          )
      ),
  )
  state = restored.state
  start_step = latest_step + 1

# Step 0 is reserved by the replicator service to indicate "no checkpoint".
for step in range(start_step, final_step + 1):
  state = jitted_train_step(state)
  checkpoint_manager.save(
      step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeSave(item=state)
      ),
  )

checkpoint_manager.wait_until_finished()
checkpoint_manager.close()

החלפה חמה

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

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

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

  • מצב ברירת המחדל: אם יש פרוסות TPU פנויות באותו אשכול, מתזמן Kubernetes מתעדף את תזמון המשימות שהופעלו מחדש לפרוסות האלה, במקום להמתין לתיקון הפרוסות שנכשלו. התזמון הזה מאפשר התאוששות מהירה יותר.
  • עומסי עבודה הטרוגניים: באשכולות שבהם פועלים כמה עומסי עבודה עם משאב PriorityClass מוגדר ב-Kubernetes, הפעלה מחדש של JobSet יכולה להפעיל החלפה חמה. אם ה-affinity של המשימה שהופעלה מחדש תואם למשאבים של משימה בעדיפות נמוכה יותר, Kubernetes מפקיע את המשימה בעדיפות הנמוכה יותר כדי שהמשימה בעדיפות הגבוהה יותר תוכל להתחיל בלי להמתין לתיקון הפרוסה. לדוגמה, אתם יכולים להגדיר את ה-Pods של ה-worker ב-Pathways עם עדיפויות שונות באמצעות משאב PriorityClass.

כדי להשתמש בעדיפויות באשכול, מגדירים מניפסט של PriorityClass ושומרים אותו בתור high-prior-job.yaml:

apiVersion: scheduling.k8s.io/v1
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

בשלב הבא, מצמידים את מחלקת העדיפות החדשה לעומס העבודה של Pathways. אם אתם שולחים את עומס העבודה באמצעות Cluster Toolkit, העבירו את --priority=high-prior-job לפקודה gcluster job submit, שמגדירה את priorityClassName גם ב-Pods של pathways-head וגם ב-Pods של worker. אם אתם כותבים מניפסט JobSet משלכם, הגדירו את priorityClassName בתבנית ה-Pod של המשימה המשוכפלת worker, שמריצה את הקונטיינר pathways-worker:

replicatedJobs:
- name: worker
  template:
    spec:
      template:
        spec:
          priorityClassName: high-prior-job

אם יצרתם את האשכול באמצעות תוכנית TPU של Cluster Toolkit עם Pathways מופעל, האשכול כבר כולל את מחלקות העדיפות very-low,‏ low,‏ medium,‏ high ו-very-high, כך שאתם יכולים להעביר אחת מהן אל --priority במקום ליצור מחלקה משלכם.

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