StableHLO Quantization

סוגי קוונטיזציה ב-StableHLO

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

הכימות ב-StableHLO פועל לפי מפרט הכימות של LiteRT, באמצעות תוכנית כימות אחידה עם תמיכה בכימות לכל טנזור וכימות לכל ציר. הוא מקבל בירושה את ביטוי הסוג שלו מQuant dialect של MLIR, ומספק דרך סטנדרטית לייצוג סוגי נתונים שעברו קוונטיזציה.

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

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

הקשר בין ערך הנקודה הצפה המקורי (real_value) לבין ערך השלם שעבר קוונטיזציה (quantized_value) בקוונטיזציה אחידה הוא:

real_value = scale * (quantized_value - zero_point)

Per-tensor Quantization

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

quant.uniform scale:zero_point>

דוגמה: !quant.uniform<i8:f32, 0.01:50>

הערך הזה מייצג מספר שלם (integer) ב-8 ביט (i8) שמשמש לאחסון מספר בשיטת נקודה צפה (floating-point) ב-32 ביט (f32) באמצעות קנה מידה של 0.01 ונקודת אפס של 50.

Per-axis Quantization

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

נניח שיש טנזור t עם גדלים של מימדים [4, 3, 2]. אנחנו בוחרים לכמת את הטנזור הזה לאורך המימד השני (quantized_dimension = 1). המשמעות היא שיהיו לנו שלושה פרוסים (כי הגודל של המימד השני הוא 3), שלכל אחד מהם יש קנה מידה משלו ונקודת אפס משלו:

t[:, 0, :]: This slice gets scale[0] and zero_point[0].
t[:, 1, :]: This slice gets scale[1] and zero_point[1].
t[:, 2, :]: This slice gets scale[2] and zero_point[2].

ב-StableHLO, סוג הכמות שנקבעת לכל ציר מבוטא כך:

quant.uniform {scale0:zero_point0, scale1:zero_point1, ...}>

כאשר האורך של scale:zero_point תואם למספר הפרוסות לאורך quantized_dimension של הטנזור המכיל.

דוגמה: tensor<4x3x2x!quant.uniform<i8:f32:1, {0.2:20, 0.1:10, 0.3:30}>>

שלבי קוונטיזציה ב-StableHLO

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

stablehlo-legalize-qdq-to-quantized-op

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

stablehlo-legalize-quantized-op-to-qdq

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

stablehlo-legalize-quant-to-math

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

stablehlo-quant-legalize-to-tosa-rescale

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

tosa-rescale-legalize-to-stablehlo

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

הערכה של תוכניות שעברו קוונטיזציה

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

בעצם, המפרש משתמש ב-stablehlo-legalize-quant-to-math כדי להמיר פעולות שעברו קוונטיזציה להטמעות המתאימות שלהן של אריתמטיקה של מספרים שלמים. במעבר הזה מוצגות פעולות שידור של CHLO לטיפול בפעולות של כפל או חילוק בקנה מידה והוספה של נקודת אפס. כדי להבטיח תאימות למתורגמן של StableHLO, הפעולות האלה של CHLO עוברות המרה לפעולות של StableHLO. השינוי הזה מוסיף פעולות שקשורות לצורות, שבהמשך עוברות קנוניזציה ואופטימיזציה באמצעות סדרה של מעברים קנוניים.

הרצף המלא של המעברים שקשורים לתהליך ההורדה הזה הוא:

stablehlo-legalize-quant-to-math
chlo-legalize-to-stablehlo
canonicalize
shape-legalize-to-stablehlo
stablehlo-canonicalize-dynamism

מקרי בדיקה עם כימות

‫StableHLO מספקת חבילה מקיפה של תרחישי בדיקה עם כימות כדי לאמת את הנכונות וההתנהגות של פעולות עם כימות. מקרי הבדיקה האלה משמשים כבדיקות יחידה, והם כוללים פעולות שונות של StableHLO בתרחישים כמותיים.

דוגמה טיפוסית למקרה בדיקה שעבר קוונטיזציה:

func.func @main() -> tensor<11xf32> {
    %operand_0 = stablehlo.constant dense<...> : tensor<11xf32>
    %operand_1 = stablehlo.constant dense<...> : tensor<11xf32>
    %golden = stablehlo.constant dense<...> : tensor<11xf32>

    %0 = stablehlo.uniform_quantize %operand_0 : (tensor<11xf32>) -> tensor<11x!quant.uniform<i8:f32, 0.3>>
    %1 = stablehlo.uniform_quantize %operand_1 : (tensor<11xf32>) -> tensor<11x!quant.uniform<i8:f32, 0.3>>

    %2 = stablehlo.add %1, %0 : tensor<11x!quant.uniform<i8:f32, 0.3>>

    %result = stablehlo.uniform_dequantize %2 : (tensor<11x!quant.uniform<i8:f32, 0.3>>) -> tensor<11xf32>

    %4 = stablehlo.custom_call @check.eq(%golden, %result) : (tensor<11xf32>, tensor<11xf32>) -> tensor<i1>

    return %3 : tensor<11xf32>
  }

וכוללת:

תרחישי הבדיקה האלה שימושיים במקרים הבאים:

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