Déploiement efficace de grands modèles de langage à faible nombre de bits sur CPU
L'innovation clé de T-MAC réside dans l'adoption d'un paradigme de calcul basé sur des tables de consultation (LUT), plutôt que le paradigme traditionnel de multiplication-accumulation (MAC). T-MAC utilise des tables de consultation pour prendre directement en charge les calculs à faible nombre de bits, éliminant ainsi les opérations de déquantification nécessaires dans d'autres systèmes, et réduisant considérablement le nombre d'opérations de multiplication et d'addition.
Les expériences ont montré que T-MAC offre des performances exceptionnelles : sur un PC AI Surface équipé de la dernière puce Qualcomm Snapdragon X E lite, le modèle BitNet-b1.58 de 3B peut générer 48 tokens par seconde, le modèle llama 7B 2 bits peut générer 30 tokens par seconde, et le modèle llama 7B 4 bits peut générer 20 tokens par seconde.
Cela surpasse même les performances du NPU !
Lors du déploiement du modèle llama-2-7b-4bit, bien que le NPU puisse générer 10,4 tokens par seconde, le CPU, avec l'aide de T-MAC, peut atteindre 12,6 tokens par seconde en utilisant seulement deux cœurs, et peut même atteindre jusqu'à 22 tokens par seconde.
Tout cela dépasse de loin la vitesse moyenne de lecture humaine, avec une amélioration de 4 à 5 fois par rapport au framework llama.cpp original.
Même sur des appareils de gamme inférieure comme le Raspberry Pi 5, T-MAC peut atteindre un taux de génération de 11 tokens par seconde pour le BitNet-b1.58 3B. T-MAC présente également des avantages significatifs en termes de consommation d'énergie : pour atteindre le même taux de génération, T-MAC ne nécessite que 1/4 à 1/6 du nombre de cœurs par rapport au llama.cpp original, réduisant ainsi la consommation d'énergie tout en laissant des ressources de calcul pour d'autres applications.
Il est à noter que les performances de calcul de T-MAC augmentent linéairement à mesure que le nombre de bits diminue, un phénomène difficile à observer dans les GPU et NPU basés sur la déquantification. Cependant, T-MAC peut atteindre 10 tokens par seconde sur un seul cœur et 28 tokens par seconde sur quatre cœurs avec 2 bits, surpassant largement les performances du NPU.
La multiplication matricielle sans multiplication, juste avec des tables de consultation (LUT)
Pour les paramètres à faible nombre de bits (poids), T-MAC groupe chaque bit séparément (par exemple, un groupe de 4 bits), multiplie ces bits avec le vecteur d'activation, pré-calcule toutes les sommes partielles possibles, puis les stocke dans une LUT.
Ensuite, T-MAC utilise des opérations de décalage et d'accumulation pour prendre en charge un nombre de bits extensible de 1 à 4. Avec cette méthode, T-MAC abandonne les instructions FMA (multiplication-addition) peu efficaces sur CPU au profit des instructions TBL/PSHUF (consultation de table) plus efficaces et moins énergivores.
Calcul centré sur les bits, remplaçant le calcul centré sur les types de données
Le calcul traditionnel basé sur la déquantification est en réalité un calcul centré sur les types de données, ce qui nécessite une personnalisation pour chaque type de données différent.
Chaque combinaison de largeur de bits pour l'activation et les poids, comme W4A16 (poids int4, activation float16) et W2A8, nécessite une disposition des poids et un noyau de calcul spécifiques.
Par exemple, la disposition W3 nécessite de séparer 2 bits et 1 bit, et d'utiliser différentes méthodes d'entrelacement ou de mélange pour l'alignement mémoire ou le décodage rapide. Ensuite, le noyau de calcul correspondant doit déballer cette disposition spécifique dans les types de données pris en charge par le matériel pour l'exécution.
T-MAC, en observant le calcul de multiplication matricielle à faible nombre de bits du point de vue des bits, ne nécessite de concevoir une structure de données optimale que pour un seul bit, puis de l'étendre à 2/3/4 bits par empilement.
En même temps, pour les vecteurs d'activation de différentes précisions (float16/float32/int8), seul le processus de construction de la table doit être modifié, sans avoir à considérer différentes structures de données lors de la consultation de la table.
De plus, avec les méthodes traditionnelles basées sur la déquantification, lorsqu'on passe de 4 bits à 3/2/1 bits, bien que l'occupation mémoire soit réduite, la charge de calcul ne diminue pas, et les performances peuvent même être pires en raison de l'augmentation des frais généraux de déquantification.
Cependant, la charge de calcul de T-MAC diminue linéairement avec la réduction du nombre de bits, offrant ainsi une meilleure accélération avec un nombre de bits plus faible, fournissant une solution de déploiement efficace pour les derniers travaux tels que BitNet, EfficientQAT et d'autres modèles à 1 bit/2 bits.
Implémentation hautement optimisée des opérateurs
Le calcul centré sur les bits présente de nombreux avantages, mais sa mise en œuvre sur CPU reste un défi considérable :
(1) Comparé à l'accès séquentiel aux données d'activation et de poids, l'accès aux tables est aléatoire. La résidence des tables dans la mémoire rapide sur puce est cruciale pour les performances finales d'inférence ;
(2) Cependant, la mémoire sur puce est limitée, et la méthode des tables de consultation (LUT) augmente l'utilisation de la mémoire sur puce par rapport au mpGEMV traditionnel. Cela est dû au fait que les tables de consultation doivent stocker les résultats de la multiplication du vecteur d'activation avec tous les motifs de bits possibles, ce qui est beaucoup plus que l'activation elle-même.
Pour relever ces défis, les chercheurs du Microsoft Research Asia ont approfondi le flux de données de calcul basé sur les tables de consultation et ont conçu des structures de données et des processus de calcul efficaces pour ce paradigme de calcul, notamment :
-
Stocker les LUT dans la mémoire sur puce pour tirer parti des instructions vectorielles de consultation de table (TBL/PSHUF) sur CPU afin d'améliorer les performances d'accès aléatoire à la mémoire.
-
Modifier l'ordre de calcul des axes de la matrice pour maximiser le taux de réutilisation des données des LUT limitées placées dans la mémoire sur puce.
-
Concevoir une méthode de tuilage matriciel optimale spécifiquement pour la consultation de table, combinée à la recherche de paramètres de tuilage optimaux avec autotvm.
-
Optimisation de la disposition des paramètres (poids) :
a) Réorganisation des poids pour maximiser l'accès séquentiel et améliorer le taux de succès du cache.
b) Entrelacement des poids pour améliorer l'efficacité du décodage.
- Optimisations spécifiques pour les CPU Intel/ARM, notamment :
a) Réorganisation des registres pour une construction rapide des tables de consultation.
b) Utilisation d'instructions de moyenne pour une accumulation rapide sur 8 bits.
Les chercheurs ont appliqué progressivement diverses optimisations sur une implémentation de base, obtenant finalement une accélération significative par rapport aux opérateurs à faible nombre de bits de l'état de l'art :