Grouped Query Attention (GQA)

Gruplu sorgu dikkati (GQA), transformer modellerine yönelik çok başlı bir dikkat türüdür; bu yöntemde sorgu başı grupları tek bir anahtar başını ve tek bir değer başını paylaşır. Örneğin Llama 3.1 8B, her katmanda 32 sorgu başını 8 anahtar-değer başıyla eşleştirir. Her sorgu başının kendi anahtar ve değerlerine sahip olduğu çok başlı dikkat (MHA) ile tüm sorgu başlarının tek bir anahtar ve değer başını paylaştığı çoklu sorgu dikkati (MQA) arasında yer alır. Google Research'ten Joshua Ainslie ve çalışma arkadaşları bu yöntemi EMNLP 2023'te GQA: Çok Başlı Kontrol Noktalarından Genelleştirilmiş Çoklu Sorgu Transformer Modelleri Eğitimi başlıklı çalışmada tanıttı. GQA modellerinin, MHA'ya yakın kaliteye ve MQA'ya yakın çözümleme hızına ulaştığını gösterdiler. Llama 2 70B, Llama 3 ailesi, Mistral 7B ve Gemma 2 gibi açık büyük dil modelleri (LLM'ler) bu yöntemi kullanır. Google'ın Gemma 2 üzerine kurulu PaliGemma 2 modeli gibi, bu modeller üzerine kurulan görü-dil modelleri (VLM'ler) de dil kod çözücülerinde aynı yöntemi kullanır.

GQA Nasıl Çalışır#

Öz-dikkatte, her token her baş için bir sorguya, anahtara ve değere yansıtılır; her baş da sorgularını anahtarlarıyla karşılaştırır. GQA, sorgu başlarının tümünü korur ancak daha az sayıda anahtar ve değer başı üretir. Her sorgu başı grubu, kendisiyle paylaşılan anahtar ve değerleri kullanarak dikkat uygular. 8 sorgu başı ve 2 anahtar-değer başıyla her anahtar-değer başı 4 sorgu başına hizmet eder:

Her baştaki dikkat hesaplaması, ölçekli noktasal çarpım dikkati olarak değişmeden kalır. Tek bir grup katmanı MQA'ya, sorgu başı sayısı kadar grup ise MHA'ya dönüştürür. PyTorch, bu düzeni doğrudan scaled_dot_product_attention öğesinin enable_gqa bağımsız değişkeni üzerinden destekler:

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 ve 256 with GQA, 1024 with MHA çıktısını verir: MHA'nın depolayacağı anahtar ve değerlerin dörtte biriyle 8 sorgu başının tümü çıktı üretir.

GQA, KV Önbelleğini Neden Küçültür?#

Metin üretimi sırasında model, önceki her tokenın anahtar ve değerlerini KV önbelleğinde tutar ve her yeni token için önbelleğin tamamını okur. Bu nedenle çözümleme hızı, işlem gücünden çok bellek bant genişliğiyle sınırlıdır. Önbellek boyutu, anahtar-değer başlarının sayısıyla orantılıdır. Llama 3.1 8B'de 32 katman vardır ve her katman 128 boyutlu 8 anahtar-değer başı içerir. Bu nedenle model, 16 bit duyarlılıkta token başına 128 KiB, 131.072 tokenlık tam bağlam penceresinde ise 16 GiB depolar. MHA ile aynı model token başına 512 KiB, tam uzunlukta ise 64 GiB gerektirirdi. Boşalan bellek; daha uzun bağlamlara, sürekli gruplama için daha büyük gruplara ve token başına daha düşük çıkarım gecikmesine ayrılır. Ultralytics LLM arayüzü, kendi sunucunda barındırılan Llama 3.1 8B gibi OpenAI uyumlu tüm sunucular için bir base_url kabul eder. Bu sunucuda 8 anahtar-değer başı, belirli bir önbellek belleğinde 32 başlı MHA düzenine kıyasla dört kat daha fazla konuşma tokenının tutulmasını sağlar.

Modelleri GQA'ya Dönüştürme#

GQA, eğitime sıfırdan başlamayı gerektirmez. Yazarlar, her grup içindeki anahtar ve değer başlarının ortalamasını alarak mevcut bir MHA kontrol noktasını dönüştürdü; ardından ilk hesaplama miktarının yaklaşık %5'ini kullanarak ön eğitime devam etti. Bu aşamaya ek eğitim adını verdiler. Tek paylaşımlı başın uç örneği olan MQA, 2019'da Noam Shazeer tarafından tanıtıldı ve en hızlı düzendir; ancak kaliteyi düşürdüğü ve eğitimi kararsızlaştırabildiği için ara sayıda grup, bu kalitenin çoğunu geri kazandırır.

GQA ve İlgili Teknikler#

• Çok başlı gizil dikkat (MLA): DeepSeek-V2 ile tanıtılmıştır; her tokenın anahtar ve değerlerini küçük, paylaşılan bir gizil vektörde sıkıştırır.
• PagedAttention: Parçalanmayı önlemek için KV önbelleğini sabit boyutlu bloklara ayırır; bu işlem, önbelleği dolduran baş sayısından bağımsızdır.
• Kayan pencere dikkati: Her katmandaki önbelleğe alınmış token sayısını sınırlar; Mistral 7B bunu GQA ile birlikte kullanır.
• Flash attention: Tam dikkati daha hızlı hesaplayan ve GQA düzenlerini yerel olarak destekleyen bir GPU çekirdeğidir.