Quamba2: Ein robustes und skalierbares Framework für die Post-Training-Quantisierung von selektiven State-Space-Modellen

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

March 28, 2025

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

cs.AI

Zusammenfassung

State Space Models (SSMs) erweisen sich zunehmend als überzeugende Alternative zu Transformern aufgrund ihres konsistenten Speicherverbrauchs und ihrer hohen Leistungsfähigkeit. Dennoch ist die Skalierung von SSMs auf Cloud-Diensten oder ressourcenbeschränkten Geräten aufgrund ihrer Speicheranforderungen und Rechenleistung eine Herausforderung. Um dies zu bewältigen, kann die Quantisierung von SSMs mit Datenformaten niedriger Bitbreite die Modellgröße verringern und von Hardwarebeschleunigung profitieren. Da SSMs anfällig für Quantisierungsfehler sind, haben sich jüngste Bemühungen darauf konzentriert, ein bestimmtes Modell oder eine bestimmte Bitbreite für Effizienz zu optimieren, ohne die Leistung zu beeinträchtigen. Allerdings sind unterschiedliche Bitbreiten-Konfigurationen für verschiedene Szenarien entscheidend, wie beispielsweise W4A8 zur Steigerung der Decodiergeschwindigkeit bei großen Batches und W4A16 zur Verbesserung der Generierungsgeschwindigkeit bei kurzen Prompt-Anwendungen für einen einzelnen Benutzer. Zu diesem Zweck präsentieren wir Quamba2, das mit W8A8, W4A8 und W4A16 sowohl für Mamba1- als auch Mamba2-Backbones kompatibel ist und die wachsende Nachfrage nach SSM-Bereitstellung auf verschiedenen Plattformen adressiert. Basierend auf der kanalordnungsbewahrenden und aktivierungspersistenten Natur von SSMs schlagen wir einen Offline-Ansatz vor, um die Eingaben einer linearen Rekurrenz in 8-Bit zu quantisieren, indem die Eingabe x sortiert und geclustert wird, kombiniert mit einer pro-Zustandsgruppen-Quantisierung für die eingabeabhängigen Parameter B und C. Um die Recheninvarianz in der SSM-Ausgabe sicherzustellen, ordnen wir die Gewichte offline gemäß der Clustering-Sequenz neu an. Die Experimente zeigen, dass Quamba2-8B mehrere state-of-the-art SSM-Quantisierungsmethoden übertrifft und eine 1,3-fache bzw. 3-fache Beschleunigung in den Vorbereitungs- und Generierungsphasen bietet, während es eine 4-fache Speicherreduzierung bei nur einem durchschnittlichen Genauigkeitsverlust von 1,6 % ermöglicht. Die Auswertung auf MMLU zeigt die Generalisierbarkeit und Robustheit unseres Frameworks. Der Code und die quantisierten Modelle werden unter folgender Adresse veröffentlicht: 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.