קוד שגיאה: E2001

קטגוריה: זמן הידור: סוג נתונים של RHS שלא נתמך בחומרה

השגיאה הזו מתרחשת כשסוג הנתונים שמשמש לאופרנד בצד שמאל (RHS) במכפלת מטריצות (למשל, 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. הפעלת Cast לסוגים מקוריים

הפתרון הכי אמין הוא להמיר ידנית את האופרנדים לסוג נתונים שנתמך בחומרה (כמו Float32 או BFloat16 ב-TPU v4+) בתוך הגרעין לפני פעולת matmul.

• למה: Float32 הוא סוג הנתונים האוניברסלי שנתמך באופן מקורי על ידי MXU בכל הדורות של TPU.
• פשרה: הפעולה הזו כרוכה בעלות של VPU (יחידת עיבוד וקטורי) – המחזורים שנדרשים לביצוע ההמרה, אבל היא מבטיחה שהליבה תפעל בחומרה הנוכחית.

2. בדיקת מצב התאימות

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

הפעולה הזו משמשת כ-polyfill, ומזריקה רצפים של אמולציה בתוכנה לפעולות שהחומרה לא תומכת בהן באופן טבעי.

מה אפשר לעשות במצב תאימות:

• Mixed-precision MatMuls: מאפשרת לערבב אופרנדים של מספרים שלמים עם צוברים של מספרים ממשיים (Float) על ידי הוספה אוטומטית של פעולות המרה (למשל, הרחבת מספרים שלמים ל-Float32 לפני פעולת הכפל של מטריצות).
• אמולציה עם דיוק נמוך: בדורות מסוימים של חומרה, המערכת מבצעת אמולציה של סוגים לא נתמכים כמו נקודה צפה (floating point) מסוג 4-bit (4E2M1FN) או נקודה צפה מסוג 8-bit (8E4M3FN) על ידי הרחבתם לסוגים נתמכים כמו BFloat16 או Float32 לפני ההפעלה.

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

3. שדרוג חומרה או בקשת תמיכה

אם האלגוריתם שלכם דורש ביצועים מקוריים באופן מוחלט עבור סוגים כמו Int4 או Float8 ללא התקורה של שידור או אמולציה, תצטרכו להריץ אותו בדור חדש יותר של TPU עם תמיכה מקורית.

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