Fonctionnement de la mise en cache

Imaginez que vous soyez en train de discuter avec quelqu’un et que, au lieu de retenir ce qu’il a dit précédemment, il doive repartir de zéro à chaque fois que vous répondez. Ce serait lent et inefficace, n’est-ce pas ?

Vous pouvez étendre cette analogie aux modèles transformers. La génération par un modèle autorégressif peut être lente car elle effectue une prédiction un jeton à la fois. Chaque nouvelle prédiction dépend de tout le contexte précédent.

Pour prédire le 1000e jeton, le modèle a besoin d'informations issues des 999 jetons précédents. Ces informations sont représentées sous forme de multiplications de matrices entre les représentations des jetons.

Pour prédire le 1001e jeton, vous avez besoin des mêmes informations issues des 999 jetons précédents, en plus de toute information provenant du 1000e jeton. Cela représente un grand nombre de multiplications de matrices que le modèle doit calculer à nouveau pour chaque jeton !

Un cache clé-valeur (KV) élimine cette inefficacité en stockant les paires kv dérivées des couches d'attention des tokens précédemment traités. Les paires kv stockées sont récupérées depuis le cache et réutilisées pour les tokens suivants, évitant ainsi le besoin de les recalculer.

La mise en cache ne doit être utilisée que pour l'inférence. Son activation pendant l'entraînement peut provoquer des erreurs inattendues.

Pour mieux comprendre comment et pourquoi la mise en cache fonctionne, examinons de plus près la structure des matrices d'attention.

Matrices d'attention

L' attention par produit scalaire normalisé est calculée comme indiqué ci-dessous pour un lot de taille b, nombre de têtes d'attention h, longueur de séquence jusqu'à présent T, et dimension par tête d'attention d_head. $\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V$

Les matrices de requête (Q), de clé (K) et de valeur (V) sont des projections des embeddings d'entrée de forme (b, h, T, d_head).

Pour l'attention causale, le masque empêche le modèle de prêter attention aux jetons futurs. Une fois un jeton traité, sa représentation ne change plus par rapport aux jetons futurs, ce qui signifie que $K_{\text{past}}$ et $V_{\text{past}}$ peuvent être mis en cache et réutilisés pour calculer la représentation du dernier jeton. $\text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{cached}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{cached}}, v_{t}])$

Lors de l'inférence, vous n'avez besoin que de la requête du dernier jeton pour calculer la représentation $x_t$ qui prédit le jeton suivant $ t+1 $. À chaque étape, les nouveaux vecteurs clé et valeur sont stockés dans le cache et ajoutés aux clés et valeurs passées. $K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)$

L'attention est calculée indépendamment dans chaque couche du modèle, et la mise en cache est effectuée couche par couche.

Reportez-vous au tableau ci-dessous pour comparer comment le cache améliore l'efficacité.

sans mise en cache	avec mise en cache
pour chaque étape, recalculer toutes les valeurs précédentes de `K` et `V`	pour chaque étape, ne calculer que les `K` et `V` actuels
le coût de l'attention par étape est quadratique par rapport à la longueur de la séquence	le coût de l'attention par étape est linéaire par rapport à la longueur de la séquence (la mémoire augmente de manière linéaire, mais le calcul par jeton reste faible)

Classe de cache

Une interface de cache KV de base prend un tenseur de clé et un tenseur de valeur pour le jeton actuel et renvoie les tenseurs K et V mis à jour. Cela est géré en interne par la méthode forward d’un modèle.

new_K, new_V = cache.update(k_t, v_t, layer_idx)
attn_output = attn_layer_idx_fn(q_t, new_K, new_V)

Lorsque vous utilisez la classe Cache de Transformers, le module d'auto-attention effectue plusieurs étapes essentielles pour intégrer les informations passées et présentes.

Le module d'attention concatène les paires kv actuelles avec les paires kv passées stockées dans le cache. Cela crée des poids d'attention de forme (new_tokens_length, past_kv_length + new_tokens_length). Les paires kv actuelles et passées sont essentiellement combinées pour calculer les scores d'attention, garantissant que le modèle prend en compte le contexte précédent et l'entrée actuelle.
Lorsque la méthode forward est appelée de manière itérative, il est crucial que la forme du masque d’attention corresponde à la longueur combinée des paires kv passées et actuelles. Le masque d’attention doit avoir la forme (batch_size, past_kv_length + new_tokens_length). Cela est généralement géré en interne par generate(), mais si vous souhaitez implémenter votre propre boucle de génération avec Cache, gardez cela à l’esprit ! Le masque d’attention doit contenir les valeurs des jetons passés et actuels.

Implémentation du stockage de cache

Les caches sont structurés sous forme d'une liste de couches, où chaque couche contient un cache de clés et un cache de valeurs. Les caches de clés et de valeurs sont des tenseurs de forme [batch_size, num_heads, seq_len, head_dim].

Les couches peuvent être de différents types (par exemple DynamicLayer, StaticLayer, StaticSlidingWindowLayer), ce qui modifie principalement la manière dont la longueur de la séquence est gérée et dont le cache est mis à jour.

La plus simple est un DynamicLayer qui grandit à mesure que de nouveaux jetons sont traités. La dimension de la longueur de séquence (seq_len) augmente avec chaque nouveau jeton :

cache.layers[idx].keys = torch.cat([cache.layers[idx].keys, key_states], dim=-2)
cache.layers[idx].values = torch.cat([cache.layers[idx].values, value_states], dim=-2)

D'autres types de couches, comme StaticLayer et StaticSlidingWindowLayer, ont une longueur de séquence fixe définie lors de la création du cache. Cela les rend compatibles avec torch.compile. Dans le cas de StaticSlidingWindowLayer, les jetons existants sont décalés hors du cache lorsqu'un nouveau jeton est ajouté.

L'exemple ci-dessous montre comment créer une boucle de génération avec DynamicCache. Comme discuté, le masque d'attention est une concaténation des valeurs des jetons passés et actuels.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache
from accelerate import Accelerator

device = Accelerator().device

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map=device)
tokenizer = AutoTokenizer.from_pretrained(model_id)

past_key_values = DynamicCache(config=model.config)
messages = [{"role": "user", "content": "Hello, what's your name."}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)

generated_ids = inputs.input_ids
max_new_tokens = 10

for _ in range(max_new_tokens):
    outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
    # Greedily sample one next token
    next_token_ids = outputs.logits[:, -1:].argmax(-1)
    generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)
    # Prepare inputs for the next generation step by leaving unprocessed tokens, in our case we have only one new token
    # and expanding attn mask for the new token, as explained above
    attention_mask = inputs["attention_mask"]
    attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
    inputs = {"input_ids": next_token_ids, "attention_mask": attention_mask}

print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])
"[INST] Hello, what's your name. [/INST]  Hello! My name is LLaMA,"

Update on GitHub