קטגוריה: זמן קומפילציה: סוג נתונים של צד ימין שלא נתמך בחומרה
השגיאה הזו מתרחשת כשסוג הנתונים שמשמש לאופרנד בצד שמאל
בכפל מטריצות (למשל, jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul או האופרטור @) לא נתמך באופן מובנה בדור ה-TPU הספציפי שבו נעשה שימוש.
הודעות שגיאה לדוגמה:
INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>'
...
The MLIR operation involved:
%13440 = "tpu.matmul"(%13435, %13437, %13439) <dimension_numbers = #tpu.dot_dimension_numbers<...>
XLA Backends: TPU
סקירה כללית
יחידת הכפל מטריצות (MXU) של TPU תומכת באופן מקורי בפעולות Float32 בכל הדורות של החומרה.
עם זאת, יש תמיכה מובנית ב-BFloat16 ובסוגים אחרים של נתונים כמותיים (למשל, Int4, Int8 או Float8) משתנה בהתאם לדור החומרה. השגיאה הזו מופעלת כשהליבה מנסה למפות כפל מטריצות ל-MXU באמצעות סוג נתונים שאין לדור הספציפי של TPU מעגלים פיזיים לביצוע שלו.
השגיאה הזו בדרך כלל מציינת שהמעבר Canonicalization של הקומפיילר – שמנסה להמיר באופן אוטומטי סוגים לא נתמכים לסוגים נתמכים (למשל, באמצעות אמולציה של תוכנה) – לא הצליח למצוא כלל המרה תקין או שנמנעה ממנו האפשרות לעשות זאת כי מצב התאימות הושבת.
ניפוי באגים
כדי לפתור את השגיאה הזו, צריך להתאים את סוגי הנתונים ליכולות של החומרה. יש לכם את האפשרויות הבאות:
1. המרת סוגים מקוריים
הפתרון הכי אמין הוא להמיר ידנית את האופרנדים לסוג נתונים שנתמך בחומרה (כמו Float32 או BFloat16 ב-TPU v4+) בתוך הליבה לפני פעולת matmul.
- למה:
Float32הוא סוג הנתונים האוניברסלי שנתמך באופן מקורי על ידי MXU בכל הדורות של TPU. - פשרה: הפעולה הזו כרוכה בעלות של VPU (יחידת עיבוד וקטורי) – המחזורים שנדרשים לביצוע ההמרה, אבל היא מבטיחה שהליבה תפעל בחומרה הנוכחית.
2. בדיקת מצב התאימות
בדרך כלל, הקומפיילר יכול לטפל אוטומטית בבעיות של אי התאמה בין סוגים במצב תאימות, שמופעל כברירת מחדל. בודקים שוב את ההגדרות של XLA
כדי לוודא ש---xla_mosaic_compat_mode לא מוגדר כ-False.
הפעולה הזו משמשת כ-polyfill, ומזריקה רצפים של אמולציה בתוכנה לפעולות שהחומרה לא תומכת בהן באופן טבעי.
מה אפשר לעשות במצב תאימות:
- Mixed-Precision MatMuls: מאפשרת לערבב אופרנדים של מספרים שלמים עם צוברים של Float
על ידי הוספה אוטומטית של פעולות המרה (למשל, הרחבת מספרים שלמים ל-
Float32לפני matmul). - אמולציה של דיוק נמוך: בדורות מסוימים של חומרה, המערכת מבצעת אמולציה של סוגים לא נתמכים כמו נקודה צפה (floating point) מסוג
4-bit(4E2M1FN) או נקודה צפה מסוג8-bit(8E4M3FN) על ידי הרחבתם לסוגים נתמכים כמוBFloat16אוFloat32לפני ההפעלה.
שימו לב: במצב הזה, התאימות מקבלת עדיפות על פני ביצועים אופטימליים, כי כדי שה-MXU יוכל לפעול על נתונים, צריך להמיר את פורמטי הנתונים באמצעות הוראות נוספות.
3. שדרוג חומרה או בקשת תמיכה
אם האלגוריתם שלכם דורש ביצועים מקוריים באופן מוחלט עבור סוגים כמו Int4 או Float8 ללא התקורה של שידור או אמולציה, תצטרכו להריץ אותו בדור חדש יותר של TPU עם תמיכה מקורית.
בקשה להוספת תכונה: אם לדעתכם החומרה שלכם תומכת בפעולה הזו, או אם חסר מהדר נתיב אמולציה תקין גם במצב תאימות, אתם יכולים להגיש בקשה להוספת תכונה. בדרך כלל אנחנו מבטיחים שהפעולות יהיו תואמות קדימה. לכן, אם ליבת המערכת שלכם פועלת בדור מסוים של TPU, היא אמורה לפעול בכל הדורות הבאים. אבל לא מובטח שיהיה בה אמולציה לדורות ישנים יותר (שחלק מההמרות שלהם יהיו יקרות מאוד).