یکی از روش ‌های رایج که مدل های متن باز ازش استفاده میکنند Grouped-Query Attention و Multi-Query Attention هست که اندازه کش KV رو به طور قابل توجهی کاهش میدن.

مدل های متن باز مثل Llama 3.3 70B و Mistral Large 2 از این روش استفاده کردن و اندازه کش KV رو به شکل چشمگیری کاهش دادن.

https://arxiv.org/pdf/1911.02150

در معماری Multi-Head Attention، هر head مجموعه منحصر به فردی از بردارهای Q، K و V رو محاسبه میکنه اما در معماری Multi-Query Attention، فقط بردارهای Q برای هر head منحصر به فرده، در حالی که بردارهای Key و Value بین همه head ها به اشتراک گذاشته میشه. 

در این حالت با توجه به فرمول قبل اندازه حافظه مورد نیاز به 468MB کاهش پیدا میکنه! اما مقاله زیر نشون میده که این روش مشکلاتی رو در وظایفی مثل خلاصه سازی متون رو میتونه به همراه داشته باشه و معماری Multi-Head Attention در این وظایف بهتر عمل میکنه.

https://arxiv.org/abs/2305.13245

همونطور که میدونید مدل هایی که از معماری transformer استفاده میکنن اغلب به حافظه بیشتری نسبت به مقدار حافظه ای که در یک GPU قرار داره، نیاز دارن بنابراین به نوعی از موازی ‌سازی برای آموزش و استنتاج نیاز داریم! موازی‌ سازی تنسور(Tensor Parallelism) روشی برای موازی ‌سازی مدل‌ های بزرگ تولید متن هست.

به طور خلاصه، موازی ‌سازی تنسور با اختصاص attention head مختلف به GPUهای مختلف کار میکنه.

برای مثال، مدل ۱۷۵ میلیارد پارامتری GPT-3 دارای ۹۶ head هست. میتونیم این مدل رو روی ۸ GPU با اختصاص ۱۲ head به هر کدام موازی کنیم.

هر GPU کل محاسبات attention رو برای head اختصاص داده شده خودش انجام میده، در نهایت همه GPUها نتایج محاسبات خودشون رو مبادله میکنن.

اما در معماری Multi-Query Attention، هر attention head روی مجموعه یکسانی از بردارهای K و V عمل میکنه و این یعنی ما مجموعه یکسانی از بردارهای K و V را روی هر GPU محاسبه و کش می‌کنیم.

https://arxiv.org/abs/1910.10683

ادامه دارد...