Fatigué de la lenteur des grands modèles lors du traitement de longs textes ? Ne vous inquiétez pas ! L’université Tsinghua a dévoilé une technologie révolutionnaire : le framework d’inférence parallèle de séquences APB, qui équivaut à installer un moteur à turbocompresseur aux grands modèles ! Des tests ont montré que cette technologie de pointe traite les textes extrêmement longs jusqu’à 10 fois plus vite que Flash Attention ! Oui, vous avez bien entendu, 10 fois plus vite !
Avec le succès fulgurant de ChatGPT et d’autres grands modèles, les capacités de « lecture » de l’IA se sont considérablement améliorées, capables de traiter facilement des textes de plusieurs dizaines de milliers de mots. Cependant, face à un volume massif d’informations, le « cerveau » des grands modèles traditionnels peut se bloquer : l’architecture Transformer, bien que puissante, repose sur un mécanisme d’attention qui fonctionne comme un « super scanner ». Plus le texte est long, plus la portée de la numérisation augmente de façon exponentielle, ralentissant ainsi la vitesse de traitement.
Pour résoudre ce problème crucial, les scientifiques de l’université Tsinghua, en collaboration avec plusieurs instituts de recherche et géants technologiques, ont adopté une approche différente en lançant le framework APB. Le secret de ce framework réside dans la combinaison astucieuse de la « parallélisation des séquences » et de « l’attention parcimonieuse ».
En termes simples, le framework APB fonctionne comme une équipe de « travail collaboratif » efficace. Il divise le long texte en petits blocs, les répartit entre plusieurs GPU travaillant en parallèle. De plus, APB fournit à chaque GPU des compétences de « compression du cache KV local » et de « communication simplifiée », leur permettant de traiter leurs tâches respectives tout en partageant efficacement les informations clés et en résolvant ensemble les problèmes complexes de dépendance sémantique dans les longs textes.
Plus surprenant encore, le framework APB n’échange pas les performances contre la vitesse. Au contraire, lors de tests sur des textes extrêmement longs de 128 k, APB a non seulement considérablement augmenté la vitesse, mais a également surpassé Flash Attention ! Même Star Attention, fortement promu par Nvidia, a été dépassé par APB, avec une amélioration de la vitesse de 1,6 fois, faisant de lui un véritable « as » polyvalent.
Cette technologie révolutionnaire a pour application la plus directe de réduire considérablement le temps de réponse du premier jeton lors du traitement de longs textes par les grands modèles. Cela signifie qu’à l’avenir, les grands modèles équipés du framework APB pourront comprendre instantanément et répondre rapidement aux longues instructions des utilisateurs, éliminant ainsi l’attente interminable du « chargement en cours… ».
Alors, comment le framework APB réussit-il à accélérer le traitement de manière aussi spectaculaire ?
Le framework APB comprend parfaitement le « point faible » du traitement des longs textes : la charge de calcul. La charge de calcul du mécanisme d’attention traditionnel est proportionnelle au carré de la longueur du texte, les longs textes étant des « trous noirs » de calcul. Pour surmonter cet obstacle, le framework APB utilise deux « techniques magiques » :
Première technique : augmenter le parallélisme, pour que « l’union fasse la force »
Le framework APB exploite pleinement les avantages du calcul distribué, répartissant les tâches de calcul sur plusieurs GPU, comme une « collaboration », ce qui multiplie naturellement l’efficacité. En particulier en ce qui concerne le parallélisme des séquences, le framework APB présente une extensibilité extrêmement élevée, sans être limité par la structure du modèle, même les textes les plus longs peuvent être facilement traités.
Deuxième technique : réduire les calculs inutiles, pour que « le bon acier soit utilisé à bon escient »
Le framework APB introduit un mécanisme d’attention parcimonieuse, ne traitant pas « tout d’un coup », mais effectuant des calculs « sélectifs ». Il fonctionne comme un expert « à l’œil perçant », se concentrant uniquement sur les informations clés du texte et ignorant les parties non pertinentes, réduisant ainsi considérablement la charge de calcul.
Cependant, ces deux techniques, « parallélisme » et « parcimonie », semblent simples, mais cachent en réalité des « secrets ». Comment réaliser un calcul d’attention parcimonieuse efficace dans un framework de parallélisme des séquences ? C’est là que réside le véritable « cœur » du framework APB.
Dans un environnement de parallélisme des séquences, chaque GPU ne possède qu’une partie des informations textuelles. Réaliser une attention parcimonieuse avec une « perception globale » est comme « toucher un éléphant aveugle », la difficulté est donc imaginable. Les méthodes précédentes telles que Star Attention et APE, soit sacrifient les performances, soit ont des applications limitées, et n’ont pas réussi à résoudre parfaitement ce problème.
Le framework APB contourne astucieusement le « piège » de la « communication à grande échelle », en adoptant une approche différente et en construisant un mécanisme d’attention parcimonieuse à faible communication pour les scénarios de parallélisme des séquences. Les composants principaux de ce mécanisme comprennent :
Un bloc d’ancrage (Anchor block) plus compact : le bloc d’ancrage fonctionne comme un « système de navigation », guidant le mécanisme d’attention vers les informations clés. Le framework APB réduit de manière innovante la taille du bloc d’ancrage, le rendant plus léger et plus flexible, réduisant ainsi les coûts de calcul.
Un bloc de transmission (Passing block) original : le bloc de transmission est le composant « âme » du framework APB. Il résout astucieusement les problèmes de dépendance sémantique à longue distance. En « compressant et empaquetant » les informations clés traitées par le GPU précédent et en les transmettant au GPU suivant, chaque « membre de l’équipe » peut avoir une « vue d’ensemble » et comprendre le contexte du long texte.
Une compression de contexte sensible à la requête : le framework APB introduit également un mécanisme « sensible à la requête », permettant au compresseur de contexte de « comprendre la question » et de filtrer et de conserver plus précisément les informations clés liées à la requête, améliorant ainsi l’efficacité et la précision.
Sur la base de ces « techniques exclusives », le framework APB construit un processus d’inférence fluide :
Segmentation du contexte : le long texte est réparti uniformément entre les différents GPU, et un bloc d’ancrage est ajouté au début, « intégrant » la question de la requête.
Compression du contexte : en utilisant Locret, les têtes conservées sont utilisées pour « compresser intelligemment » le cache KV.
Communication efficace : l’opérateur AllGather est utilisé pour « transmettre » le cache KV compressé aux GPU suivants, créant le bloc de transmission.
Calcul rapide : un noyau Flash Attention spécial est utilisé, ainsi qu’un masque d’attention optimisé, pour effectuer des calculs efficaces. Le bloc de transmission se « retire » après le calcul et ne participe pas aux calculs suivants.
Les résultats expérimentaux démontrent de manière éloquente les performances exceptionnelles du framework APB. Lors de tests sur plusieurs modèles tels que Llama-3.1-8B-instruct, Qwen-2.5-14B-instruct et Yi-34B-200K, ainsi que sur plusieurs benchmarks tels qu’InfiniteBench et RULER, le framework APB a surclassé tous les autres, atteignant le meilleur équilibre entre performances et vitesse.
Il est particulièrement important de noter que l’avantage en vitesse du framework APB devient de plus en plus évident à mesure que la longueur du texte augmente, réalisant un effet « plus c’est long, plus c’est rapide ». Le secret réside dans le fait que la charge de calcul du framework APB est bien inférieure à celle des autres méthodes, et l’écart s’élargit avec l’augmentation de la longueur du texte.
Une analyse plus approfondie du temps de pré-remplissage montre que la technique de parallélisation des séquences permet à elle seule de réduire considérablement le temps de calcul de l’attention et du FFN (réseau neuronal feed-forward). Le mécanisme d’attention parcimonieuse du framework APB comprime encore plus le temps de calcul de l’attention. Comparé à Star Attention, le framework APB utilise astucieusement le bloc de transmission pour transmettre les dépendances sémantiques à longue distance, réduisant considérablement la taille du bloc d’ancrage et réduisant efficacement les coûts supplémentaires du FFN, réalisant un effet « avoir le beurre et l’argent du beurre ».
Plus encourageant encore, le framework APB présente une compatibilité exceptionnelle, capable de s’adapter de manière flexible à différents environnements distribués et à différentes échelles de modèles, maintenant des performances et une efficacité élevées et stables dans diverses conditions « rigoureuses ».
On peut prévoir qu’avec l’arrivée du framework APB, le « goulot d’étranglement » de l’inférence de longs textes par les grands modèles sera complètement éliminé, et l’espace d’imagination des applications de l’IA sera considérablement élargi. À l’avenir, que ce soit pour le service client intelligent, l’analyse financière, la recherche scientifique ou la création de contenu, nous entrerons dans une nouvelle ère de l’IA « plus rapide, plus puissante et plus intelligente » !
Adresse du projet : https://github.com/thunlp/APB
Adresse de l’article : https://arxiv.org/pdf/2502.12085