مکانیسم های توجه پیشرفته – II

نویسنده(های): آریون داس

در ابتدا منتشر شد به سمت هوش مصنوعی.

توجه فلش.
می توانید به پیشینیان آن مراجعه کنید اینجا: حافظه پنهان KV، توجه پنجره کشویی، MHA، MQA، آموزش، & GQA. این روش ها به کار گرفته شد کاهش نیازهای حافظه و محاسبه، اما در هزینه کاهش کیفیت. ما در اینجا به چند راه برای رفع آن می پردازیم.

چرا؟

ما باید عملکرد فعلی را درک کنیم پردازنده گرافیکی برای به دست آوردن شهود پشت توجه فلاش.

مشاهده کنید که چگونه پردازنده گرافیکی SRAM بسیار سریعتر از HBM است. اکنون، در اینجا نگاهی گذرا به آنچه تا کنون انجام داده‌ایم آورده شده است:

1) ضرب ماتریس (Q,K)
الف) خواندن Q، K به SRAM (1)
ب) A = QxK را محاسبه کنید
ج) نوشتن A به HBM (2)

2) ماسک
الف) خواندن A به SRAM (3)
ب) ماسک A به A
ج) A’ را به HBM بنویسید (4)

3) سافت مکس
الف) خواندن A’ به SRAM (5)
ب) SoftMax A’ تا A”
ج) نوشتن A” به HBM (6)

مشاهده کنید که چند عملیات خواندن و نوشتن برای چنین محاسبات ساده ای انجام می شود. این ایده آل نیست. سخت افزار موجود به طور بهینه استفاده نمی شود.

توجه فلش

|| کاغذ ||

این مقاله به یک جزء مهم گم شده اشاره می کند – ساختن مکانیسم توجه I/O آگاه است.
هدف اصلی آنها این بود که از خواندن و نوشتن ماتریس توجه به و از HBM خودداری کنید. این مورد نیاز داشت:

1) محاسبه کاهش SoftMax بدون دسترسی به کل ورودی.
2) ذخیره نکردن ماتریس توجه میانی بزرگ برای عبور به عقب.

اکنون عملیات ما به این صورت است:

1) Q,K را به SRAM بخوانید
2) محاسبه کنید A = QxK
3) ماسک الف به الف
4) سافت مکس الف تا الف
5) A” را در HBM بنویسید

برای در نظر گرفتن آن، توجه فلش روی GPT-2 منجر به افزایش سرعت 7.6 برابری در محاسبه توجه می شود.

کاشی کاری

توجه کنید که چگونه حافظه SRAM 1/2000 برابر حافظه HBM است. از این رو، نویسندگان راه حلی را برای قابل محاسبه کردن عملیات در فضای کوچک و در عین حال سریع حافظه پیشنهاد می کنند.

هدف این است که کل ماتریس وزن را به قطعاتی تقسیم کنید که می توانند در حافظه SRAM قرار بگیرند، به چندین GPU اجازه می دهد تا آنها را به صورت موازی پردازش کنند. به طور موثر، به جای محاسبه توجه برای کل دنباله در یک حرکت، الگوریتم هر کاشی را به صورت متوالی یا ترجیحاً موازی پردازش می کند (زیرا اینجاست که بهینه سازی واقعی کار می کند).
این کار در مراحل زیر قابل انجام است:

مرحله 1: تکه تکه شدن تانسورهای پرس و جو، کلید و ارزش

با توجه به دنباله ای به طول n و هر کاشی به اندازه t، دنباله به قطعات تقسیم می شود به طوری که:
س: کاشی هایی با اندازه tx d_k
K: کاشی هایی با اندازه tx d_k
V: کاشی هایی با اندازه tx d_v

مرحله 2: محاسبه توجه جزئی در هر کاشی

برای هر جفت کاشی Q و K، FlashAttention توجه نقطه-محصول مقیاس شده را محاسبه می کند:
Q x K^T را برای کاشی محاسبه کنید: یک ماتریس امتیاز توجه به اندازه tx t می دهد.
نرمال سازی SoftMax را اعمال کنید: برای اطمینان از اینکه نمرات توجه جزئی کاشی در بافت کاشی نرمال شده است.
با کاشی V ضرب کنید: امتیازهای softmaxed برای کاشی سپس برای وزن کردن کاشی V مربوطه استفاده می شود.

مرحله 3: نتایج را جمع آوری کنید
امتیازات توجه از هر کاشی به صورت کاشی به کاشی انباشته می شود، که محاسبه نهایی توجه معادل QK^T را تضمین می کند، اما با نیازهای حافظه بسیار کمتر.

این کمک می کند کاهش سربار حافظه، استفاده بهینه از پهنای باند، & پردازش موازی را امکان پذیر می کند. رجوع شود به وبلاگ کریستین میل برای جزئیات پیاده سازی CUDA C++.

بهینه سازی مخرج SOFTMAX

تثبیت کننده سافت مکس

با نگاهی به فرمول softmax متوجه می شویم،
p_i = exp(l_i) / sum_j(exp(l_j))،
که در آن p احتمالات و lها لوجیت هستند.

مقادیر نمایی ممکن است در هر جهت منفجر شوند، بنابراین نویسندگان عملیات softmax را با ثبات می کنند کم کردن حداکثر logit (m) از همه logitها، اطمینان حاصل شود که شمارنده در داخل قرار دارد [0, 1].
p_i = exp(l_i – m) / sum_j(exp(l_j – m)).

افزایشی / SoftMax آنلاین

من این را با استفاده از یک قیاس توضیح می دهم.

یک استارتاپ جدید را تصور کنید. این نوید درآمد بالا و سود پروژه را می دهد. متأسفانه در چند سال اول با ضرر و زیان مواجه می شود و نمی تواند درآمد مطلوبی داشته باشد. بدون سود با این حال، در دراز مدت، عملکرد فوق العاده خوبی دارد و همان طور که وعده داده بود، سود بالایی ایجاد می کند. (بله، مانند آمازون.)

به طور مشابه، در اینجا ما برای محاسبه کل مخرج برای هر بلوک / کاشی تلاش نمی کنیم. در عوض، ما softmax را برای آن بلوک خاص محاسبه می کنیم. به عنوان مثال، برای بلوک با بردارهای x1 و x2،

آیا این درست است؟

خیر، اگر softmax را برای کل محاسبات توجه در نظر بگیریم.
بله، اگر فقط بلوک/کاشی را در i = 2 در نظر بگیریم.

اکنون، نگاهی به محاسبه در حال تکامل تا آخرین مرحله زمانی بیندازید.

این مقدار softmax در هر مُهر زمانی غیر از آخرین، ناقص است، اما در نهایت در آخرین مهر زمانی به مقدار مورد نیاز می رسیم. من فکر می کنم این یک راه حل عالی است که توسط محققان پیشنهاد شده است. آن را به کاهش نیازهای حافظه کمک می کند.

من سعی کردم مراحل مختلف این الگوریتم باورنکردنی را حاشیه نویسی کنم. با این حال، من به شدت توصیه می‌کنم که برای شهود، پیاده‌سازی را در مقاله انجام دهید.

موفق باشید آن را در یک حرکت درک کنید! اگر مایل به بحث در مورد جزئیات پیچیده هستید، از اینجا به من ایمیل بزنید: [email protected]

این مکانیسم های توجه پیشرفته را به پایان می رساند (از 12 نوامبر 2024). اگر مکانیسم های توجه بیشتری به وجود بیاید، مقالات بیشتری را به این مجموعه اضافه خواهم کرد.

Github: 🐈‍⬛
لینکدین: در
توییتر: X

منتشر شده از طریق به سمت هوش مصنوعی