طاها اکبری
طاها اکبری
خواندن ۲ دقیقه·۳ ماه پیش

Multi-Query attention و Group-Query attention

یک معماری که توی LLM های اخیر مورد توجه قرار گرفته معماری Group Query attention (GQA) هستش. جریان اینه که در زمان آموزش با توجه به این که همه توکن ها رو داریم می‌تونیم امبدینگ ها و پروجکشن‌هاشون رو همزمان و به شکل ضرب دو تا ماتریس حساب کنیم. ولی در زمان inference نیاز داریم که یک توکن رو توی هر زمان تولید کنیم و دوباره به شبکه بدیم که از موازی‌سازی به شکل ضرب ماتریس جلو‌گیری می‌کنه.

یک GPU وقتی دو تا ماتریس رو می‌خواد در هم ضرب کنه هر دو تا رو میاره توی کش و ضرب می‌کنه و باعث میشه که دسترسی مموریمون در حد همون دو تا ماتریس باشه. زمانی که توی حالت inference هستیم دیگه نمی‌تونیم عملیات رو با استفاده از ضرب ماتریس انجام بدیم و در واقع بعد از هر projection روی هر کدوم از query, key, value میریم سراغ لایه های بعدی attention و ماتریس های projection از کش خارج میشن و مجبوریم دوباره لودشون کنیم.

این لود کردن در هنگام inference باعث میشه که هر بار مجبور بشیم ماتریس های projection رو از مموری لود کنیم که باعث میشه که تعداد عملیات خواندن از مموری نسبت به تعداد عملیات حسابی زیاد بشه که توی پردازنده‌هایی مثل GPU که عملیات پردازشی رو سریع انجام میدن خوندن مموری به باتلنک تبدیل میشه.

برای حل این مشکل راه‌حلی که ارائه شده اینه که به جای این که از یک ماتریس project برای key, value برای هر کدوم از head های attention استفاده کنیم از یک ماتریس مشترک برای همه key ها و value ها استفاده کنیم. این باعث میشه که گلوگاه لود کردن ماتریس های key و value عملا از بین بره. این روش به Multi-Query attention معروفه که در مقابل Multi-Head attention قرار می‌گیره.

مشکلی که وجود داره اینه که در این روش عملا مقداری از عملکرد مدل فدا میشه و لذا پیشنهاد میشه که Head ها رو با هم دیگه گروه کنیم و برای هر گروه یک ماتریس تصویر key و value جدا در نظر بگیریم. این روش که عملا بین دو روش Multi-Query attention و Multi-Head attention قرار گرفته به Group-Query attention معروفه که در بعضی از LLM های نسبتا جدید مثلا Mistral هم استفاده شده.


transformerllm
شاید از این پست‌ها خوشتان بیاید