יצירת פרופיל של חישובי JAX באמצעות XProf

‫XProf היא דרך מצוינת להשיג פרופילים של עקבות ביצועים של התוכנית שלכם ולהציג אותם בצורה חזותית, כולל פעילות ב-GPU וב-TPU. התוצאה הסופית נראית בערך כך:

צילום פרוגרמטי

אפשר להוסיף לקוד כלים לתיעוד של פרופילר עבור קוד JAX באמצעות השיטות jax.profiler.start_trace ו-jax.profiler.stop_trace. שולחים קריאה אל jax.profiler.start_trace עם הספרייה שאליה רוצים לכתוב את קובצי המעקב. זו צריכה להיות אותה ספרייה --logdir שבה השתמשתם כדי להפעיל את XProf. אחר כך אפשר להשתמש ב-XProf כדי להציג את העקבות.

לדוגמה, כדי ליצור מעקב פרופיל:

import jax

jax.profiler.start_trace("/tmp/profile-data")

# 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()

שימו לב לשיחה jax.block_until_ready. אנחנו משתמשים בזה כדי לוודא שהביצוע במכשיר מתועד בנתוני המעקב. בקטע שליחה אסינכרונית מוסבר למה זה נחוץ.

אפשר גם להשתמש במנהל ההקשר jax.profiler.trace כחלופה ל-start_trace ול-stop_trace:

import jax

with jax.profiler.trace("/tmp/profile-data"):
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

צפייה בנתוני המעקב

אחרי שמבצעים מעקב, אפשר לראות אותו בממשק המשתמש של XProf.

אפשר להפעיל את ממשק המשתמש של הכלי ליצירת פרופילים ישירות באמצעות הפקודה העצמאית XProf, על ידי הפניה שלה לספריית היומן:

$ xprof --port=8791 /tmp/profile-data
Attempting to start XProf server:
  Log Directory: /tmp/profile-data
  Port: 8791
  Worker Service Address: 0.0.0.0:50051
  Hide Capture Button: False
XProf at http://localhost:8791/ (Press CTRL+C to quit)

עוברים לכתובת ה-URL שצוינה (לדוגמה, http://localhost:8791/) בדפדפן כדי לראות את הפרופיל.

המעקב הזמין מופיע בתפריט הנפתח 'סשנים' בצד ימין. בוחרים את הסשן שרוצים לבדוק, ואז בתפריט הנפתח 'כלים' בוחרים באפשרות 'כלי לצפייה בנתוני מעקב'. עכשיו אמור להופיע ציר זמן של ההפעלה. אפשר להשתמש במקשי WASD כדי לנווט בנתוני המעקב, וללחוץ או לגרור כדי לבחור אירועים ולקבל פרטים נוספים. במאמרי העזרה של הכלי Trace Viewer מופיעים פרטים נוספים על השימוש בכלי.

צילום ידני באמצעות XProf

בהמשך מפורטות הוראות ליצירת קובץ מעקב של N שניות שהופעל באופן ידני מתוכנית שפועלת.

1. מפעילים שרת XProf:

   xprof --logdir /tmp/profile-data/
   

   אפשר לטעון את XProf בכתובת <http://localhost:8791/>. אפשר לציין יציאה אחרת באמצעות הדגל --port.

2. בתוכנית או בתהליך של Python שרוצים ליצור להם פרופיל, מוסיפים את השורה הבאה איפשהו קרוב להתחלה:

   import jax.profiler
   jax.profiler.start_server(9999)
   

   הפעולה הזו מפעילה את שרת הפרופיל שאליו מתחבר XProf. צריך להפעיל את שרת הפרופילר לפני שממשיכים לשלב הבא. כשמסיימים להשתמש בשרת, אפשר להתקשר אל jax.profiler.stop_server() כדי לכבות אותו.

   אם רוצים ליצור פרופיל של קטע קוד בתוכנית שפועלת לאורך זמן (למשל, לולאת אימון ארוכה), אפשר להוסיף את הקוד הזה בתחילת התוכנית ולהפעיל אותה כרגיל. אם רוצים ליצור פרופיל של תוכנית קצרה (למשל, מיקרו-בנצ'מרק), אפשר להפעיל את שרת הפרופילר במעטפת IPython ולהריץ את התוכנית הקצרה באמצעות %run אחרי שמפעילים את הלכידה בשלב הבא. אפשרות נוספת היא להפעיל את שרת הפרופילר בתחילת התוכנית ולהשתמש ב-time.sleep() כדי שיהיה לכם מספיק זמן להתחיל את הלכידה.

3. פותחים את <http://localhost:8791/> ולוחצים על הלחצן 'יצירת פרופיל' בפינה הימנית העליונה. מזינים את הערך localhost:9999 ככתובת ה-URL של שירות הפרופיל (זו הכתובת של שרת הפרופילר שהפעלתם בשלב הקודם). מזינים את מספר המילישניות שרוצים ליצור פרופיל עבורן ולוחצים על 'CAPTURE' (לכידה).

4. אם הקוד שאתם רוצים ליצור לו פרופיל לא פועל כבר (למשל, אם הפעלתם את שרת הפרופילר במעטפת Python), מריצים אותו בזמן שהלכידה פועלת.

5. אחרי שהלכידה מסתיימת, XProf אמור להתרענן אוטומטית. (לא כל תכונות הפרופיל של XProf מחוברות ל-JAX, ולכן יכול להיות שבהתחלה ייראה כאילו לא נלכד כלום). מימין, בקטע 'כלים', בוחרים באפשרות 'כלי לצפייה בנתוני מעקב'.

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

XProf ו-Tensorboard

‫XProf הוא הכלי הבסיסי שמפעיל את הפונקציונליות של יצירת פרופילים ולכידת עקבות ב-Tensorboard. כל עוד xprof מותקן, הכרטיסייה 'פרופיל' תופיע ב-Tensorboard. השימוש באפשרות הזו זהה להפעלת XProf באופן עצמאי, כל עוד ההפעלה מתבצעת עם אותה ספריית יומנים. ההרשאה הזו כוללת פונקציונליות של צילום פרופיל, ניתוח וצפייה. ‫XProf מחליף את הפונקציונליות של tensorboard_plugin_profile שהומלצה בעבר.

$ tensorboard --logdir=/tmp/profile-data
[...]
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.19.0 at http://localhost:6006/ (Press CTRL+C to quit)

הוספה של אירועי מעקב מותאמים אישית

כברירת מחדל, האירועים בכלי לצפייה בנתוני מעקב הם בעיקר פונקציות פנימיות ברמה נמוכה של JAX. אפשר להוסיף אירועים ופונקציות משלכם באמצעות jax.profiler.TraceAnnotation ו-jax.profiler.annotate_function בקוד.

הגדרת אפשרויות של כלי הפרופיל

השיטה start_trace מקבלת פרמטר אופציונלי profiler_options, שמאפשר שליטה מדויקת בהתנהגות של הכלי ליצירת פרופילים. הפרמטר הזה צריך להיות מופע של jax.profiler.ProfileOptions.

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

import jax

options = jax.profiler.ProfileOptions()
options.python_tracer_level = 0
options.host_tracer_level = 0
jax.profiler.start_trace("/tmp/profile-data", profiler_options=options)

# 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()

אפשרויות כלליות

1. ‫host_tracer_level: הגדרת רמת המעקב לפעילויות בצד המארח.

   ערכים נתמכים:

   • ‫0: משבית את המעקב אחר המארח (CPU) באופן מלא.
   • ‫1: מאפשר מעקב רק אחרי אירועי TraceMe שהוגדרו על ידי המשתמש.
   • ‫2: כולל עקבות ברמה 1 ופרטים ברמה גבוהה על ביצוע התוכנית, כמו פעולות XLA יקרות (ברירת מחדל).
   • ‫3: כולל עקבות ברמה 2, בתוספת פרטים מפורטים יותר על ביצוע התוכנית ברמה נמוכה, כמו פעולות XLA זולות.

2. ‫device_tracer_level: קובעת אם מעקב אחרי המכשיר מופעל.

   ערכים נתמכים:

   • ‫0: השבתה של מעקב אחרי המכשיר.
   • ‫1: הפעלת מעקב אחר המכשיר (ברירת מחדל).

3. ‫python_tracer_level: קובעת אם מעקב Python מופעל.

   ערכים נתמכים:

   • ‫0: השבתה של מעקב אחר קריאות לפונקציות Python (ברירת מחדל).
   • ‫1: הפעלת מעקב אחר Python.

אפשרויות מתקדמות להגדרה אישית

אפשרויות TPU

1. ‫tpu_trace_mode: מציין את המצב של מעקב TPU.

   ערכים נתמכים:

   • ‫TRACE_ONLY_HOST: המשמעות היא שמתבצע מעקב רק אחרי פעילויות בצד המארח (CPU), ולא נאספים נתוני מעקב מהמכשיר (TPU/GPU).
   • ‫TRACE_ONLY_XLA: המשמעות היא שמתבצע מעקב רק אחרי פעולות ברמת XLA במכשיר.
   • ‫TRACE_COMPUTE: מעקב אחר פעולות חישוב במכשיר.
   • ‫TRACE_COMPUTE_AND_SYNC: האפשרות הזו עוקבת אחרי פעולות חישוב ואירועי סנכרון במכשיר.

   אם לא מציינים את tpu_trace_mode, ברירת המחדל של trace_mode היא TRACE_ONLY_XLA.

2. ‫tpu_num_sparse_cores_to_trace: מציין את מספר ליבות ה-TPU הדלילות שרוצים לעקוב אחריהן.

3. ‫tpu_num_sparse_core_tiles_to_trace: מציין את מספר המשבצות בכל ליבה דלילה למעקב ב-TPU.

4. ‫tpu_num_chips_to_profile_per_task: מציין את מספר שבבי ה-TPU ליצירת פרופיל לכל משימה.

אפשרויות GPU

אלה האפשרויות שזמינות לפרופיל GPU:

• ‫gpu_max_callback_api_events: הגדרת המספר המקסימלי של אירועים שנאספים על ידי CUPTI callback API. ברירת המחדל היא 2*1024*1024.
• ‫gpu_max_activity_api_events: מגדיר את המספר המקסימלי של אירועים שנאספים על ידי CUPTI activity API. ברירת המחדל היא 2*1024*1024.
• ‫gpu_max_annotation_strings: הגדרת המספר המקסימלי של מחרוזות הערות שאפשר לאסוף. ברירת המחדל היא 1024*1024.
• ‫gpu_enable_nvtx_tracking: הפעלת מעקב NVTX ב-CUPTI. ברירת המחדל היא False.
• ‫gpu_enable_cupti_activity_graph_trace: הפעלת מעקב אחר תרשים פעילות של CUPTI עבור תרשימי CUDA. ברירת המחדל היא False.
• ‫gpu_pm_sample_counters: מחרוזת של מדדים של ביצועי GPU, מופרדים באמצעות פסיקים, לאיסוף באמצעות תכונת הדגימה של PM ב-CUPTI (לדוגמה, "sm__cycles_active.avg.pct_of_peak_sustained_elapsed"). דגימת PM מושבתת כברירת מחדל. במאמרי העזרה של NVIDIA בנושא CUPTI מפורטים המדדים הזמינים.
• ‫gpu_pm_sample_interval_us: הגדרת מרווח הדגימה במיקרו-שניות לדגימת CUPTI PM. ברירת המחדל היא 500.
• ‫gpu_pm_sample_buffer_size_per_gpu_mb: הגדרת גודל מאגר הזיכרון של המערכת לכל מכשיר במגה-בייט עבור דגימת CUPTI PM. ברירת המחדל היא 64MB. הערך המקסימלי שנתמך הוא 4GB.
• ‫gpu_num_chips_to_profile_per_task: מציין את מספר מכשירי ה-GPU ליצירת פרופיל לכל משימה. אם לא מציינים ערך, מגדירים את הערך ל-0 או מגדירים ערך לא תקין, כל המעבדים הגרפיים הזמינים ינותחו. אפשר להשתמש באפשרות הזו כדי להקטין את הגודל של אוסף נתוני המעקב.
• ‫gpu_dump_graph_node_mapping: אם האפשרות הזו מופעלת, פרטי מיפוי של צומת בתרשים CUDA נשמרים בקובץ המעקב. ברירת המחדל היא False.

לדוגמה:

options = ProfileOptions()
options.advanced_configuration = {"tpu_trace_mode" : "TRACE_ONLY_HOST", "tpu_num_sparse_cores_to_trace" : 2}

הפונקציה מחזירה InvalidArgumentError אם נמצאו מפתחות או ערכי אפשרויות לא מוכרים.