Quamba2 : Un cadre robuste et évolutif de quantification post-entraînement pour les modèles à espace d'états sélectifs

Quamba2: A Robust and Scalable Post-training Quantization Framework for Selective State Space Models

March 28, 2025

Auteurs: Hung-Yueh Chiang, Chi-Chih Chang, Natalia Frumkin, Kai-Chiang Wu, Mohamed S. Abdelfattah, Diana Marculescu

cs.AI

Résumé

Les modèles d'espace d'états (SSMs) émergent comme une alternative convaincante aux Transformers en raison de leur utilisation de mémoire constante et de leurs performances élevées. Cependant, leur déploiement à grande échelle sur des services cloud ou des appareils à ressources limitées est difficile en raison de leurs exigences en matière de stockage et de puissance de calcul. Pour surmonter cela, la quantification des SSMs avec des formats de données à faible largeur de bits peut réduire la taille du modèle et bénéficier de l'accélération matérielle. Étant donné que les SSMs sont sensibles aux erreurs induites par la quantification, des efforts récents se sont concentrés sur l'optimisation d'un modèle ou d'une largeur de bits spécifique pour l'efficacité sans sacrifier les performances. Cependant, des configurations de largeur de bits distinctes sont essentielles pour différents scénarios, comme W4A8 pour augmenter la vitesse de décodage en grands lots, et W4A16 pour améliorer la vitesse de génération dans des applications à court prompt pour un seul utilisateur. À cette fin, nous présentons Quamba2, compatible avec W8A8, W4A8 et W4A16 pour les architectures Mamba1 et Mamba2, répondant à la demande croissante de déploiement de SSMs sur diverses plateformes. En nous basant sur la préservation de l'ordre des canaux et la persistance de l'activation des SSMs, nous proposons une approche hors ligne pour quantifier les entrées d'une récurrence linéaire en 8 bits par tri et regroupement pour l'entrée x, combinée avec une quantification par groupe d'états pour les paramètres dépendants de l'entrée B et C. Pour garantir l'invariance de calcul dans la sortie du SSM, nous réorganisons les poids hors ligne selon la séquence de regroupement. Les expériences montrent que Quamba2-8B surpasse plusieurs méthodes de quantification de SSM de pointe et offre des accélérations de 1,3 fois et 3 fois respectivement dans les étapes de pré-remplissage et de génération, tout en offrant une réduction de mémoire de 4 fois avec seulement une baisse de précision moyenne de 1,6 %. L'évaluation sur MMLU démontre la généralisabilité et la robustesse de notre cadre. Le code et les modèles quantifiés seront publiés à l'adresse : https://github.com/enyac-group/Quamba.

English

State Space Models (SSMs) are emerging as a compelling alternative to Transformers because of their consistent memory usage and high performance. Despite this, scaling up SSMs on cloud services or limited-resource devices is challenging due to their storage requirements and computational power. To overcome this, quantizing SSMs with low bit-width data formats can reduce model size and benefit from hardware acceleration. As SSMs are prone to quantization-induced errors, recent efforts have focused on optimizing a particular model or bit-width for efficiency without sacrificing performance. However, distinct bit-width configurations are essential for different scenarios, like W4A8 for boosting large-batch decoding speed, and W4A16 for enhancing generation speed in short prompt applications for a single user. To this end, we present Quamba2, compatible with W8A8, W4A8, and W4A16 for both Mamba1 and Mamba2 backbones, addressing the growing demand for SSM deployment on various platforms. Based on the channel order preserving and activation persistence of SSMs, we propose an offline approach to quantize inputs of a linear recurrence in 8-bit by sorting and clustering for input x, combined with a per-state-group quantization for input-dependent parameters B and C. To ensure compute-invariance in the SSM output, we rearrange weights offline according to the clustering sequence. The experiments show that Quamba2-8B outperforms several state-of-the-art SSM quantization methods and delivers 1.3times and 3times speed-ups in the pre-filling and generation stages, respectively, while offering 4times memory reduction with only a 1.6% average accuracy drop. The evaluation on MMLU shows the generalizability and robustness of our framework. The code and quantized models will be released at: https://github.com/enyac-group/Quamba.