Grouped Query Attention (GQA)

L’attention à requêtes groupées (GQA) est une forme d’attention multi-tête pour les Transformers, dans laquelle des groupes de têtes de requête partagent une tête de clé et une tête de valeur. Llama 3.1 8B, par exemple, associe 32 têtes de requête à 8 têtes clé-valeur dans chaque couche. La GQA se situe entre l’attention multi-tête (MHA), où chaque tête de requête possède ses propres clés et valeurs, et l’attention multi-requête (MQA), où toutes les têtes de requête partagent une seule tête de clé et une seule tête de valeur. Joshua Ainslie et ses collègues de Google Research l’ont présentée dans GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, à EMNLP 2023, montrant que les modèles GQA atteignent une qualité proche de MHA à une vitesse de décodage proche de MQA. Des grands modèles de langage (LLM) ouverts tels que Llama 2 70B, la famille Llama 3, Mistral 7B et Gemma 2 l’utilisent. Les modèles vision-langage (VLM) qui reposent sur ces modèles, comme PaliGemma 2 de Google basé sur Gemma 2, en héritent dans leurs décodeurs de langage.

Fonctionnement de la GQA#

Dans l’auto-attention, chaque jeton est projeté en une requête, une clé et une valeur pour chaque tête, et chaque tête compare ses requêtes à ses clés. La GQA conserve l’ensemble des têtes de requête, mais projette moins de têtes de clé et de valeur ; chaque groupe de têtes de requête utilise les clés et les valeurs partagées de son groupe. Avec 8 têtes de requête et 2 têtes clé-valeur, chaque tête clé-valeur sert 4 têtes de requête :

Le calcul de l’attention au sein de chaque tête reste une attention à produit scalaire mis à l’échelle. Avec un seul groupe, la couche devient MQA ; avec autant de groupes que de têtes de requête, elle devient MHA. PyTorch prend directement en charge cette disposition avec l’argument enable_gqa de 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")

Il affiche Output (1, 8, 16, 64), same as copied heads: True et 256 with GQA, 1024 with MHA : les 8 têtes de requête produisent des sorties à partir d’un quart des clés et des valeurs que MHA stockerait.

Pourquoi la GQA réduit le cache KV#

Pendant la génération de texte, un modèle conserve les clés et les valeurs de chaque jeton précédent dans son cache KV et lit l’intégralité du cache pour chaque nouveau jeton. La vitesse de décodage est donc davantage limitée par la bande passante mémoire que par les opérations arithmétiques. La taille du cache varie en fonction du nombre de têtes clé-valeur. Llama 3.1 8B comporte 32 couches, chacune avec 8 têtes clé-valeur de 128 dimensions ; il stocke donc 128 Kio par jeton en précision 16 bits, soit 16 Gio pour sa fenêtre de contexte complète de 131 072 jetons. Avec MHA, le même modèle nécessiterait 512 Kio par jeton et 64 Gio à pleine longueur. La mémoire libérée peut servir à des contextes plus longs, à des lots plus importants pour le traitement continu par lots et à réduire la latence d’inférence par jeton. L’interface LLM d’Ultralytics accepte un base_url pour tout serveur compatible avec OpenAI, comme un Llama 3.1 8B hébergé sur place. Grâce à ses 8 têtes clé-valeur, le serveur peut conserver quatre fois plus de jetons de conversation dans une mémoire de cache donnée qu’avec une disposition MHA à 32 têtes.

Conversion de modèles en GQA#

La GQA ne nécessite pas d’entraînement à partir de zéro. Les auteurs ont converti un point de contrôle MHA existant en calculant la moyenne des têtes de clé et de valeur de chaque groupe, puis ont poursuivi le préentraînement avec environ 5 % du calcul initial, une étape qu’ils appellent entraînement complémentaire. MQA, le cas extrême où une seule tête est partagée, a été introduite par Noam Shazeer en 2019. C’est la disposition la plus rapide, mais elle dégrade la qualité et peut rendre l’entraînement instable ; un nombre intermédiaire de groupes permet de récupérer la majeure partie de cette qualité.

GQA et techniques connexes#

• Attention latente multi-tête (MLA) : présentée avec DeepSeek-V2, elle compresse les clés et les valeurs de chaque jeton dans un petit vecteur latent partagé.
• PagedAttention : alloue le cache KV en blocs de taille fixe pour éviter la fragmentation, quel que soit le nombre de têtes qui le remplissent.
• Attention à fenêtre glissante : limite le nombre de jetons mis en cache par couche ; Mistral 7B la combine avec la GQA.
• Attention Flash : noyau GPU qui calcule l’attention exacte plus rapidement et prend en charge nativement les dispositions GQA.