Quamba2: Un Framework Robusto e Scalabile per la Quantizzazione Post-Addestramento di Modelli a Spazio di Stati Selettivi

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

March 28, 2025

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

cs.AI

Abstract

I modelli a spazio di stato (SSM) stanno emergendo come un'alternativa convincente ai Transformer grazie al loro utilizzo di memoria costante e alle elevate prestazioni. Nonostante ciò, scalare gli SSM su servizi cloud o dispositivi con risorse limitate è impegnativo a causa dei loro requisiti di archiviazione e potenza computazionale. Per superare questa sfida, la quantizzazione degli SSM con formati di dati a bassa larghezza di bit può ridurre le dimensioni del modello e trarre vantaggio dall'accelerazione hardware. Poiché gli SSM sono suscettibili a errori indotti dalla quantizzazione, recenti sforzi si sono concentrati sull'ottimizzazione di un modello o di una larghezza di bit specifica per migliorare l'efficienza senza sacrificare le prestazioni. Tuttavia, configurazioni distinte di larghezza di bit sono essenziali per scenari diversi, come W4A8 per aumentare la velocità di decodifica in batch di grandi dimensioni e W4A16 per migliorare la velocità di generazione in applicazioni con prompt brevi per un singolo utente. A tal fine, presentiamo Quamba2, compatibile con W8A8, W4A8 e W4A16 per entrambe le architetture Mamba1 e Mamba2, rispondendo alla crescente domanda di distribuzione degli SSM su varie piattaforme. Basandoci sulla conservazione dell'ordine dei canali e sulla persistenza dell'attivazione degli SSM, proponiamo un approccio offline per quantizzare gli ingressi di una ricorrenza lineare a 8 bit ordinando e raggruppando l'input x, combinato con una quantizzazione per gruppo di stati per i parametri dipendenti dall'input B e C. Per garantire l'invarianza computazionale nell'output dell'SSM, riorganizziamo i pesi offline in base alla sequenza di raggruppamento. Gli esperimenti mostrano che Quamba2-8B supera diversi metodi di quantizzazione SSM all'avanguardia e offre un aumento di velocità di 1,3 volte e 3 volte rispettivamente nelle fasi di pre-riempimento e generazione, riducendo la memoria di 4 volte con un calo di accuratezza medio di solo l'1,6%. La valutazione su MMLU dimostra la generalizzabilità e la robustezza del nostro framework. Il codice e i modelli quantizzati saranno rilasciati su: 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.