Optimiser les LLM pour la vitesse et la mémoire

Les grands modèles de langage (LLM), tels que GPT3/4, Falcon et Llama, progressent rapidement dans leur capacité à traiter des tâches centrées sur l'humain, s'imposant comme des outils essentiels dans les industries modernes basées sur la connaissance. Le déploiement de ces modèles dans des tâches réelles reste toutefois difficile :

Pour afficher des capacités de compréhension et de génération de texte proches de celles de l’humain, les LLM doivent actuellement être composés de milliards de paramètres (voir Kaplan et al, Wei et. al). Cela amplifie par conséquent les exigences en mémoire pour l’inférence.
Dans de nombreuses tâches réelles, les LLM doivent recevoir des informations contextuelles étendues. Cela nécessite que le modèle soit capable de gérer des séquences d’entrée très longues lors de l’inférence.

Le cœur de ces défis réside dans l'augmentation des capacités de calcul et de mémoire des LLM, en particulier lors du traitement de séquences d'entrée étendues.

Dans ce guide, nous passons en revue les techniques efficaces pour le déploiement efficace des LLM :

Précision réduite : Les recherches ont montré que le fonctionnement à une précision numérique réduite, à savoir 8 bits et 4 bits, peut offrir des avantages de calcul sans dégradation significative des performances du modèle.
Flash Attention : Flash Attention est une variante de l'algorithme d'attention qui offre non seulement une approche plus efficace en termes de mémoire, mais également une efficacité accrue grâce à une utilisation optimisée de la mémoire GPU.
Innovations architecturales : Étant donné que les LLM sont toujours déployés de la même manière lors de l'inférence, à savoir la génération de texte autorégressive avec un contexte d'entrée long, des architectures de modèles spécialisées ont été proposées permettant une inférence plus efficace. Les avancées les plus importantes en matière d'architecture de modèles sont Alibi, les embeddings rotationnels, l'attention multi-requête (MQA) et l'attention groupée par requête (GQA).

Tout au long de ce guide, nous proposons une analyse de la génération auto-régressive du point de vue des tenseurs. Nous examinons les avantages et les inconvénients de l’adoption d’une précision inférieure, nous explorons en détail les derniers algorithmes d’attention et nous discutons des architectures de LLM améliorées. En procédant ainsi, nous exécutons des exemples pratiques illustrant chacune des améliorations de fonctionnalités.

1. Précision réduite

Les besoins en mémoire des LLM se comprennent le mieux en considérant le LLM comme un ensemble de matrices et de vecteurs de poids, et les entrées textuelles comme une séquence de vecteurs. Dans la suite, le terme poids désignera l'ensemble des matrices et vecteurs de poids du modèle.

Au moment de la rédaction de ce guide, les LLM comptent au moins quelques milliards de paramètres. Chaque paramètre est ainsi constitué d'un nombre décimal, par exemple 4.5689, généralement stocké au format float32, bfloat16 ou float16. Cela nous permet de calculer facilement la quantité de mémoire nécessaire pour charger le LLM en mémoire :

Charger les poids d'un modèle ayant X milliards de paramètres nécessite environ 4 * X Go de VRAM en précision float32

Aujourd'hui, les modèles sont toutefois rarement entraînés en précision float32 complète, mais généralement en précision bfloat16 ou, moins fréquemment, en précision float16. La règle empirique devient donc :

Charger les poids d'un modèle ayant X milliards de paramètres nécessite environ 2 * X Go de VRAM en précision bfloat16/float16

Pour les entrées de texte courtes (moins de 1024 jetons), l’exigence mémoire pour l’inférence est très largement dominée par l’exigence mémoire nécessaire au chargement des poids. Par conséquent, pour l’instant, supposons que l’exigence mémoire pour l’inférence soit égale à l’exigence mémoire pour charger le modèle dans la VRAM du GPU.

Voici quelques exemples de la quantité approximative de VRAM nécessaire pour charger un modèle en bfloat16 :

GPT3 nécessite 2 * 175 Go = 350 Go de VRAM
Bloom nécessite 2 * 176 Go = 352 Go de VRAM
Llama-2-70b nécessite 2 * 70 Go = 140 Go de VRAM
Falcon-40b nécessite 2 * 40 Go = 80 Go de VRAM
MPT-30b nécessite 2 * 30 Go = 60 Go de VRAM
bigcode/starcoder nécessite 2 * 15,5 = 31 Go de VRAM

Au moment de la rédaction de ce document, la puce GPU la plus puissante du marché est l’A100 & H100, offrant 80 Go de VRAM. La plupart des modèles listés précédemment nécessitent plus de 80 Go rien que pour être chargés et exigent donc obligatoirement le tensor parallelism et/ou le pipeline parallelism.

🤗 Transformers prend désormais en charge le parallélisme de tenseurs pour les modèles pris en charge disposant de base_tp_plan dans leurs classes de configuration respectives. En savoir plus sur le parallélisme de tenseurs ici. De plus, si vous souhaitez écrire des modèles de manière compatible avec le parallélisme de tenseurs, n'hésitez pas à consulter la bibliothèque text-generation-inference.

Le parallélisme de pipeline naïf est pris en charge nativement. Pour cela, il suffit de charger le modèle avec device="auto", ce qui placera automatiquement les différentes couches sur les GPU disponibles, comme expliqué ici. Notez toutefois que, bien qu’efficace, ce parallélisme de pipeline naïf ne résout pas les problèmes d’inactivité des GPU. Pour cela, un parallélisme de pipeline plus avancé est requis, comme expliqué ici.

Si vous avez accès à un nœud 8 x 80 Go A100, vous pouvez charger BLOOM comme suit

!pip install transformers accelerate bitsandbytes optimum

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", pad_token_id=0)

En utilisant device_map="auto", les couches d'attention seraient réparties de manière égale sur tous les GPU disponibles.

Dans ce guide, nous utiliserons bigcode/octocoder car il peut être exécuté sur une puce de GPU A100 de 40 Go. Notez que toutes les optimisations de mémoire et de vitesse que nous appliquerons par la suite sont tout aussi applicables aux modèles nécessitant un parallélisme de modèle ou de tenseurs.

Étant donné que le modèle est chargé en précision bfloat16, en utilisant notre règle empirique ci-dessus, nous devrions nous attendre à ce que la quantité de mémoire nécessaire pour exécuter l’inférence avec bigcode/octocoder soit d’environ 31 Go de VRAM. Essayons.

Nous chargeons d'abord le modèle et le tokenizer, puis nous passons les deux à l'objet pipeline de Transformers.

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", dtype=torch.bfloat16, device_map="auto", pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer:"

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Sortie :

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

Parfait, nous pouvons maintenant utiliser directement le résultat pour convertir des octets en gigaoctets.

def bytes_to_giga_bytes(bytes):
  return bytes / 1024 / 1024 / 1024

Appelons torch.accelerator.memory.max_memory_allocated pour mesurer l'allocation de mémoire maximale de l'accélérateur.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Sortie :

29.0260648727417

Assez proche de notre estimation approximative ! On peut voir que le nombre n’est pas tout à fait exact, car le passage des octets aux kilo-octets nécessite une multiplication par 1024 plutôt que par 1000. La formule approximative peut donc aussi être comprise comme un calcul de « au plus X Go ». Notons que si nous avions tenté d’exécuter le modèle en précision float32 complète, il aurait fallu pas moins de 64 Go de VRAM.

Presque tous les modèles sont désormais entraînés en bfloat16, il n’y a donc aucune raison d’exécuter le modèle en précision float32 complète si votre GPU prend en charge bfloat16. Float32 ne donnera pas de meilleurs résultats d’inférence que la précision utilisée pour entraîner le modèle.

Si vous ne savez pas dans quel format les poids du modèle sont stockés sur le Hub, vous pouvez toujours consulter la configuration du point de contrôle sous "dtype", par exemple ici. Il est recommandé de définir le modèle sur le même type de précision que celui indiqué dans la configuration lors du chargement avec from_pretrained(..., dtype=...), sauf si le type d’origine est float32, auquel cas on peut utiliser float16 ou bfloat16 pour l’inférence.

Définissons une fonction flush(...) pour libérer toute la mémoire allouée afin de pouvoir mesurer précisément la mémoire GPU allouée au pic.

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.accelerator.empty_cache()
  torch.accelerator.reset_peak_memory_stats()

Appelons-le maintenant pour la prochaine expérience.

flush()

À partir de la bibliothèque Accelerate, vous pouvez également utiliser une méthode utilitaire indépendante du dispositif appelée release_memory, qui prend en compte divers backends matériels tels que XPU, MLU, NPU, MPS, et d'autres.

from accelerate.utils import release_memory
# ...

release_memory(model)

Et si votre GPU ne dispose pas de 32 Go de VRAM ? Il a été démontré que les poids du modèle peuvent être quantifiés en 8 bits ou 4 bits sans perte de performance significative (voir Dettmers et al.). Le modèle peut même être quantifié en 3 ou 2 bits avec une perte de performance acceptable, comme le montre le récent article GPTQ 🤯.

Sans entrer dans trop de détails, les schémas de quantisation visent à réduire la précision des poids tout en essayant de conserver les résultats d’inférence du modèle aussi précis que possible (alias aussi proches que possible de bfloat16). Notez que la quantisation fonctionne particulièrement bien pour la génération de texte, car il ne s’agit que de choisir le jeu de jetons suivants les plus probables et que l’on s’en soucie peu des valeurs exactes de la distribution des logits du jeton suivant. L’essentiel est que la distribution des logits du jeton suivant reste globalement la même, afin qu’une opération argmax ou topk donne les mêmes résultats.

Il existe diverses techniques de quantification, que nous n’aborderons pas en détail ici, mais en général, toutes les techniques de quantification fonctionnent de la manière suivante :

1. Quantifier tous les poids vers la précision cible
1. Chargez les poids quantifiés et passez la séquence d'entrée de vecteurs en précision bfloat16
1. Désquantifier dynamiquement les poids en bfloat16 pour effectuer le calcul avec leurs vecteurs d'entrée en précision bfloat16

En résumé, cela signifie que matrice d'entrée-poids multiplications, avec $X$ étant la matrice d'entrée, $W$ étant une matrice de poids et $Y$ étant la sortie : $Y = X * W$

sont modifiées en $Y = X * \text{dequantize}(W)$

pour chaque multiplication de matrices. La déquantification et la re-quantification sont effectuées séquentiellement pour toutes les matrices de poids lorsque les entrées traversent le graphe du réseau.

Par conséquent, le temps d’inférence n’est souvent pas réduit lors de l’utilisation de poids quantifiés, mais augmente au contraire. Assez de théorie, essayons ! Pour quantiser les poids avec Transformers, vous devez vous assurer que la bibliothèque bitsandbytes est installée.

!pip install bitsandbytes

Nous pouvons ensuite charger des modèles en quantification 8 bits en ajoutant simplement le drapeau load_in_8bit=True à from_pretrained.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_8bit=True), pad_token_id=0)

Maintenant, exécutons à nouveau notre exemple et mesurons l’utilisation de la mémoire.

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Sortie :

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

Parfait, nous obtenons le même résultat qu’auparavant, donc aucune perte de précision ! Voyons combien de mémoire a été utilisée cette fois-ci.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Sortie :

15.219234466552734

Beaucoup moins ! Nous sommes passés à un peu plus de 15 Go et pouvons donc exécuter ce modèle sur des GPU grand public comme le 4090. Nous observons un très bel gain en efficacité mémoire et une dégradation quasi nulle de la sortie du modèle. Cependant, nous pouvons également noter un léger ralentissement lors de l'inférence.

Nous supprimons les modèles et vidons à nouveau la mémoire.

del model
del pipe

flush()

Voyons quelle consommation de mémoire GPU de pointe la quantification 4 bits permet d’obtenir. La quantification du modèle en 4 bits se fait avec la même API que précédemment – cette fois-ci en passant load_in_4bit=True au lieu de load_in_8bit=True.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_4bit=True), pad_token_id=0)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Sortie :

Here is a Python function that transforms bytes to Giga bytes:\n\n```\ndef bytes_to_gigabytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single argument

Nous obtenons presque le même texte de sortie qu’auparavant - seul le python manque juste avant l’exemple de code. Voyons combien de mémoire a été nécessaire.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Sortie :

9.543574333190918

Seulement 9,5 Go ! Ce n'est vraiment pas beaucoup pour un modèle de plus de 15 milliards de paramètres.

Bien que nous observions une dégradation très faible de la précision pour notre modèle ici, la quantification 4 bits peut en pratique souvent conduire à des résultats différents par rapport à la quantification 8 bits ou à une inférence complète en bfloat16. C’est à l’utilisateur de tester.

Notez également que l'inférence était à nouveau un peu plus lente par rapport à la quantification 8 bits, en raison de la méthode de quantification plus agressive utilisée pour la quantification 4 bits, ce qui entraîne $\text{quantize}$ et $\text{dequantize}$ à prendre plus de temps pendant l'inférence.

del model
del pipe

flush()

Dans l’ensemble, nous avons constaté que l’exécution d’OctoCoder en précision 8 bits réduisait la VRAM GPU requise de 32 Go à seulement 15 Go, et que l’exécution du modèle en précision 4 bits réduisait encore la VRAM GPU requise à un peu plus de 9 Go.

La quantification 4 bits permet d’exécuter le modèle sur des GPU tels que RTX3090, V100 et T4, qui sont assez accessibles pour la plupart des personnes.

Pour plus d’informations sur la quantification et pour voir comment quantiser des modèles afin de nécessiter encore moins de mémoire VRAM GPU que la quantification 4 bits, nous recommandons d’examiner l’implémentation GPT-QModel.

En conclusion, il est important de rappeler que la quantification des modèles échange une efficacité mémoire accrue contre la précision et, dans certains cas, le temps d'inférence.

Si la mémoire GPU n’est pas une contrainte pour votre cas d’usage, il n’y a souvent pas besoin d’envisager la quantification. Cependant, de nombreuses GPU ne peuvent tout simplement pas exécuter des LLM sans méthodes de quantification, et dans ce cas, les schémas de quantification 4 bits et 8 bits sont des outils extrêmement utiles.

Pour des informations d’utilisation plus détaillées, nous vous recommandons vivement de consulter la documentation sur la quantification de Transformers. Passons maintenant à la manière dont nous pouvons améliorer l’efficacité computationnelle et mémoire en utilisant de meilleurs algorithmes et une architecture de modèle améliorée.

2. Flash Attention

Les LLM les plus performants d’aujourd’hui partagent plus ou moins la même architecture fondamentale, composée de couches feed-forward, de couches d’activation, de couches de normalisation par couche et, surtout, de couches d’auto-attention.

Les couches d'auto-attention sont au cœur des grands modèles de langage (LLM), car elles permettent au modèle de comprendre les relations contextuelles entre les jetons d'entrée. Cependant, la consommation mémoire GPU maximale des couches d'auto-attention augmente de manière quadratique à la fois en complexité de calcul et de mémoire avec le nombre de jetons d'entrée (également appelé longueur de séquence) que nous notons ci-dessous par $N$ . Bien que cela ne soit pas vraiment perceptible pour les séquences d'entrée courtes (jusqu'à 1000 jetons d'entrée), cela devient un problème sérieux pour les séquences d'entrée plus longues (autour de 16000 jetons d'entrée).

Examinons cela de plus près. La formule pour calculer la sortie $\mathbf{O}$ d'une couche d'auto-attention pour une entrée $\mathbf{X}$ de longueur $N$ est la suivante : $\textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \mathbf{Q} = \mathbf{W}_q \mathbf{X}, \mathbf{V} = \mathbf{W}_v \mathbf{X}, \mathbf{K} = \mathbf{W}_k \mathbf{X}$

$\mathbf{X} = (\mathbf{x}_1, ... \mathbf{x}_{N})$ constitue ainsi la séquence d'entrée de la couche d'attention. Les projections $\mathbf{Q}$ et $\mathbf{K}$ seront chacune composées de $N$ vecteurs, ce qui donne à $\mathbf{QK}^T$ une taille de $N^2$ .

Les LLM possèdent généralement plusieurs têtes d'attention, effectuant ainsi plusieurs calculs d'auto-attention en parallèle. En supposant que le LLM dispose de 40 têtes d'attention et fonctionne en précision bfloat16, nous pouvons calculer la mémoire nécessaire pour stocker les $\mathbf{QK^T}$ matrices à $40 * 2 * N^2$ octets. Pour $N=1000$ seuls environ 50 Mo de VRAM sont nécessaires, cependant, pour $N=16000$ il faudrait 19 Go de VRAM, et pour $N=100,000$ il faudrait près de 1 To juste pour stocker les $\mathbf{QK}^T$ matrices.

En résumé, l'algorithme d'auto-attention par défaut devient rapidement prohibitif en termes de mémoire pour les contextes d'entrée de grande taille.

À mesure que les LLM progressent en compréhension et en génération de texte, ils sont appliqués à des tâches de plus en plus complexes. Alors que les modèles traitaient autrefois la traduction ou la résumés de quelques phrases, ils gèrent désormais des pages entières, exigeant la capacité de traiter des longueurs d'entrée étendues.

Comment éliminer les exigences de mémoire exorbitantes pour les longueurs d'entrée importantes ? Nous avons besoin d'une nouvelle méthode pour calculer le mécanisme d'auto-attention qui se débarrasse de la $\mathbf{QK}^T$ matrice. Tri Dao et al. ont développé exactement ce type de nouvel algorithme et l'ont appelé Flash Attention.

En résumé, Flash Attention découple le calcul de $\mathbf{V} \times \text{Softmax}(\mathbf{QK}^T)$ et calcule à la place de plus petits morceaux de la sortie en itérant sur plusieurs étapes de calcul de softmax : $\textbf{O}_i \leftarrow s^a_{ij} * \textbf{O}_i + s^b_{ij} * \mathbf{V}_{j} \times \text{Softmax}(\mathbf{QK}^T_{i,j}) \text{ for multiple } i, j \text{ iterations}$

avec $s^a_{ij}$ et $s^b_{ij}$ étant des statistiques de normalisation softmax qui doivent être recalculées pour chaque $i$ et $j$ .

Veuillez noter que l’ensemble de Flash Attention est un peu plus complexe et est considérablement simplifié ici, car entrer dans trop de détails dépasse le cadre de ce guide. Le lecteur est invité à consulter le bien rédigé article Flash Attention pour plus de détails.

L'essentiel à retenir ici est le suivant :

En suivant les statistiques de normalisation softmax et en utilisant des mathématiques astucieuses, Flash Attention fournit des sorties numériquement identiques par rapport à la couche d'auto-attention par défaut, avec un coût mémoire qui n'augmente que linéairement avec $N$ .

En examinant la formule, on pourrait intuitivement penser que Flash Attention est beaucoup plus lent que la formule d'auto-attention par défaut, car davantage de calculs sont nécessaires. En effet, Flash Attention requiert plus de FLOPs que l'attention normale, car les statistiques de normalisation du softmax doivent être constamment recalculées (voir l'article pour plus de détails si vous êtes intéressé)

Cependant, Flash Attention est beaucoup plus rapide lors de l’inférence par rapport à l’attention par défaut, grâce à sa capacité à réduire considérablement les exigences sur la mémoire à large bande passante mais plus lente du GPU (VRAM), en se concentrant plutôt sur la mémoire sur puce plus rapide (SRAM).

En substance, Flash Attention s'assure que toutes les opérations intermédiaires d'écriture et de lecture peuvent être effectuées à l'aide de la mémoire on-chip SRAM rapide, plutôt que d'avoir à accéder à la mémoire VRAM plus lente pour calculer le vecteur de sortie $\mathbf{O}$ .

En pratique, il n'y a actuellement aucune raison de ne pas utiliser Flash Attention s'il est disponible. L'algorithme produit des résultats mathématiquement identiques, tout en étant à la fois plus rapide et plus efficace en termes de mémoire.

3. Innovations architecturales

Jusqu'à présent, nous avons examiné comment améliorer l'efficacité computationnelle et mémoire en :

Casting des poids vers un format de précision inférieure
Remplacer l'algorithme d'auto-attention par une version plus efficace en termes de mémoire et de calcul

Examinons maintenant comment modifier l’architecture d’un LLM pour qu’il soit le plus efficace et performant possible sur des tâches nécessitant des entrées textuelles longues, par exemple :

Questions-réponses augmentées par récupération,
Résumé,
Chat

Notez que chat exige non seulement que le LLM gère des entrées textuelles longues, mais aussi qu’il soit capable de traiter efficacement le dialogue aller-retour entre l’utilisateur et l’assistant (comme ChatGPT).

Une fois entraîné, il est difficile de modifier l’architecture fondamentale du LLM. Il est donc important de réfléchir en amont aux tâches du LLM et d’optimiser en conséquence l’architecture du modèle. Deux composants importants de l’architecture du modèle deviennent rapidement des goulets d’étranglement en termes de mémoire et/ou de performances pour les séquences d’entrée de grande taille.

Les embeddings positionnels
Le cache clé-valeur

Passons en revue chaque composant en détail

3.1 Amélioration des embeddings positionnels des LLM

L'auto-attention met chaque jeton en relation avec les jetons des autres. Par exemple, la $\text{Softmax}(\mathbf{QK}^T)$ matrice de la séquence d'entrée textuelle “« Hello », « I », « love », « you » pourrait ressembler à ceci :

Chaque jeton de mot reçoit une masse de probabilité qui détermine à quel point il porte son attention sur tous les autres jetons de mots et est donc mis en relation avec l’ensemble des autres jetons. Par exemple, le mot « love » porte son attention sur le mot « Hello » à 5 %, sur « I » à 30 % et sur lui-même à 65 %.

Un LLM basé sur l'auto-attention, mais sans embeddings de position, aurait de grandes difficultés à comprendre les positions relatives des entrées textuelles. C'est parce que le score de probabilité calculé par $\mathbf{QK}^T$ relie chaque jeton de mot à chaque autre jeton de mot en $O(1)$ calculs, indépendamment de leur distance positionnelle relative. Par conséquent, pour un LLM sans embeddings de position, chaque jeton semble avoir la même distance par rapport à tous les autres jetons, e.g. ce qui rendrait la distinction entre “« Hello I love you » et “« You love I hello » très difficile.

Pour que le LLM comprenne l’ordre des phrases, un indice supplémentaire est nécessaire et est généralement appliqué sous la forme de encodages positionnels (ou également appelés embeddings positionnels). Les encodages positionnels encodent la position de chaque jeton dans une représentation numérique que le LLM peut exploiter pour mieux comprendre l’ordre des phrases.

Les auteurs de l'article Attention Is All You Need ont introduit des embeddings positionnels sinusoïdaux $\mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N$ . où chaque vecteur $\mathbf{p}_i$ est calculé comme une fonction sinusoïdale de sa position $i$ . Les encodages positionnels sont ensuite simplement ajoutés aux vecteurs de la séquence d'entrée $\mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N$ = $\mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N$ permettant ainsi au modèle d'apprendre mieux l'ordre des phrases.

Au lieu d'utiliser des embeddings de position fixes, d'autres (tels que Devlin et al.) ont utilisé des encodages positionnels appris pour lesquels les embeddings de position $\mathbf{P}$ sont appris pendant l'entraînement.

Les embeddings de position sinusoïdaux et appris étaient autrefois les méthodes prédominantes pour encoder l'ordre des phrases dans les LLM, mais plusieurs problèmes liés à ces encodages positionnels ont été identifiés :

Les embeddings de position sinusoïdaux et appris sont tous deux des embeddings de position absolus, i.e. c'est-à-dire qu'ils encodent un embedding unique pour chaque identifiant de position : $0, \ldots, N$ . Comme l'ont montré Huang et al. et Su et al., les embeddings de position absolus entraînent de mauvaises performances des LLM pour les entrées de texte longues. Pour les entrées de texte longues, il est avantageux que le modèle apprenne la distance de position relative entre les jetons d'entrée plutôt que leur position absolue.
Lors de l'utilisation d'embeddings de position appris, le LLM doit être entraîné sur une longueur d'entrée fixe $N$, ce qui rend difficile l'extrapolation à une longueur d'entrée supérieure à celle utilisée pour l'entraînement.

Récemment, les plongements positionnels relatifs, capables de résoudre les problèmes mentionnés ci-dessus, sont devenus plus populaires, notamment :

Les deux RoPE et ALiBi soutiennent qu'il est préférable d'indiquer l'ordre des phrases au LLM directement dans l'algorithme d'auto-attention, car c'est là que les jetons de mots sont mis en relation les uns avec les autres. Plus précisément, l'ordre des phrases doit être indiqué en modifiant le $\mathbf{QK}^T$ calcul.

Sans entrer dans trop de détails, RoPE note que les informations positionnelles peuvent être encodées dans les paires requête-clé, e.g. $\mathbf{q}_i$ et $\mathbf{x}_j$ en faisant pivoter chaque vecteur d'un angle $\theta * i$ et $\theta * j$ respectivement avec $i, j$ décrivant la position de chaque vecteur dans la phrase : $\mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}}_i^T \mathbf{R}_{\theta, i -j} \mathbf{{x}}_j.$

$\mathbf{R}_{\theta, i - j}$ représente donc une matrice de rotation. $\theta$ n'est pas appris pendant l'entraînement, mais est plutôt fixé à une valeur prédéfinie qui dépend de la longueur maximale de la séquence d'entrée lors de l'entraînement.

En procédant ainsi, le score de probabilité entre $\mathbf{q}_i$ et $\mathbf{q}_j$ n'est affecté que si $i \ne j$ et ne dépend que de la distance relative $i - j$ indépendamment des positions spécifiques de chaque vecteur $i$ et $j$ .

RoPE est utilisé dans plusieurs des LLM les plus importants d’aujourd’hui, tels que :

En alternative, ALiBi propose un schéma d'encodage de position relative beaucoup plus simple. La distance relative entre les jetons d'entrée est ajoutée sous forme d'entier négatif, mis à l'échelle par une valeur prédéfinie m à chaque entrée requête-clé de la $\mathbf{QK}^T$ matrice, juste avant le calcul du softmax.

Comme le montre l’article ALiBi, ce codage positionnel relatif simple permet au modèle de conserver de bonnes performances même sur des séquences d’entrée de texte très longues.

ALiBi est utilisé dans plusieurs des LLM les plus importants d’aujourd’hui, tels que :

Les deux RoPE et ALiBi encodages de position peuvent extrapoler vers des longueurs d'entrée non vues pendant l'entraînement, alors qu'il a été démontré que l'extrapolation fonctionne beaucoup mieux par défaut pour ALiBi par rapport à RoPE. Pour ALiBi, il suffit d'augmenter les valeurs de la matrice de position triangulaire inférieure pour correspondre à la longueur de la séquence d'entrée. Pour RoPE, conserver le même $\theta$ utilisé lors de l'entraînement conduit à de mauvais résultats lorsque l'on transmet des entrées textuelles beaucoup plus longues que celles vues pendant l'entraînement, c.f Press et al.. Cependant, la communauté a trouvé quelques astuces efficaces qui adaptent $\theta$, permettant ainsi aux RoPE embeddings de position de fonctionner bien pour les séquences d'entrée textuelles extrapolées (voir ici).

RoPE et ALiBi sont des embeddings de position relative qui ne sont pas appris lors de l'entraînement, mais qui reposent plutôt sur les intuitions suivantes :

Les indices positionnels concernant les entrées textuelles doivent être fournis directement à la $\mathbf{QK}^T$ matrice de la couche d'auto-attention.
Le LLM doit être incité à apprendre un encodage positionnel à distance relative constante.
Plus les jetons d'entrée de texte sont éloignés les uns des autres, plus la probabilité de leur requête-valeur est faible. À la fois RoPE et ALiBi réduisent la probabilité requête-clé des jetons éloignés les uns des autres. RoPE la réduit en diminuant leur produit vectoriel en augmentant l'angle entre les vecteurs requête-clé. ALiBi la réduit en ajoutant de grands nombres négatifs au produit vectoriel.

En conclusion, les LLM destinés à être déployés sur des tâches nécessitant le traitement de grandes entrées textuelles sont mieux entraînés avec des embeddings positionnels relatifs, tels que RoPE et ALiBi. Notez également que même si un LLM utilisant RoPE et ALiBi a été entraîné uniquement sur une longueur fixe, par exemple $N_1 = 2048$ il peut tout de même être utilisé en pratique avec des entrées textuelles bien plus grandes que $N_1$, comme $N_2 = 8192 > N_1$, en extrapolant les embeddings positionnels.

3.2 Le cache clé-valeur

La génération de texte auto-régressive avec des LLM fonctionne en insérant itérativement une séquence d’entrée, en échantillonnant le jeton suivant, en l’ajoutant à la séquence d’entrée, et en répétant ce processus jusqu’à ce que le LLM produise un jeton indiquant que la génération est terminée.

Consultez le tutoriel Génération de texte de Transformers pour une explication plus visuelle du fonctionnement de la génération auto-régressive.

Exécutons un court extrait de code pour montrer comment fonctionne l’auto-régression en pratique. Nous prendrons simplement le jeton suivant le plus probable via torch.argmax.

input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits = model(input_ids)["logits"][:, -1:]
  next_token_id = torch.argmax(next_logits,dim=-1)

  input_ids = torch.cat([input_ids, next_token_id], dim=-1)
  print("shape of input_ids", input_ids.shape)

generated_text = tokenizer.batch_decode(input_ids[:, -5:])
generated_text

Sortie :

shape of input_ids torch.Size([1, 21])
shape of input_ids torch.Size([1, 22])
shape of input_ids torch.Size([1, 23])
shape of input_ids torch.Size([1, 24])
shape of input_ids torch.Size([1, 25])
[' Here is a Python function']

Comme nous pouvons le voir, à chaque fois, nous augmentons les jetons d'entrée du texte par le jeton tout juste échantillonné.

À quelques exceptions près, les LLM sont entraînés à l'aide de l'objectif de modélisation de langage causal et masquent donc le triangle supérieur de la matrice des scores d'attention - c'est pourquoi, dans les deux schémas ci-dessus, les scores d'attention sont laissés vides (alias ont une probabilité de 0). Pour un rappel rapide sur la modélisation de langage causal, vous pouvez consulter le blog Illustrated Self Attention.

Par conséquent, les jetons ne dépendent jamais des jetons suivants, plus précisément le vecteur $\mathbf{q}_i$ n'est jamais mis en relation avec les vecteurs clés ou valeurs $\mathbf{k}_j, \mathbf{v}_j$ si $j > i$. En revanche, $\mathbf{q}_i$ ne porte attention qu'aux vecteurs clé-valeur précédents $\mathbf{k}{m < i}, \mathbf{v}{m < i} \text{ , pour } m \in {0, \ldots i - 1}$. Afin de réduire les calculs inutiles, il est donc possible de mettre en cache les vecteurs clé-valeur de chaque couche pour toutes les étapes précédentes.

Dans la suite, nous demanderons au LLM d’utiliser le cache clé-valeur en le récupérant et en le transmettant à chaque passe avant. Dans Transformers, nous pouvons récupérer le cache clé-valeur en passant le drapeau use_cache à l’appel forward, puis le transmettre avec le jeton actuel.

past_key_values = None # past_key_values is the key-value cache
generated_tokens = []
next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True).to_tuple()
  next_logits = next_logits[:, -1:]
  next_token_id = torch.argmax(next_logits, dim=-1)

  print("shape of input_ids", next_token_id.shape)
  print("length of key-value cache", past_key_values.get_seq_length())  # past_key_values are of shape [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

Sortie :

shape of input_ids torch.Size([1, 1])
length of key-value cache 20
shape of input_ids torch.Size([1, 1])
length of key-value cache 21
shape of input_ids torch.Size([1, 1])
length of key-value cache 22
shape of input_ids torch.Size([1, 1])
length of key-value cache 23
shape of input_ids torch.Size([1, 1])
length of key-value cache 24
[' Here', ' is', ' a', ' Python', ' function']

Comme on peut le voir, lors de l'utilisation du cache clé-valeur, les jetons d'entrée texte ne sont pas allongés, mais restent un vecteur d'entrée unique. En revanche, la longueur du cache clé-valeur est augmentée d'un à chaque étape de décodage.

L'utilisation du cache clé-valeur signifie que le $\mathbf{QK}^T$ est essentiellement réduit à $\mathbf{q}_c\mathbf{K}^T$ avec $\mathbf{q}_c$ étant la projection de requête du jeton d'entrée actuellement transmis, qui est toujours un simple vecteur.

L'utilisation du cache clé-valeur présente deux avantages :

Augmentation significative de l'efficacité computationnelle, car moins de calculs sont effectués par rapport au calcul complet de la matrice $\mathbf{QK}^T$ . Cela entraîne une augmentation de la vitesse d'inférence
La mémoire maximale requise n'augmente pas de manière quadratique avec le nombre de jetons générés, mais seulement de manière linéaire.

On doit toujours utiliser le cache clé-valeur, car il produit des résultats identiques et offre un gain de vitesse significatif pour les séquences d’entrée plus longues. Transformers active le cache clé-valeur par défaut lors de l’utilisation du pipeline de texte ou de la méthode generate. Nous avons un guide entier dédié aux caches ici.

Notez que, malgré notre recommandation d’utiliser des caches clé-valeur, la sortie de votre LLM peut être légèrement différente lorsque vous les utilisez. Il s’agit d’une propriété des noyaux de multiplication de matrices eux-mêmes — vous pouvez en savoir plus à ce sujet ici.

3.2.1 Conversation multi-tours

Le cache clé-valeur est particulièrement utile pour des applications telles que le chat, où plusieurs passes de décodage auto-régressif sont nécessaires. Examinons un exemple.

User: How many people live in France?
Assistant: Roughly 75 million people live in France
User: And how many are in Germany?
Assistant: Germany has ca. 81 million inhabitants

Dans ce chat, le LLM exécute le décodage auto-régressif deux fois :

La première fois, le cache clé-valeur est vide et le prompt d'entrée est "User: How many people live in France?" et le modèle génère de manière auto-régressive le texte "Roughly 75 million people live in France" tout en augmentant le cache clé-valeur à chaque étape de décodage.
La deuxième fois, l’invite d’entrée est "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?". Grâce au cache, tous les vecteurs clé-valeur des deux premières phrases sont déjà calculés. Par conséquent, l’invite d’entrée ne se compose que de "User: And how many in Germany?". Lors du traitement de l’invite d’entrée raccourcie, ses vecteurs clé-valeur calculés sont concaténés au cache clé-valeur de la première décodage. La réponse du deuxième Assistant, "Germany has ca. 81 million inhabitants", est ensuite générée de manière auto-régressive avec un cache clé-valeur composé des vecteurs clé-valeur encodés de "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many are in Germany?".

Deux points doivent être notés ici :

Conserver tout le contexte est crucial pour les LLM déployés en chat afin que le LLM comprenne tout le contexte précédent de la conversation. Par exemple, pour l’exemple ci-dessus, le LLM doit comprendre que l’utilisateur fait référence à la population lorsqu’il pose la question "And how many are in Germany".
Le cache clé-valeur est extrêmement utile pour le chat, car il nous permet de faire croître continuellement l’historique de conversation encodé, au lieu de devoir réencoder l’historique de conversation à partir de zéro (comme ce serait le cas, par exemple, avec une architecture encodeur-décodeur).

Dans transformers, un appel generate renvoie past_key_values lorsque return_dict_in_generate=True est passé, en plus de la valeur par défaut use_cache=True. Notez que ce n'est pas encore disponible via l'interface pipeline.

# Generation as usual
prompt = system_prompt + "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(**model_inputs, max_new_tokens=60, return_dict_in_generate=True)
decoded_output = tokenizer.batch_decode(generation_output.sequences)[0]

# Piping the returned `past_key_values` to speed up the next conversation round
prompt = decoded_output + "\nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(
  **model_inputs,
  past_key_values=generation_output.past_key_values,
  max_new_tokens=60,
  return_dict_in_generate=True
)
tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]

Sortie :

 is a modified version of the function that returns Mega bytes instead.

def bytes_to_megabytes(bytes):
   return bytes / 1024 / 1024

Answer: The function takes a number of bytes as input and returns the number of

Parfait, aucun temps supplémentaire n'est consacré au recalcul des mêmes clés et valeurs pour la couche d'attention ! Il y a toutefois un point à noter. Bien que la mémoire de pointe requise pour la $\mathbf{QK}^T$ matrice soit considérablement réduite, le maintien du cache clé-valeur en mémoire peut devenir très coûteux en mémoire pour les séquences d'entrée longues ou les conversations multi-tours. N'oubliez pas que le cache clé-valeur doit stocker les vecteurs clé-valeur pour tous les vecteurs d'entrée précédents $\mathbf{x}_i \text{, for } i \in \{1, \ldots, c - 1\}$ pour toutes les couches d'auto-attention et pour toutes les têtes d'attention.

Calculons le nombre de valeurs flottantes à stocker dans le cache clé-valeur pour le LLM bigcode/octocoder que nous avons utilisé précédemment. Le nombre de valeurs flottantes correspond à deux fois la longueur de la séquence, multipliée par le nombre de têtes d’attention, multipliée par la dimension des têtes d’attention, et multipliée par le nombre de couches. Le calcul de cette valeur pour notre LLM, avec une longueur de séquence d’entrée hypothétique de 16000, donne :

config = model.config
2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head

Sortie :

7864320000

Environ 8 milliards de valeurs flottantes ! Stocker 8 milliards de valeurs flottantes en précision float16 nécessite environ 15 Go de RAM, soit à peu près la moitié de la taille des poids du modèle ! Les chercheurs ont proposé deux méthodes permettant de réduire considérablement le coût mémoire du stockage du cache clé-valeur, qui sont explorées dans les sous-sections suivantes.

3.2.2 Multi-Query-Attention (MQA)

Multi-Query-Attention a été proposé dans l’article Fast Transformer Decoding: One Write-Head is All You Need de Noam Shazeer. Comme l’indique le titre, Noam a découvert que, au lieu d’utiliser n_head poids de projections clé-valeur, on peut utiliser une seule paire de poids de projection tête-valeur partagée entre toutes les têtes d’attention, sans que les performances du modèle ne se dégradent de manière significative.

En utilisant une seule paire de poids de projection tête-valeur, les vecteurs clé-valeur $\mathbf{k}_i, \mathbf{v}_i$ doivent être identiques pour toutes les têtes d'attention, ce qui signifie que nous n'avons besoin de stocker qu'une seule paire de projection clé-valeur dans le cache au lieu de n_head paires.

La plupart des LLM utilisent entre 20 et 100 têtes d'attention, ce qui permet à MQA de réduire considérablement la consommation mémoire du cache clé-valeur. Pour le LLM utilisé dans ce notebook, nous avons donc pu réduire la consommation mémoire requise de 15 Go à moins de 400 Mo pour une longueur de séquence d'entrée de 16000.

En plus des économies de mémoire, MQA améliore également l'efficacité computationnelle, comme expliqué ci-dessous. Lors du décodage auto-régressif, de grands vecteurs clé-valeur doivent être rechargés, concaténés avec la paire de vecteurs clé-valeur actuelle, puis fournis à la $\mathbf{q}_c\mathbf{K}^T$ à chaque étape. Pour le décodage auto-régressif, la bande passante mémoire requise pour le rechargement constant peut devenir un goulot d'étranglement temporel important. En réduisant la taille des vecteurs clé-valeur, moins de mémoire doit être accédée, réduisant ainsi le goulot d'étranglement de la bande passante mémoire. Pour plus de détails, veuillez consulter l'article de Noam.

L'aspect important à comprendre ici est que réduire le nombre de têtes d'attention clé-valeur à 1 n'a de sens que si un cache clé-valeur est utilisé. La consommation mémoire maximale du modèle pour un seul passage avant sans cache clé-valeur reste inchangée, car chaque tête d'attention possède toujours un vecteur de requête unique, de sorte que chaque tête d'attention a toujours une matrice $\mathbf{QK}^T$ différente.

MQA a connu une large adoption par la communauté et est désormais utilisé par de nombreux LLM parmi les plus populaires :

De plus, le point de contrôle utilisé dans ce notebook - bigcode/octocoder - fait appel à MQA.

3.2.3 Grouped-Query-Attention (GQA)

Grouped-Query-Attention, proposé par Ainslie et al. de Google, a montré que l'utilisation de MQA peut souvent entraîner une dégradation de la qualité par rapport à l'utilisation de projections multi-têtes de clé-valeur classiques. L'article soutient qu'il est possible de préserver davantage les performances du modèle en réduisant de manière moins drastique le nombre de poids de projection des têtes de requête. Au lieu d'utiliser un seul poids de projection clé-valeur, il convient d'utiliser n < n_head poids de projection clé-valeur. En choisissant n comme une valeur nettement inférieure à n_head, par exemple 2, 4 ou 8, la quasi-totalité des gains en mémoire et en vitesse apportés par MQA peut être conservée, tout en sacrifiant moins de capacité du modèle et donc, a priori, moins de performances.

De plus, les auteurs de GQA ont constaté que les points de contrôle de modèles existants peuvent être uptrained pour adopter une architecture GQA avec seulement 5 % du calcul de pré-entraînement initial. Bien que 5 % du calcul de pré-entraînement initial puissent toujours représenter une quantité considérable, l'uptraining GQA permet aux points de contrôle existants de rester utiles pour des séquences d'entrée plus longues.

GQA a été proposé il y a peu de temps, ce qui explique son adoption encore limitée au moment de la rédaction de ce notebook. L’application la plus notable de GQA est Llama-v2.

En conclusion, il est fortement recommandé d’utiliser GQA ou MQA si le LLM est déployé avec un décodage auto-régressif et doit traiter de longues séquences d’entrée, comme c’est le cas par exemple pour le chat.

Conclusion

La communauté de recherche invente constamment de nouvelles approches ingénieuses pour accélérer le temps d’inférence des LLM de plus en plus volumineux. À titre d’exemple, l’une de ces pistes de recherche prometteuses est le décodage spéculatif, où les « jetons faciles » sont générés par des modèles de langage plus petits et plus rapides, et seuls les « jetons difficiles » sont générés par le LLM lui-même. Aborder le sujet en détail dépasse le cadre de ce notebook, mais vous pouvez en savoir plus dans ce article de blog.

La raison pour laquelle les LLM massifs tels que GPT3/4, Llama-2-70b, Claude, PaLM peuvent s’exécuter si rapidement dans des interfaces de chat telles que Hugging Face Chat ou ChatGPT est en grande partie due aux améliorations susmentionnées en matière de précision, d’algorithmes et d’architecture. À l’avenir, les accélérateurs tels que les GPU, les TPU, etc… ne feront que s’accélérer et permettre plus de mémoire, mais il faut néanmoins toujours s’assurer d’utiliser les meilleurs algorithmes et architectures disponibles pour obtenir le meilleur rapport qualité-prix 🤗

Update on GitHub