פתרון בעיות ב-JAX – TPU
במדריך הזה מופיעים קישורים למידע על פתרון בעיות ב-JAX, כדי לעזור לכם לזהות ולפתור בעיות שאתם עלולים להיתקל בהן במהלך אימון של מודלים של JAX ב-Cloud TPU.
מדריך כללי יותר לתחילת העבודה עם Cloud TPU זמין במאמר מדריך למתחילים: יצירת מופע TPU.
בעיות כלליות ב-JAX
אם נתקלתם בבעיות במהלך פיתוח מודל האימון או האימון באמצעות JAX, תוכלו לעיין בשאלות הנפוצות בנושא JAX.
למידע נוסף על שגיאות כלליות בתכנות שאתם עשויים להיתקל בהן כשאתם כותבים אפליקציה לאימון באמצעות JAX, תוכלו לעיין במאמר בנושא שגיאות ב-JAX.
ביצועי פרופיל JAX
כדי להבין איך נעשה שימוש במשאבי ה-TPU, אפשר להשתמש בכלים שמתוארים במאמר יצירת פרופיל של ביצועי JAX.
פתרון בעיות שקשורות לזיכרונות
אפשר לעקוב אחרי השימוש בזיכרון באמצעות כלי ליצירת פרופיל של זיכרון המכשיר ב-JAX, אבל אי אפשר לנהל ישירות את השימוש בזיכרון.
אפשר להשתמש בפרופילר של זיכרון המכשיר ב-JAX כדי:
- לבדוק אילו מערכים וקבצים הפעלה נמצאים בזיכרון ה-TPU בזמן נתון, או
- איתור דליפות זיכרון
אי אפשר לציין איך הזיכרון של ה-TPU מוקצה לפעולות ספציפיות. מידע נוסף על בעיות בביצועים של TPU שספציפיות ל-JAX זמין במאמר הערות לגבי ביצועים של שימוש ב-TPU עם JAX.
פתרון בעיות שקשורות ל-TPU
בקטעים הבאים מוסבר איך לפתור כמה בעיות נפוצות שאתם עשויים להיתקל בהן כשמריצים תוכנית JAX ב-TPU.
איך אפשר לוודא שה-TPU פועל?
כל הפעולות יבוצעו ב-TPU כל עוד JAX לא מדפיס את ההודעה "No GPU/TPU found, falling back to CPU" (לא נמצא GPU/TPU, המערכת חוזרת לשימוש במעבד).
כדי לוודא שמעבד ה-TPU פעיל, אפשר לבדוק את jax.devices(), שבו אמורים להופיע כמה מכשירי TPU, או לבדוק באופן פרוגרמטי באמצעות: assert jax.devices()[0].platform == 'tpu'.
RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No TPU Platform available.
מוצגת הודעת השגיאה RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No
TPU Platform available או האזהרה הבאה ב-TPU VM:
/tmp/tpu_logs/tpu_driver.WARNING
W1118 17:40:20.985243 23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:VERSION
השגיאה הזו יכולה להצביע על כך שאתם מריצים את הגרסה הלא נכונה של TPU VM. צריך לוודא שמותקנת אצלכם גרסת זמן הריצה הנוכחית של JAX ולנסות שוב.
פתרון בעיות ב-TPU וב-GKE
כדי לקבל עזרה בפתרון בעיות, מפעילים רישום מפורט ביומן במניפסט של עומס העבודה ב-GKE, ואז שולחים את היומנים לתמיכה של GKE.
TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0
בקטעים הבאים מוסבר על הודעות שגיאה שקשורות להגדרות של TPU ו-GKE ועל פתרונות לבעיות.
אין נקודות קצה זמינות לשירות jobset-webhook-service
השגיאה הזו מציינת שחבילת העבודות לא הותקנה כמו שצריך. בודקים אם פועלים פודים של Kubernetes בפריסת jobset-controller-manager. מידע נוסף זמין במסמכי התיעוד בנושא פתרון בעיות ב-JobSet.
האתחול של TPU נכשל: החיבור נכשל
חשוב לוודא שגרסת הצומת של GKE היא 1.30.4-gke.1348000 ומעלה (אין תמיכה ב-GKE 1.31).