העברת עומסי עבודה (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 Cloud. משתמשים במעבד (CPU) גדול כמו שרת ה-Proxy ומנהל המשאבים של IFRT. אנחנו ממליצים על n2-standard-64 ייעודי אחד לפחות, שמגיע עם 64 מעבדי vCPU וזיכרון בנפח 256GB.

PathwaysUtils

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

  • "שרת proxy" לקצה העורפי של JAX: הקצה העורפי המותאם אישית הזה מאפשר לאפליקציית JAX להשתמש בתשתית של Pathways על ידי הגדרת משתנה הסביבה JAX_PLATFORMS=proxy.
  • כלי פרופיילינג משולבים: יכולות פרופיילינג שמאפשרות להבין את הביצועים של האפליקציה. באמצעות ממשקי API סטנדרטיים של JAX לפרופיילינג, כמו jax.profiler.start_trace ו-jax.profiler.start_server, אתם יכולים ליצור פרופיל לא רק לקוד JAX אלא גם לרכיבי Pathways הבסיסיים, וכך לקבל תצוגה הוליסטית של ההפעלה בסביבת הענן.
  • ‫Distributed Checkpointing with Orbax: כלי מותאם אישית לטיפול בנקודות שמירה ב-Orbax, שמאפשר להשתמש בנקודות שמירה מבוזרות ולשחזר את נקודות השמירה כשמשתמשים בספריית Orbax בסביבת Pathways. השילוב הזה אמור לפעול בלי שתצטרכו לשנות את הקוד הקיים של נקודות הבדיקה ב-Orbax, כל עוד הוא מייבא את pathwaysutils.
  • ‫Elastic Training Primitives: מספק פרימיטיבים בסיסיים של אימון גמיש, שבהם אפשר להשתמש כדי ליצור תהליכי עבודה חזקים וניתנים להתאמה של אימון באמצעות Pathways. הפרימיטיבים האלה מאפשרים למשימות האימון להתאים את עצמן באופן דינמי לשינויים במשאבים הזמינים, וכך לשפר את היעילות והעמידות בסביבות הענן.

נקודת ביקורת

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

‫Python באותו מיקום

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

הנתונים נטענים

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

כשמריצים טוען נתונים מרובה-תהליכים (כמו Grain) בקונטיינר של עומס העבודה הראשי ב-GKE, הקונטיינר משתף את הצומת הראשי עם ה-sidecars של pathways-proxy ו-pathways-rm. בגלל הארכיטקטורה המשותפת הזו, צריך להגדיר את הדברים הבאים:

  • הגדרת מספר מפורש של תהליכי עבודה: בתוך קונטיינר, os.cpu_count() מדווח על סך כל המעבדים הווירטואליים של המכונה הווירטואלית המארחת, ולא על מגבלת המעבד של הקונטיינר. צריך להעביר מספר מפורש של קובצי worker (לא יותר ממגבלת המעבד של הקונטיינר) במקום להסיק אותו מ-os.cpu_count().
  • טעינת נפח /dev/shm שמגובה בזיכרון: טועני נתונים מרובי תהליכים משתמשים בזיכרון משותף של POSIX‏ (/dev/shm) לתקשורת בין תהליכים. מוודאים שמוגדר אמצעי אחסון emptyDir (medium: Memory) שמגובה בזיכרון ב-/dev/shm, כדי שהאצוות שאוחזרו מראש לא יחרגו ממגבלת הזיכרון המשותף של מאגר ברירת המחדל.

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

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

ניהול גרסאות של Jax

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

מטמון הידור

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

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

המטמון הזה דומה למטמון ה-JAX compilation שמושבת על ידי pathwaysutils.initialize() בעומסי עבודה של Pathways.

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

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

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

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

  • פרוגרמטי
    • איך לתעד פרופילים באופן פרוגרמטי מקוד JAX
  • באופן ידני
    • תיעוד פרופילים על פי דרישה אחרי הפעלת שרת ה-profiler מקוד ה-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
  • ‫Pathways Resource manager: server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • עובדים ב-Pathways: ‏server.*<tpu-node-name>.xplane.pb

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

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

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

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

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

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

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

לכידת פרופיל ידנית

כדי לתעד באופן ידני את פרטי הפרופיל, צריך להפעיל את שרת ה-profiler מקוד 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://BUCKET/PREFIX/"

curl -X POST http://localhost:JAX_PROFILER_PORT/profiling \
     -H "Content-Type: application/json" \
     -d "{\"duration_ms\": $((DURATION_IN_SECS * 1000)), \"repository_path\": \"${OUTPUT_DIR}\"}"

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

דגלי XLA

כשמשתמשים ב-Pathways, צריך להגדיר את הדגלים של XLA בקונטיינר של pathways-proxy.

כשמשתמשים ב-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 כך שישליך את ה-HLO למיקום מסוים ב-Cloud Storage באמצעות --pathways-proxy-args:

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

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