Comment accélérer vos modèles d’IA sans abîmer votre mémoire

Table des liens

Résumé et 1 Introduction

2. Travaux connexes

3. Contexte

3.1. Notation et 3.2. Inférence du modèle de langage

3.3. Multi-Requête, Multi-tête et l'attention multi-requête généralisée

4. Attention bifurquée sensible au contexte et 4.1. Motivation

4.2. Formulation et 4.3. Complexité des E/S mémoire

5. Expériences

5.1. Comparaison des capacités d'attention multi-têtes, multi-requêtes et multi-groupes

5.2. Latences des modèles à capacités équivalentes

5.3. Applications

6. Conclusion et références

A. FAQ

B. Travaux connexes

C. Installation

D. Attention multi-groupe Famille

E. Attention bifurquée sensible au contexte

F. Applications : Résultats supplémentaires

G. Compatibilité avec les techniques de décodage spéculatif et de décodage rapide

Abstrait

Dans notre étude, nous présentons l'attention bifurquée , une méthode développée pour l'inférence de modèles de langage dans des contextes d'échantillonnage par lots à contexte unique. Cette approche vise à réduire les coûts d'E/S de mémoire redondants, un facteur important de latence pour les tailles de lots élevées et les longueurs de contexte longues. L'attention bifurquée y parvient en divisant le mécanisme d'attention pendant le décodage incrémental en deux opérations GEMM distinctes, en se concentrant sur le cache KV du pré-remplissage et le processus de décodage. Cette méthode garantit un calcul précis et maintient la charge de calcul habituelle (FLOP) des mécanismes d'attention standard, mais avec des E/S de mémoire réduites. L'attention bifurquée est également compatible avec le mécanisme d'attention multi-requêtes connu pour réduire les E/S de mémoire pour le cache KV, ce qui permet en outre une taille de lot et une longueur de contexte plus élevées. L'efficacité qui en résulte conduit à une latence plus faible, améliorant l'adéquation aux applications en temps réel, par exemple, en permettant la génération de réponses massivement parallèles sans augmenter considérablement la latence, améliorant les performances lorsqu'elle est intégrée à des techniques de post-traitement telles que le reclassement.

1. Introduction

L'avènement des grands modèles de langage (LLM) a inauguré une nouvelle ère de l'apprentissage automatique, affichant des performances remarquables sur un large éventail de tâches (Brown et al., 2020 ; OpenAI, 2023 ; Chowdhery et al., 2022 ; Touvron et al., 2023 ; Chen et al., 2021 ; Hoffmann et al., 2022 ; Li et al., 2022 ; Microsoft ; Amazon, 2022 ; Nijkamp et al., 2023). Malgré leurs capacités impressionnantes, le déploiement de ces modèles à grande échelle dans des applications pratiques pose des défis importants, notamment en termes de latence et d'efficacité d'inférence. L'amélioration de ces aspects est essentielle, car ils influencent directement les ressources de calcul nécessaires pour générer des prédictions et permettre la mise en œuvre pratique de ces modèles avancés dans divers secteurs.

Un scénario d'inférence particulièrement exigeant est l'échantillonnage par lots à contexte unique, où l'objectif est de générer plusieurs complétions à partir d'un seul contexte. Cette tâche est couramment rencontrée dans de nombreuses applications telles que les outils IDE d'édition de code qui fournissent plusieurs recommandations, ou dans les cas où un classement parmi plusieurs générations est nécessaire pour des performances optimales (via des mesures de classement telles que la probabilité logarithmique moyenne, le vote majoritaire, etc.). Le décodage incrémentiel d'un tel scénario d'échantillonnage nécessite beaucoup d'E/S de mémoire, ce qui devient un goulot d'étranglement de latence pour les lots et les longueurs de contexte élevés.

Dans cette étude, nous étudions deux stratégies compatibles pour relever les défis des E/S de mémoire dans l'inférence des transformateurs : (1) une étude des requêtes multiples et de leurs compromis, et (2) une nouvelle technique appelée attention bifurquée sensible au contexte.

Notre étude commence par une analyse de l'attention multi-requêtes généralisée (Ainslie et al., 2023), qui inclut la multi-requête (Shazeer, 2019), ainsi que le mécanisme d'attention multi-têtes établi (Vaswani et al., 2017) pour le compromis entre performances et latence. Nos résultats montrent une mise à l'échelle fluide des performances avec l'augmentation de la taille du modèle pour une valeur fixe du nombre de groupes g pour la multi-requête généralisée[1]. La diminution de g entraîne un décalage vers le haut des courbes de mise à l'échelle de la perte de validation par rapport à la taille du modèle. La relation cohérente entre la compression du cache, la taille du modèle et la perte de validation nous permet de faire un compromis entre l'efficacité de l'inférence et la taille du modèle, c'est-à-dire qu'elle nous permet de sélectionner une compression plus élevée pour les cas d'utilisation nécessitant une efficacité élevée, tout en faisant correspondre les performances de l'attention multi-têtes en compensant avec une taille de modèle plus grande.

Deuxièmement, nous introduisons l'attention bifurquée sensible au contexte, une technique qui bifurque toute attention de la famille de requêtes multiples généralisées en composants de contexte et de décodage lors du décodage incrémental. Une telle bifurcation implique le même nombre de FLOP et produit des résultats identiques par rapport à l'attention d'origine, mais peut réduire considérablement le coût des E/S de mémoire et donc la latence dans les scénarios de lots et de longueur de contexte élevés. Cette approche permet la génération de plusieurs complétions en temps réel sans entraîner de coûts de latence supplémentaires importants, ou permet des tailles de lots beaucoup plus élevées conduisant à des performances de classement améliorées. Par exemple, pour le modèle multi-têtes CodeGen 16B (Nijkamp et al., 2022) avec une longueur de contexte de 2 000, nous sommes en mesure d'augmenter la taille du lot à 128 avec une attention bifurquée, par rapport à une taille de lot de seulement 5 sans, ce qui entraîne une augmentation du pass@k (Chen et al., 2021) de 59,0 % à 84,6 %, ou du pass@top3 via le log-p moyen augmentant de 55,2 % à 58,1 %.

[1] Des valeurs inférieures des groupes d'attention g conduisent à une compression plus élevée des tenseurs clé-valeur, comme dans le cas de requêtes multiples où g = 1, améliorant ainsi l'efficacité de l'inférence et la latence en raison du cache KV réduit par rapport au cas de plusieurs têtes où g = h, le nombre de têtes d'attention de requête.