העברת עומסי עבודה (workloads) של JAX ל-Pathways

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

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

ודאו שיש לכם:

אינדקס תהליכים

‫JAX עם Pathways מתייחס לכל המכשירים באשכול Pathways כמקומיים. כך ניהול המכשירים פשוט יותר, ו-JAX יכול להשתמש בכל המשאבים הזמינים. בפועל, זה אומר:

  • הערך של jax.process_index() הוא תמיד 0 בכל המכשירים.
  • ‫jax.devices() ו-jax.local_devices() מחזירים את כל מכשירי ה-TPU בכל המשימה.

סוג החומרה ומיקום משותף

כדי להשיג את הביצועים הכי טובים, כדאי למקם את כל רכיבי ה-Pathways ואת משרת המשתמש באותו אזור ענן Google Cloud . להשתמש במעבד גדול כמו IFRT proxy ו-resource manager. מומלץ להשתמש לפחות ב-n2-standard-64 ייעודי, עם 64 ליבות וירטואליות וזיכרון של 256GB.

PathwaysUtils

‫Pathways-utils הוא מאגר GitHub מבוסס-Python שמספק כלי עזר וכלים חיוניים שמאפשרים לייעל את הפריסה וההפעלה של עומסי עבודה של JAX בארכיטקטורה של Pathways on Cloud. החבילה הזו מטפלת בהתאמות הנדרשות לסביבת הענן, ומאפשרת למפתחי JAX להתמקד בתהליכי העבודה העיקריים של למידת מכונה עם מינימום הגדרות ספציפיות לפלטפורמה. באופן ספציפי, הוא מציע:

  • עורף קצה (backend) של JAX מסוג proxy: עורף הקצה המותאם אישית הזה מאפשר לאפליקציית JAX להשתמש בתשתית של Pathways על ידי הגדרת משתנה הסביבה JAX_PLATFORMS=proxy.
  • כלי פרופילים משולבים: יכולות פרופילים שמאפשרות לכם להבין את הביצועים של האפליקציה. באמצעות ממשקי API סטנדרטיים של JAX לפרופילים, כמו jax.profiler.start_trace ו-jax.profiler.start_server, אפשר ליצור פרופילים לא רק של קוד JAX, אלא גם של רכיבי Pathways הבסיסיים, וכך לקבל תמונה הוליסטית של הביצוע בסביבת הענן.
  • ‫Distributed Checkpointing with Orbax: ‏ handler מותאם אישית של Orbax checkpoint שמאפשר לכם להשתמש ב-distributed checkpoints ולשחזר את ה-checkpoints שלכם כשאתם משתמשים בספריית Orbax בסביבת Pathways. המטרה של השילוב הזה היא לפעול בלי לדרוש שינויים בקוד הקיים של Orbax checkpointing, כל עוד הוא מייבא את pathwaysutils.
  • פרימיטיבים גמישים לאימון: מספק פרימיטיבים בסיסיים גמישים לאימון שאפשר להשתמש בהם כדי ליצור תהליכי עבודה חזקים וניתנים להרחבה לאימון באמצעות Pathways. הפרימיטיבים האלה מאפשרים לעבודות ההדרכה שלכם להסתגל באופן דינמי לשינויים במשאבים הזמינים, וכך לשפר את היעילות והעמידות בסביבות ענן.

Checkpointing

Orbax נבדק באופן יסודי עם Pathways לצורך יצירת נקודות ביקורת מבוזרות ושחזור באמצעות Cloud Storage. כשמגדירים את משתנה הסביבה ENABLE_PATHWAYS_PERSISTENCE=1 וקוראים ל-import pathwaysutils; pathwaysutils.initialize() ב-train.py, נרשם ArrayHandler מותאם אישית שמטפל ביעילות בפעולות של נקודות ביקורת דרך ה-proxy של IFRT, וכך מאפשר לעובדי Pathways במאיצים לשמור ולשחזר נתונים ישירות.

Python שמוצב באותו מיקום

‫Colocated Python הוא API של JAX בקוד פתוח שמאפשר להריץ קוד Python שצוין על ידי המשתמש ישירות במארחי TPU או GPU. השימוש בו פשוט יותר ב-JAX עם כמה בקרי TPU. כך אפשר לבצע משימות שדורשות יותר כוח מחשוב, כמו טעינת נתונים ויצירת נקודות ביקורת, כדי להימנע מהעברת נתונים בין הלקוח למכונות TPU. כדי להגדיר את אשכול Pathways להרצת Python JAX API במיקום משותף, פועלים לפי ההוראות שבקובץ ה-README של Python במיקום משותף. בהוראות האלה מוסבר איך להפעיל תהליך Python sidecar שמוצב באותו מיקום עם תהליכי העבודה של Pathways.

טעינת נתונים

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

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

אפשר למצוא הטמעה לדוגמה של העברת צינור קלט שמבוסס על TFDS בהטמעה של RemoteIterator ב-multihost_dataloading.py. ההטמעה הזו פועלת גם ב-JAX עם כמה בקרי Pathways וגם ב-Pathways באופן מבוזר באמצעות Python JAX API שמוצב באותו מיקום.

ניהול גרסאות ב-Jax

הגרסאות של Pathways קשורות באופן הדוק לגרסאות של JAX כדי להבטיח תאימות ויציבות. כדי למנוע בעיות אפשריות, חשוב לוודא שחפצי ה-Pathways וגרסת JAX תואמים. בכל גרסת Pathways מצוינות בבירור גרסאות JAX התואמות באמצעות תג מהצורה jax-<version>.

מטמון הידור

מטמון קומפילציה מתמיד של Pathways הוא תכונה שמאפשרת לשרתי Pathways לאחסן קובצי הפעלה של XLA שעברו קומפילציה במיקום מתמיד, כמו Cloud Storage, כדי למנוע קומפילציה מיותרת. התכונה הזו מופעלת כברירת מחדל. מיקום המטמון מועבר כדגל --gcs_scratch_locationלמאגרי המשאבים ולמאגרי העובדים של Pathways. כדי לצמצם את עלויות האחסון המשויכות, המטמון מצרף מדיניות מחזור חיים למיקום ב-Cloud Storage. יש מגבלה של 50 כללי מדיניות לכל קטגוריה ב-Cloud Storage. לכן, אנחנו ממליצים להשתמש במיקום משותף ב-Cloud Storage בכל עומסי העבודה.

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

המטמון הזה דומה למטמון הקומפילציה של JAX, שמושבת על ידי pathwaysutils.initialize() עבור עומסי עבודה של Pathways.

כדי להשתמש במטמון של קומפילציה, צריך את ההרשאות הבאות ב-Cloud Storage:

  • ‫storage.buckets.get: כדי לאחזר מטא-נתונים של קטגוריה.
  • ‫storage.buckets.update: חיוני ל-Pathways כדי להגדיר מדיניות מחזור חיים של אובייקטים כדי לאכוף TTL לצורך פינוי מהמטמון.
  • ‫storage.objects.list: כדי להציג רשימה של אובייקטים קיימים במטמון בתוך הקטגוריה.
  • ‫storage.objects.create: כדי לכתוב קובצי הפעלה חדשים שעברו קומפילציה למטמון.
  • ‫storage.objects.get: כדי לקרוא קובצי הפעלה ששמורים במטמון מהקטגוריה.

יצירת פרופילים

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

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

בשני המקרים, הפרופילים נכתבים לקטגוריה של Cloud Storage. ייווצרו כמה קובצי מעקב בקטגוריה של Cloud Storage, יכול להיות שמתחת לתיקיות שונות של חותמות זמן, למשל:

  • תהליך Python הראשי שהפעיל את המעקב (בדרך כלל מכונת ה-VM של ה-Notebook): <jax-client-vm-name>.xplane.pb
  • שרת proxy של Pathways IFRT: client.<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • מנהל המשאבים של תוכניות הלימודים: server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • עובדים ב-Pathways: server.*<tpu-node-name>.xplane.pb

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

# verify trace files are present
gsutil ls -l -r gs://BUCKET/PREFIX

# View on tensorboard
tensorboard --logdir=gs://BUCKET/PREFIX

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

  • ‫BUCKET : קטגוריה של Cloud Storage לאחסון קובצי המעקב
  • ‫PREFIX: נתיב בתוך הקטגוריה של Cloud Storage לאחסון קובצי ה-Trace

לכידת פרופיל פרוגרמטית

לכידת פרופיל מתוך הקוד. הפרופילים נשמרים בתיקייה gs://<bucket>/<prefix> בתוך ספרייה עם חותמת זמן

import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_trace("gs://BUCKET/PREFIX")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

לכידת פרופיל באופן ידני

כדי ללכוד באופן ידני את פרטי הפרופיל, צריך להפעיל את שרת הפרופיל מקוד Python:

import jax
import pathwaysutils

pathwaysutils.initialize()
jax.profiler.start_server(jax_profiler_port)

# Your JAX code
jax.profiler.stop_server() # this is functionally a no-op

בזמן ששרת הפרופיל פועל, אפשר לצלם פרופיל ולייצא את הנתונים למיקום היעד ב-Cloud Storage:

export DURATION_IN_SECS=6
export OUTPUT_DIR=gs://<var>BUCKET</var>/<var>PREFIX</var>

curl -d "{\"duration_ms\":\"${DURATION_IN_SECS} * 1000 }}\", \"repository_path\":\"${OUTPUT_DIR}\"}" -H "Content-Type: application/json" -X POST http://localhost:<jax_profiler_port>

אפשר למצוא מידע על התזמון של שיטות של לקוח proxy של IFRT כמו Compile ו-Execute במעקב של התוכנית. האירועים האלה, שמפרטים את האינטראקציות עם שרת ה-gRPC של IFRT Proxy במהלך ההידור וההרצה, מופיעים בשרשור שנקרא GrpcClientSessionUserFuturesWorkQueue. בדיקת השרשור הזה בנתוני המעקב מאפשרת לקבל תובנות לגבי הביצועים של הפעולות האלה.

XLA Flags

כשמשתמשים ב-Pathways, צריך להגדיר את דגלי ה-XLA בpathways-proxycontainer.

כשמשתמשים ב-Cluster Toolkit (הפקודה gcluster job submit), מגדירים את דגלי XLA באמצעות --pathways-proxy-args:

--pathways-proxy-args="--xla_flag_1=value1 --xla_flag_2=value2"

HLO Dump

כדי לבדוק את הקלט של High Level Optimizer ‏ (HLO) שמועבר לקומפיילר XLA, אפשר להגדיר את Pathways כך שיבצע dump של ה-HLO למיקום ספציפי ב-Cloud Storage באמצעות --pathways-proxy-args:

--pathways-proxy-args="--xla_dump_to=gs://BUCKET_NAME/HLO_PREFIX/"

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