יצירת פרופיל של המודל במכונות וירטואליות של Cloud TPU

פרופילים הם אחד הכלים העיקריים לאופטימיזציה של ביצועי המודלים ב-Cloud TPU. כלי הפרופיל הליבה נקרא XProf, והוא זמין במאגר OpenXLA/XProf ב-GitHub. ‫XProf תומך בפרופילים של כל ה-frameworks שמבוססים על XLA, כולל JAX, ‏ Pytorch XLA ו-Tensorflow/Keras.

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

כדי להשתמש ב-XProf, צריך קודם לצלם פרופיל של המודל. יש שתי דרכים ליצור פרופילים:

  1. צילום פרוגרמטי
  2. צילום לפי דרישה (הידוע גם כצילום ידני)

בשיטה של תיעוד פרוגרמטי, צריך להוסיף הערות לקוד המודל כדי לציין איפה בקוד רוצים לתעד פרופילים. בדרך כלל, יוצרים פרופיל לכמה שלבי אימון, או יוצרים פרופיל לבלוק קוד ספציפי במודל. יש דרכים שונות ללכידת עקבות ב-frameworks השונים JAX, ‏ Pytorch XLA ו-TensorFlow – אפשר להשתמש בהתחלה או בסיום של מעקב מבוסס-API, או במנהל הקשר. במסגרות ברמה גבוהה יותר, כמו MaxText, אפשר להפעיל את איסוף הפרופילים פשוט על ידי הפעלת דגל profiler=xplane כשמריצים את MaxText.

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

כדי להפעיל את יצירת הפרופיל לפי דרישה, עדיין צריך להפעיל את שרת XProf בקוד. הפעלת profiler.start_server תתחיל שרת XProf בעומס העבודה של ה-ML, שיאזין לטריגר של לכידת פרופילים לפי דרישה כדי להתחיל ללכוד פרופילים. אפשר להפעיל צילום פרופיל לפי דרישה דרך ממשק המשתמש של Tensorboard או דרך CLI באמצעות הכלי XProfiler.

במאמרים הבאים מוסבר איך להפעיל לכידה פרוגרמטית ולכידה לפי דרישה עבור מסגרות שונות:

גם כשמבצעים לכידת פרופילים באופן אוטומטי וגם כשמבצעים אותה לפי דרישה, צריך לציין את המיקום שבו רוצים לאחסן את הפרופילים שנלכדו. אפשר לאחסן את הפרופילים בספרייה במכונה הווירטואלית של TPU עם נתיב כמו /profiles/run-name, או להשתמש ב-Cloud Storage עם נתיב כמו gs://bucket-name/run-name/.

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

/profiles/run-name/plugins/profile/session1/

או

gs://bucket-name/run-name/plugins/profile/session1/.

סשן מייצג לכידה אחת של פרופילים במהלך הרצת המודל. לדוגמה, אם אתם מצלמים פרופילים בהרצת אימון משלבים 1 עד 3, ומאוחר יותר מצלמים פרופילים משלבים 8 עד 10, אלה פרופילים של אותה הרצה, אבל הצילום הראשון משלבים 1 עד 3 יהיה session1 והצילום השני משלבים 8 עד 10 יהיה session2.

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

הדמיה של פרופילים באמצעות TensorBoard

כדי להציג פרופילים של XProf ב-Tensorboard, צריך להתקין את הפלאגין tensorboard-plugin-profile.

כדי להציג את הפרופילים שתועדו ב- Google Cloud, מומלץ להשתמש ב-TensorBoard במכונה וירטואלית של TPU או בכלי XProfiler כדי לארח את TensorBoard במכונה וירטואלית או בתרמיל GKE.

אם צילמתם את הפרופילים במכונה וירטואלית של TPU, תוכלו להציג אותם באמצעות TensorBoard באותה מכונה וירטואלית של TPU.

כדי להתקין את XProf ו-TensorBoard במכונת TPU, אפשר להריץ את הפקודות הבאות:

pip install tensorboard_plugin_profile tensorboard
pip install xprof

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

$ tensorboard --logdir=profiles/run-name

או

$ tensorboard --logdir=profiles

כדי לטעון כמה פרופילים מכמה הרצות, צריך להפנות את TensorBoard לספריית הבסיס שמכילה את כל ההרצות והפרופילים, ולא לתת-הספרייה שמכילה פרופילים של הרצה ספציפית.

ספריית Cloud-Diagnostics-XProf להצגת פרופילים של Tensorboard ב Google Cloud

הספרייה cloud-diagnostics-xprof (שנקראת גם XProfiler Library) מקלה על אירוח של TensorBoard ועל הצגת פרופילים באופן חזותי ב- Google Cloud. מידע נוסף זמין במאגר cloud-diagnostics-xprof ב-GitHub.

מאגר cloud-diagnostics-xprof מספק את השיפורים הבאים בהשוואה להרצה של XProf ו-TensorBoard באופן מקומי:

  • הגדרה ואריזה של יחסי תלות ב-XProf וב-TensorBoard
  • אחסון הפרופילים ב-Cloud Storage יכול להיות שימושי לשמירה לטווח ארוך ולניתוח אחרי ההרצה (פרופילים מקומיים שצולמו יימחקו אחרי שהחוקר יסיים את ההרצה)
  • טעינה מהירה של פרופילים גדולים ושל כמה פרופילים על ידי הקצאת TensorBoard במכונה וירטואלית של Compute Engine או בתרמיל GKE, עם אפשרות לשנות את סוג המכונה על סמך צורכי המשתמש במהירות הטעינה ובעלות
  • יצירת קישור לשיתוף פרופילים ולשיתוף פעולה עם חברי הצוות ומהנדסי Google
  • קל יותר ליצור פרופילים של עומסי עבודה ב-GKE וב-Compute Engine לפי דרישה, כדי לבחור כל מארח שמריץ את עומס העבודה וליצור פרופילים

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

כדי להגדיר את ספריית XProfiler, צריך להגדיר את Google Cloud CLI ואת סביבת Python וירטואלית, ואז פשוט מריצים את הפקודה:

pip install cloud-diagnostics-xprof

כל הרכיבים התלויים של XProf ו-TensorBoard יותקנו.

לאחר מכן, כדי ליצור מכונת VM או פוד GKE לאירוח TensorBoard, מריצים את הפקודה הבאה:

xprofiler create -z $ZONE -l $GCS_PATH

או

xprofiler create --GKE -z $ZONE -l $GCS_PATH

מחליפים את $ZONE בכל אזור ואת $GCS_PATH בנתיב של עקבות הפרופיל. אפשר לציין את ספריית הבסיס שמכילה את כל עקבות הפרופיל שלכם לכמה הפעלות, או קבוצה ספציפית של עקבות פרופיל להפעלה אחת.

לדוגמה, אם אתם מתעדים פרופילים בספריות הבאות:

gs://bucket-name/run1/plugins/profile/session1/<profile.xplane.pb>
gs://bucket-name/run1/plugins/profile/session2/<profile.xplane.pb>
gs://bucket-name/run2/plugins/profile/session1/<profile.xplane.pb>

אפשר להגדיר את הנתיב של GCS לספרייה הראשית (gs://bucket-name):

xprofiler create -z $ZONE -l gs://bucket-name/

בממשק המשתמש של TensorBoard, כל הפרופילים של כל ההרצות והסשנים מוצגים כ-run1/session1, run1/session2, run2/session1.

כברירת מחדל, הפקודה xprofiler create תיצור מכונה וירטואלית ב-Compute Engine, ספציפית מכונה וירטואלית מסוג c4-highmem-8. אפשר לשנות את סוג המכונה באמצעות הדגל -m. אם רוצים ליצור מכונה של TensorBoard ב-pod של GKE במקום במכונת VM של Compute Engine, אפשר להעביר את הדגל –GKE אל xprofiler create. אירוח של מופע TensorBoard בתא GKE עשוי להקל על ניהול TensorBoard יחד עם שאר עומס העבודה שנפרס ב-GKE.

המכונה הווירטואלית או ה-pod של GKE שמארחים את TensorBoard מאפשרים טעינה של פרופילים גדולים ושל כמה פרופילים הרבה יותר מהר מאשר TensorBoard שמארח את עצמו באופן מקומי. Google Cloud על סמך הבדיקות שלנו, בפעם הראשונה שתשתמשו במכונת ה-VM c4-highmem-8 שמוגדרת כברירת מחדל, יופיעו פרופילים של טעינה של 1GB תוך כמה דקות. בנוסף, אתם יכולים לבחור את סוג המכונה המתאים על סמך הביצועים והעלויות שאתם צריכים.

אחרי שמריצים את הפקודה xprofiler create, הפלט אמור להיראות כך:

Instance for gs://<bucket> has been created.
You can access it via the following:
1. https://<id>-dot-us-<region>.notebooks.googleusercontent.com.
2. xprofiler connect -z <zone> -l gs://bucket-name -m ssh
Instance is hosted at xprof-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.

האפשרות הראשונה היא קישור שאפשר ללחוץ עליו כדי לראות את פרופילי XProf ב-TensorBoard. הקישור הזה מאפשר לשתף פרופילים עם הצוות שלכם וגם עם מהנדסי Google שעוזרים באופטימיזציה של הביצועים ב- Google Cloud.

אתם קובעים למי תהיה גישה לקישור בהתאם להרשאות שהוגדרו לקטגוריה של Cloud Storage שבה מאוחסנים נתוני הפרופיל. אם הקישור לא פועל מסיבה כלשהי, אפשר גם להתחבר למופע TensorBoard באמצעות SSH כדי להציג את הפרופילים באמצעות הפקודה xprofiler connect.

אם הפעלתם את שרת Cloud Profiler בקוד של עומס העבודה (שנדרש גם לצילום פרופילים פרוגרמטי וגם לצילום פרופילים לפי דרישה) ואתם רוצים לבצע פרופילים לפי דרישה, אתם יכולים לעשות זאת בשתי דרכים:

א. לוחצים על הלחצן Capture profile (צילום פרופיל) ב-TensorBoard. עכשיו אפשר לבחור כל מכשיר מארח שהעומס שלכם פועל עליו כדי ליצור פרופיל. אנחנו תומכים בלכידה לפי דרישה של עומסי עבודה שפועלים ב-GKE או ב-Compute Engine.

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

לפרטים נוספים על השימוש בספרייה cloud-diagnostics-xprof, אפשר לעיין בדף הזה ב-GitHub.