יצירת פרופיל של המודל במכונות וירטואליות של Cloud TPU
יצירת פרופילים היא אחד הכלים העיקריים לאופטימיזציה של ביצועי המודלים ב-Cloud TPU. כלי הפרופיילינג העיקרי נקרא XProf, והוא זמין במאגר OpenXLA/XProf ב-GitHub. ב-XProf יש תמיכה בפרופיילינג של כל ה-frameworks שמבוססים על XLA, כולל JAX, PyTorch ו-Tensorflow/Keras.
צילום פרופילים
כדי להשתמש ב-XProf, קודם צריך לתעד פרופיל של המודל. יש שתי דרכים לתעד פרופילים:
- צילום פרוגרמטי
- צילום על פי דרישה (נקרא גם צילום ידני)
כדי להשתמש בתיעוד פרוגרמטי, צריך להוסיף הערות לקוד המודל כדי לציין איפה בקוד רוצים לתעד פרופילים. בדרך כלל, מתעדים פרופיל לכמה שלבי אימון, או יוצרים פרופיל של בלוק קוד ספציפי במודל. יש דרכים שונות לתעד מעקבים במסגרות השונות JAX, PyTorch ו-TensorFlow – מעקב התחלה או עצירה מבוסס-API או מבוסס-מנהל הקשר. במסגרות ברמה גבוהה יותר כמו MaxText, אתם יכולים להפעיל תיעוד פרופיל פשוט על ידי הפעלת דגל profiler=xplane כשאתם מריצים את MaxText.
אתם יכולים להשתמש באיסוף פרופילים על פי דרישה כשאתם רוצים לאסוף פרופילים באופן אד-הוק, או כשלא הפעלתם איסוף פרופילים פרוגרמטי. האפשרות הזו שימושית כשמזהים בעיה במדדי המודל במהלך ההרצה, ורוצים לתעד פרופילים באותו רגע למשך זמן מסוים כדי לאבחן את הבעיה.
כדי להפעיל תיעוד פרופיל על פי דרישה, עדיין צריך להפעיל את שרת XProf בתוך הקוד. הפעלת profiler.start_server תפעיל שרת XProf בעומס העבודה של למידת המכונה, שמאזין לטריגר של תיעוד לפי דרישה כדי להתחיל לתעד פרופילים. אפשר להפעיל תיעוד פרופיל על פי דרישה דרך ממשק המשתמש של Tensorboard או דרך CLI באמצעות הכלי XProfiler.
בקישורים הבאים מוסבר איך להפעיל תיעוד פרוגרמטי ותיעוד על פי דרישה בסביבות פיתוח שונות:
גם באיסוף פרופילים פרוגרמטי וגם באיסוף פרופילים על פי דרישה, צריך לציין את המיקום שבו רוצים לאחסן את הפרופילים שנאספו. אתם יכולים לאחסן את הפרופילים בספרייה ב-TPU VM עם נתיב שדומה ל-/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 VM, אתם יכולים לצפות בהם באמצעות TensorBoard באותו TPU VM.
כדי להתקין את XProf ו-TensorBoard ב-TPU VM, אפשר להריץ את הפקודות הבאות:
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 או ב-Pod של GKE, עם אפשרות לשנות את סוג המכונה בהתאם לצורכי המשתמש מבחינת מהירות הטעינה והעלות
- יצירת קישור לשיתוף פרופילים ולשיתוף פעולה עם חברי הצוות ועם מהנדסי Google
- יצירת פרופילים קלה יותר על פי דרישה של עומסי עבודה ב-GKE וב-Compute Engine, כדי לבחור כל מארח שמריץ את עומס העבודה שלכם לצורך תיעוד פרופילים
לפני שמשתמשים בספריית XProfiler, צריך לתעד פרופילים של קוד עומס העבודה באופן פרוגרמטי או על ידי הפעלת שרת הפרופילים, כדי שתוכלו לבצע תיעוד פרופילים על פי דרישה בשלב מאוחר יותר.
כדי להגדיר את ספריית XProfiler, צריך להגדיר את Google Cloud CLI וסביבה וירטואלית של Python, ואז פשוט להריץ את הפקודה:
pip install cloud-diagnostics-xprof
הפעולה הזו תתקין את כל יחסי התלות הנדרשים של XProf ו-TensorBoard.
כדי ליצור מכונה וירטואלית או pod של 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>
אפשר להגדיר את הנתיב של Cloud Storage לתיקיית השורש (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 במקום במכונה וירטואלית של Compute Engine, אפשר להעביר את הדגל –GKE אל xprofiler create. אם מארחים מכונה של TensorBoard ב-Pod של GKE, קל יותר לנהל את TensorBoard יחד עם שאר עומס העבודה שנפרס ב-GKE.
ה-VM או ה-Pod של GKE שמארחים את TensorBoard מאפשרים לטעון פרופילים גדולים ומספר פרופילים ב- Google Cloud הרבה יותר מהר מאשר ב-TensorBoard שמתארח באופן מקומי. על סמך בדיקות הביצועים שלנו, פרופילים של טעינה בסדר גודל של 1GB יופיעו תוך כמה דקות בטעינה הראשונה באמצעות ה-VM שמוגדר כברירת מחדל, c4-highmem-8. בנוסף, אתם יכולים לבחור את סוג המכונה שמתאים לצרכים שלכם מבחינת ביצועים ועלות.
אחרי שמריצים את הפקודה 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, מסגרת, שמות של מארח VM/Pod ומשך התיעוד באלפיות השנייה. אלה אותם פרטים שצריך להזין בממשק המשתמש של TensorBoard.
בדף הזה ב-GitHub מוסבר איך משתמשים בספרייה cloud-diagnostics-xprof.