כיצד להאיץ את דגמי הבינה המלאכותית שלך - מבלי לטגן את הזיכרון שלך

טבלת קישורים

תקציר ומבוא 1

2. עבודה קשורה

3. רקע

3.1. סימון ו-3.2. מסקנות מודל שפה

3.3. ריבוי שאילתות, רב ראשים ותשומת לב מרובת שאילתות כללית

4. תשומת לב מפוצלת המודעת להקשר ו-4.1. מוֹטִיבָצִיָה

4.2. ניסוח ו-4.3. זיכרון IO מורכבות

5. ניסויים

5.1. השוואת יכולות של ריבוי ראשים, ריבוי שאילתות ותשומת לב מרובה קבוצות

5.2. השהיות של מודלים שווים ליכולות

5.3. יישומים

6. מסקנה והפניות

א. שאלות נפוצות

ב. עבודה קשורה

ג. התקנה

ד. משפחת קשב רב קבוצתית

ה. תשומת לב מפוצלת המודעת להקשר

ו. יישומים: תוצאות נוספות

ז. תאימות עם טכניקות פענוח ספקולטיבי ופענוח מהיר

תַקצִיר

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

1. הקדמה

הופעתם של מודלים של שפות גדולות (LLMs) הובילה עידן חדש של למידת מכונה, המציגים ביצועים יוצאי דופן במגוון רחב של משימות (Brown et al., 2020; OpenAI, 2023; Chowdhery et al., 2022; Touvron et al., 2023; Hoffman et al., 202 et al.; Hoffmann et al., 2021. Li et al., 2022; אמזון, 2022; למרות היכולות המרשימות שלהם, הפריסה של מודלים בקנה מידה גדול אלה ביישומים מעשיים מציבה אתגרים משמעותיים, במיוחד במונחים של חביון מסקנות ויעילות. שיפור היבטים אלה הוא קריטי, שכן הם משפיעים ישירות על משאבי החישוב הנדרשים ליצירת תחזיות ולאפשר יישום מעשי של מודלים מתקדמים אלה בתעשיות שונות.

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

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

החקירה שלנו מתחילה בניתוח של תשומת הלב המוכללת מרובת שאילתות (Ainslie et al., 2023), הכוללת ריבוי שאילתות (Shazeer, 2019), כמו גם את מנגנון הקשב הרב-ראשי שהוקם (Vaswani et al., 2017) עבור ביצועים והשהייה. הממצאים שלנו מראים קנה מידה חלק של ביצועים עם הגדלת גודל המודל עבור ערך קבוע של מספר הקבוצות g עבור ריבוי שאילתות כלליות[1]. הורדת g גורמת להזזה כלפי מעלה של עקומות קנה המידה של אובדן האימות לעומת גודל המודל. הקשר העקבי בין דחיסת המטמון, גודל הדגם ואובדן האימות מאפשר לנו להשוות בין יעילות מסקנות לגודל המודל, כלומר, מאפשר לנו לבחור דחיסה גבוהה יותר עבור מקרי שימוש הדורשים יעילות גבוהה, תוך התאמת הביצועים של תשומת לב מרובה ראשים על ידי פיצוי בגודל דגם גדול יותר.

שנית, אנו מציגים קשב מפוצל מודע להקשר, טכניקה שמפצלת כל תשומת לב במשפחת השאילתות המוכללת להקשר ורכיבי פענוח במהלך פענוח מצטבר. התפצלות כזו כוללת את אותו מספר של FLOPs ומניבה תוצאות זהות בהשוואה לתשומת הלב המקורית, אך יכולה להפחית משמעותית את עלות ה-IO של זיכרון ובכך להשהות בתרחישי אצווה ואורך הקשר גבוהים. גישה זו מאפשרת יצירת השלמות מרובות בזמן אמת מבלי לגרור עלויות חביון נוספות, או מאפשרת גדלי אצווה גבוהים בהרבה המובילים לשיפור ביצועי הדירוג. לדוגמה, עבור מודל CodeGen 16B מרובה ראשים (Nijkamp et al., 2022) עם אורך הקשר של 2k, אנו מסוגלים להגדיל את גודל האצווה ל-128 עם תשומת לב מפוצלת, בהשוואה לגודל אצווה של 5 בלבד ללא, וכתוצאה מכך ה-pass@k (Chen et al., 2021) גדל מ-59.06% ל-log@top ממוצע ל-59.06%, או מ-59.06% 55.2% עד 58.1%.

[1] ערכים נמוכים יותר של קבוצות קשב g מובילים לדחיסה גבוהה יותר של טנסור המפתח-ערכי, כמו במקרה של ריבוי שאילתות שבו g = 1, ומכאן שיפור יעילות ההסקה וההשהיה עקב מטמון KV מופחת בהשוואה למקרה הרב-ראשי שבו g = h, מספר ראשי הקשב של השאילתה.