Grouped Query Attention (GQA)

الانتباه إلى الاستعلامات المجمّعة (GQA) نوع من الانتباه متعدد الرؤوس في المحوّلات، تتشارك فيه مجموعات رؤوس الاستعلام رأسًا واحدًا للمفتاح ورأسًا واحدًا للقيمة. فعلى سبيل المثال، يقرن Llama 3.1 8B 32 رأس استعلام بثمانية رؤوس للمفتاح والقيمة في كل طبقة. ويقع هذا الأسلوب بين الانتباه متعدد الرؤوس (MHA)، حيث يكون لكل رأس استعلام مفاتيحه وقيمه الخاصة، والانتباه متعدد الاستعلامات (MQA)، حيث تتشارك كل رؤوس الاستعلام رأسًا واحدًا للمفتاح ورأسًا واحدًا للقيمة. قدم Joshua Ainslie وزملاؤه في Google Research هذا الأسلوب في بحث GQA: تدريب نماذج محوّلات معممة متعددة الاستعلامات انطلاقًا من نقاط تحقق متعددة الرؤوس في EMNLP 2023، وأظهروا أن نماذج GQA تحقق جودة قريبة من MHA وسرعة فك ترميز قريبة من MQA. وتستخدمه نماذج اللغة الكبيرة (LLMs) المفتوحة مثل Llama 2 70B وعائلة Llama 3 وMistral 7B وGemma 2، كما ترثه مفككات اللغة في نماذج الرؤية واللغة (VLMs) المبنية عليها، مثل PaliGemma 2 من Google المبني على Gemma 2.

آلية عمل GQA#

في الانتباه الذاتي، يُسقط كل رمز إلى استعلام ومفتاح وقيمة لكل رأس، ويقارن كل رأس استعلاماته بمفاتيحه. يحتفظ GQA بمجموعة رؤوس الاستعلام كاملة، لكنه يسقط عددًا أقل من رؤوس المفاتيح والقيم، وتستخدم كل مجموعة من رؤوس الاستعلام المفاتيح والقيم المشتركة لمجموعتها. مع ثمانية رؤوس استعلام ورأسي مفتاح وقيمة، يخدم كل رأس للمفتاح والقيمة أربعة رؤوس استعلام:

لا يتغير حساب الانتباه داخل كل رأس؛ فهو يظل انتباهًا دقيقًا بجداء نقطي مُقاس. ومع مجموعة واحدة تصبح الطبقة MQA، ومع عدد مجموعات يساوي عدد رؤوس الاستعلام تصبح MHA. يدعم PyTorch هذا التخطيط مباشرةً عبر الوسيط enable_gqa في scaled_dot_product_attention:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(1, 8, 16, 64)  # 8 query heads, 16 tokens, 64 dimensions per head
k = torch.randn(1, 2, 16, 64)  # only 2 key heads
v = torch.randn(1, 2, 16, 64)  # and 2 value heads

out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

# Equivalent multi-head attention: copy each key-value head to its group of 4 query heads
k_full, v_full = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
reference = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=True)

print(f"Output {tuple(out.shape)}, same as copied heads: {torch.allclose(out, reference, atol=1e-6)}")
print(f"Cached values per token per layer: {k[0, :, 0].numel() * 2} with GQA, {k_full[0, :, 0].numel() * 2} with MHA")

يطبع Output (1, 8, 16, 64), same as copied heads: True و256 with GQA, 1024 with MHA: تنتج رؤوس الاستعلام الثمانية كلها مخرجات باستخدام ربع عدد المفاتيح والقيم التي كان سيخزنها MHA.

كيف يقلّص GQA ذاكرة التخزين المؤقت للمفاتيح والقيم#

أثناء توليد النص، يحتفظ النموذج بمفاتيح وقيم كل رمز سابق في ذاكرة التخزين المؤقت للمفاتيح والقيم (KV)، ويقرأ الذاكرة المؤقتة كاملةً مع كل رمز جديد؛ لذا تحدد نطاق الذاكرة سرعة فك الترميز أكثر من العمليات الحسابية. ويتناسب حجم الذاكرة المؤقتة مع عدد رؤوس المفاتيح والقيم. يتضمن Llama 3.1 8B 32 طبقة، تضم كل منها 8 رؤوس للمفاتيح والقيم بأبعاد 128، لذا يخزن 128 كيلوبايت لكل رمز باستخدام دقة 16 بت، أو 16 جيبيبايت عند بلوغ نافذة السياق القصوى البالغة 131,072 رمزًا. وباستخدام MHA، يحتاج النموذج نفسه إلى 512 كيلوبايت لكل رمز و64 جيبيبايت عند بلوغ الحد الأقصى للطول. وتُستخدم الذاكرة المتاحة لزيادة طول السياق، أو زيادة أحجام الدفعات من أجل تجميع الدفعات المستمر، أو خفض زمن استجابة الاستدلال لكل رمز. تقبل واجهة LLM من Ultralytics وسيطًا base_url لأي خادم متوافق مع OpenAI، مثل خادم Llama 3.1 8B مستضاف ذاتيًا؛ إذ تتيح رؤوس المفاتيح والقيم الثمانية للخادم الاحتفاظ بأربعة أضعاف عدد رموز المحادثة في ذاكرة التخزين المؤقت نفسها مقارنة بتخطيط MHA ذي 32 رأسًا.

تحويل النماذج إلى GQA#

لا يتطلب GQA التدريب من الصفر. فقد حوّل المؤلفون نقطة تحقق MHA موجودة عبر حساب متوسط رؤوس المفاتيح والقيم ضمن كل مجموعة، ثم واصلوا التدريب المسبق باستخدام نحو 5% من الحوسبة الأصلية، وهي خطوة يسمونها التدريب التكميلي. وقدم Noam Shazeer أسلوب MQA، وهو الحد الأقصى عند مشاركة رأس واحد، في عام 2019؛ ويُعد أسرع تخطيط، لكنه خفّض الجودة وقد يجعل التدريب غير مستقر، بينما يستعيد استخدام عدد متوسط من المجموعات معظم تلك الجودة.

GQA والتقنيات ذات الصلة#

• الانتباه الكامن متعدد الرؤوس (MLA): قُدّم مع DeepSeek-V2، ويضغط مفاتيح كل رمز وقيمه في متجه كامن صغير مشترك.
• PagedAttention: يخصص ذاكرة التخزين المؤقت للمفاتيح والقيم في كتل ثابتة الحجم لتجنب التجزئة، بصرف النظر عن عدد الرؤوس التي تشغلها.
• انتباه النافذة المنزلقة: يحدد حدًا أقصى لعدد الرموز المخزنة مؤقتًا لكل طبقة، ويجمع Mistral 7B بينه وبين GQA.
• الانتباه السريع: نواة GPU تحسب الانتباه الدقيق بسرعة أكبر وتدعم تخطيطات GQA أصلًا.