מדוע דגמי AI מסוימים מהירים יותר מאחרים

טבלת קישורים

תקציר ומבוא 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. מסקנה והפניות

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

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

ג. התקנה

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

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

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

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

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

בספרות, ישנן דרכים מרובות לשיפור השהיה ו/או השהיה של מסקנות. קוונטיזציה מפחיתה את השימוש בזיכרון על ידי שימוש בייצוגים ברוחב סיביות נמוך כגון int8, int4 ו-fp8 (Wei et al., 2023; Yao et al., 2022; Dettmers et al., 2022; Frantar et al., 2022; Kuzmin et al., 2022; Xiao et al., 2022). קוונטיזציה כאשר היא מיושמת רק על פרמטרי מודל מציעה תוצאות פוחתות כמו עם אורכי רצף ארוכים יותר וגדלים גדולים של אצווה, שבהם גישה לזיכרון ומחשוב הקשורים לתשומת הלב של מוצר נקודה שולטים בהשהיית ההסקה הכוללת.

תשומת לב דלילה (Beltagy et al., 2020; Child et al., 2019; Zaheer et al., 2020) נחקרה בהרחבה כדרך להפחית את מורכבות הקשב להקשרים ארוכים יותר ולהסקת מסקנות מהירה יותר. פופ וחב'. (2022) חוקרת את יעילות ההסקה הגנרטיבית של מודלים של שפות גדולות על ידי שימוש בטכניקות חלוקה רב-ממדיות המותאמות עבור TPUs (collective einsum) כדי להשיג גבול Pareto על חביון וניצול FLOPs של מודל. המאמר גם מראה שתשומת לב מרובת שאילתות מאפשרת קנה מידה של עד פי 32 אורך הקשר גדול יותר עם דגש על היעילות בגודל אצווה גבוה. Paged attention (Kwon et al., 2023) משפר את ניהול הזיכרון של מטמון KV על ידי חלוקתו לבלוקים ושימוש בטבלת בלוקים למטרות מיפוי. גישה זו מתאימה ביעילות לשינויים דינמיים של עומס עבודה ומפחיתה את דרישות אחסון הזיכרון באמצעות שיתוף מטמון ה-KV של ההנחיה על פני רצפי פלט מרובים. עם זאת, זה לא מפחית את קריאות הזיכרון של KV cache.

פענוח ספקולטיבי והגרסאות שלו משתמשות במודל טיוטה קטן יותר כדי להציע אסימונים עוקבים מרובים, המעובדים במקביל על ידי המודל הראשי כדי לקבל או לדחות אסימונים כאלה (Chen et al., 2023; Leviathan et al., 2022; Li et al., 2024; Cai et al., 202., 202). הרעיון המרכזי הוא לאפשר פענוח של מספר אסימונים בכל שלב, ובכך להפחית את השימושים ב-IO בזיכרון של הדגם הראשי. עם זאת, ההשהיה של הפענוח עדיין תהיה נשלטת על ידי רוחב פס קלט/פלט של KV מטמון בגדלי הקשר גדולים, כאשר תשומת לב מפוצלת יכולה לשפר את מהירות הפענוח עוד יותר. בקיצור, פענוח מצטבר מתמקד בהורדת ה-IO של הזיכרון המופחת של טעינת המודל, בעוד שתשומת לב מרובת שאילתות ותשומת לב מפוצלת מורידה את ה-IO של הזיכרון של KV cache.

3. רקע

3.1. סִמוּן

אנו משתמשים בסימון הבא לאורך המאמר.

3.2. מסקנות מודל שפה

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

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

במהלך שלב קידוד ההקשר, מספר פעולות הנקודה הצפה ביחס לפעולות קלט/פלט זיכרון (IO) גבוה, המתאים למשטר המחובר למחשב שבו ההשהיה מושפעת מה-FLOPs. עם זאת, במהלך פענוח מצטבר שבו אנו מבצעים תשומת לב לאסימוני שאילתה בודדים, זה נופל למשטר מוגבל בזיכרון שבו מספר החישובים לכל גישה לזיכרון הוא בערך 1 ל-1 (ראה נספח D.1 לפרטים). הזיכרון IO מתייחס לפעולות הקריאה והכתיבה מזיכרון רוחב הפס הגבוה (HBM) (Jia et al., 2018) ל-SRAM המהיר על-שבב שבו מתבצע החישוב בפועל. זיכרון ה-IO של הפענוח המצטבר עצמו מורכב משני מרכיבים: (1) טעינת פרמטר הדגם ו-(2) טעינת KV מטמון. רכיב (1) קבוע ללא תלות באורך ההקשר m או גודל אצווה b כאשר רכיב (2) תלוי הן ב-m והן ב-b ושולט בזיכרון הכולל IO אם m או b גבוהים, מה שיכול להפוך לצוואר בקבוק משמעותי להסקת מסקנות. העבודה שלנו מתמקדת בעיקר בהפחתת רכיב (2).